Updated pr.

oracle · Dec 6, 2024 · cc66c42 · cc66c42
1 parent 238b6d9
commit cc66c42
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/tests/unitary/with_extras/langchain/chat_models/test_oci_data_science.py b/tests/unitary/with_extras/langchain/chat_models/test_oci_data_science.py
@@ -10,6 +10,9 @@
 from typing import Any, AsyncGenerator, Dict, Generator
 from unittest import mock
 
+from ads.llm.langchain.plugins.chat_models.oci_data_science import (
+    DEFAULT_INFERENCE_ENDPOINT_CHAT,
+)
 import pytest
 
 
@@ -126,6 +129,7 @@ def test_invoke_vllm(*args: Any) -> None:
     output = llm.invoke(CONST_PROMPT)
     assert isinstance(output, AIMessage)
     assert output.content == CONST_COMPLETION
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT_CHAT}
 
 
 @pytest.mark.requires("ads")
@@ -138,6 +142,7 @@ def test_invoke_tgi(*args: Any) -> None:
     output = llm.invoke(CONST_PROMPT)
     assert isinstance(output, AIMessage)
     assert output.content == CONST_COMPLETION
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT_CHAT}
 
 
 @pytest.mark.requires("ads")
@@ -162,6 +167,7 @@ def test_stream_vllm(*args: Any) -> None:
     assert output is not None
     if output is not None:
         assert str(output.content).strip() == CONST_COMPLETION
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT_CHAT}
 
 
 async def mocked_async_streaming_response(
@@ -194,3 +200,4 @@ async def test_stream_async(*args: Any) -> None:
     ):
         chunks = [str(chunk.content) async for chunk in llm.astream(CONST_PROMPT)]
     assert "".join(chunks).strip() == CONST_COMPLETION
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT_CHAT}
diff --git a/tests/unitary/with_extras/langchain/llms/test_oci_model_deployment_endpoint.py b/tests/unitary/with_extras/langchain/llms/test_oci_model_deployment_endpoint.py
@@ -10,6 +10,9 @@
 from typing import Any, AsyncGenerator, Dict, Generator
 from unittest import mock
 
+from ads.llm.langchain.plugins.llms.oci_data_science_model_deployment_endpoint import (
+    DEFAULT_INFERENCE_ENDPOINT,
+)
 import pytest
 
 if sys.version_info < (3, 9):
@@ -118,6 +121,7 @@ def test_invoke_vllm(*args: Any) -> None:
     llm = OCIModelDeploymentVLLM(endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME)
     output = llm.invoke(CONST_PROMPT)
     assert output == CONST_COMPLETION
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
 
 
 @pytest.mark.requires("ads")
@@ -135,6 +139,7 @@ def test_stream_tgi(*args: Any) -> None:
         count += 1
     assert count == 4
     assert output.strip() == CONST_COMPLETION
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
 
 
 @pytest.mark.requires("ads")
@@ -147,6 +152,7 @@ def test_generate_tgi(*args: Any) -> None:
     )
     output = llm.invoke(CONST_PROMPT)
     assert output == CONST_COMPLETION
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
 
 
 @pytest.mark.asyncio
@@ -170,3 +176,4 @@ async def test_stream_async(*args: Any) -> None:
     ):
         chunks = [chunk async for chunk in llm.astream(CONST_PROMPT)]
     assert "".join(chunks).strip() == CONST_COMPLETION
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}