Updated pr.

Author: lu-ohai

tests/unitary/with_extras/langchain/chat_models/test_oci_data_science.py

@@ -126,10 +126,10 @@ def mocked_requests_post(url: str, **kwargs: Any) -> MockResponse:
 def test_invoke_vllm(*args: Any) -> None:
     """Tests invoking vLLM endpoint."""
     llm = ChatOCIModelDeploymentVLLM(endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME)
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT_CHAT}
     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")
@@ -139,10 +139,10 @@ def test_invoke_vllm(*args: Any) -> None:
 def test_invoke_tgi(*args: Any) -> None:
     """Tests invoking TGI endpoint using OpenAI Spec."""
     llm = ChatOCIModelDeploymentTGI(endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME)
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT_CHAT}
     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")
@@ -154,6 +154,7 @@ def test_stream_vllm(*args: Any) -> None:
     llm = ChatOCIModelDeploymentVLLM(
         endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME, streaming=True
     )
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT_CHAT}
     output = None
     count = 0
     for chunk in llm.stream(CONST_PROMPT):
@@ -167,7 +168,6 @@ 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(
@@ -193,11 +193,11 @@ async def test_stream_async(*args: Any) -> None:
     llm = ChatOCIModelDeploymentVLLM(
         endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME, streaming=True
     )
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT_CHAT}
     with mock.patch.object(
         llm,
         "_aiter_sse",
         mock.MagicMock(return_value=mocked_async_streaming_response()),
     ):
         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}

tests/unitary/with_extras/langchain/llms/test_oci_model_deployment_endpoint.py

@@ -119,9 +119,9 @@ async def mocked_async_streaming_response(
 def test_invoke_vllm(*args: Any) -> None:
     """Tests invoking vLLM endpoint."""
     llm = OCIModelDeploymentVLLM(endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME)
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
     output = llm.invoke(CONST_PROMPT)
     assert output == CONST_COMPLETION
-    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
 
 
 @pytest.mark.requires("ads")
@@ -132,14 +132,14 @@ def test_stream_tgi(*args: Any) -> None:
     llm = OCIModelDeploymentTGI(
         endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME, streaming=True
     )
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
     output = ""
     count = 0
     for chunk in llm.stream(CONST_PROMPT):
         output += chunk
         count += 1
     assert count == 4
     assert output.strip() == CONST_COMPLETION
-    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
 
 
 @pytest.mark.requires("ads")
@@ -150,9 +150,9 @@ def test_generate_tgi(*args: Any) -> None:
     llm = OCIModelDeploymentTGI(
         endpoint=CONST_ENDPOINT, api="/generate", model=CONST_MODEL_NAME
     )
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
     output = llm.invoke(CONST_PROMPT)
     assert output == CONST_COMPLETION
-    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
 
 
 @pytest.mark.asyncio
@@ -169,11 +169,11 @@ async def test_stream_async(*args: Any) -> None:
     llm = OCIModelDeploymentTGI(
         endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME, streaming=True
     )
+    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}
     with mock.patch.object(
         llm,
         "_aiter_sse",
         mock.MagicMock(return_value=mocked_async_streaming_response()),
     ):
         chunks = [chunk async for chunk in llm.astream(CONST_PROMPT)]
     assert "".join(chunks).strip() == CONST_COMPLETION
-    assert llm.headers == {"route": DEFAULT_INFERENCE_ENDPOINT}