Update tests.

Author: qiuosier

libs/oci/tests/unit_tests/chat_models/test_oci_data_science.py

@@ -1,196 +1,30 @@
 """Test Chat model for OCI Data Science Model Deployment Endpoint."""
 
-import sys
-from typing import Any, AsyncGenerator, Dict, Generator
 from unittest import mock
 
+import httpx
 import pytest
-from langchain_core.messages import AIMessage, AIMessageChunk
-from requests.exceptions import HTTPError
 
-from langchain_oci.chat_models import (
-    ChatOCIModelDeploymentTGI,
-    ChatOCIModelDeploymentVLLM,
-)
-
-CONST_MODEL_NAME = "odsc-vllm"
-CONST_ENDPOINT = "https://oci.endpoint/ocid/predict"
-CONST_PROMPT = "This is a prompt."
-CONST_COMPLETION = "This is a completion."
-CONST_COMPLETION_ROUTE = "/v1/chat/completions"
-CONST_COMPLETION_RESPONSE = {
-    "id": "chat-123456789",
-    "object": "chat.completion",
-    "created": 123456789,
-    "model": "mistral",
-    "choices": [
-        {
-            "index": 0,
-            "message": {
-                "role": "assistant",
-                "content": CONST_COMPLETION,
-                "tool_calls": [],
-            },
-            "logprobs": None,
-            "finish_reason": "length",
-            "stop_reason": None,
-        }
-    ],
-    "usage": {"prompt_tokens": 115, "total_tokens": 371, "completion_tokens": 256},
-    "prompt_logprobs": None,
-}
-CONST_COMPLETION_RESPONSE_TGI = {"generated_text": CONST_COMPLETION}
-CONST_STREAM_TEMPLATE = (
-    'data: {"id":"chat-123456","object":"chat.completion.chunk","created":123456789,'
-    '"model":"odsc-llm","choices":[{"index":0,"delta":<DELTA>,"finish_reason":null}]}'
-)
-CONST_STREAM_DELTAS = ['{"role":"assistant","content":""}'] + [
-    '{"content":" ' + word + '"}' for word in CONST_COMPLETION.split(" ")
-]
-CONST_STREAM_RESPONSE = (
-    content
-    for content in [
-        CONST_STREAM_TEMPLATE.replace("<DELTA>", delta).encode()
-        for delta in CONST_STREAM_DELTAS
-    ]
-    + [b"data: [DONE]"]
-)
-
-CONST_ASYNC_STREAM_TEMPLATE = (
-    '{"id":"chat-123456","object":"chat.completion.chunk","created":123456789,'
-    '"model":"odsc-llm","choices":[{"index":0,"delta":<DELTA>,"finish_reason":null}]}'
-)
-CONST_ASYNC_STREAM_RESPONSE = (
-    CONST_ASYNC_STREAM_TEMPLATE.replace("<DELTA>", delta).encode()
-    for delta in CONST_STREAM_DELTAS
-)
-
-pytestmark = pytest.mark.skipif(
-    sys.version_info < (3, 9), reason="Requires Python 3.9 or higher"
-)
-
-
-class MockResponse:
-    """Represents a mocked response."""
-
-    def __init__(self, json_data: Dict, status_code: int = 200):
-        self.json_data = json_data
-        self.status_code = status_code
-
-    def raise_for_status(self) -> None:
-        """Mocked raise for status."""
-        if 400 <= self.status_code < 600:
-            raise HTTPError()  # type: ignore[call-arg]
-
-    def json(self) -> Dict:
-        """Returns mocked json data."""
-        return self.json_data
-
-    def iter_lines(self, chunk_size: int = 4096) -> Generator[bytes, None, None]:
-        """Returns a generator of mocked streaming response."""
-        return CONST_STREAM_RESPONSE
-
-    @property
-    def text(self) -> str:
-        """Returns the mocked text representation."""
-        return ""
-
-
-def mocked_requests_post(url: str, **kwargs: Any) -> MockResponse:
-    """Method to mock post requests"""
-
-    payload: dict = kwargs.get("json", {})
-    messages: list = payload.get("messages", [])
-    prompt = messages[0].get("content")
-
-    if prompt == CONST_PROMPT:
-        return MockResponse(json_data=CONST_COMPLETION_RESPONSE)
-
-    return MockResponse(
-        json_data={},
-        status_code=404,
-    )
+from langchain_oci import ChatOCIModelDeployment
 
 
 @pytest.mark.requires("ads")
 @pytest.mark.requires("langchain_openai")
-@mock.patch("ads.common.auth.default_signer", return_value=dict(signer=None))
-@mock.patch("requests.post", side_effect=mocked_requests_post)
-def test_invoke_vllm(*args: Any) -> None:
-    """Tests invoking vLLM endpoint."""
-    llm = ChatOCIModelDeploymentVLLM(endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME)
-    assert llm._headers().get("route") == CONST_COMPLETION_ROUTE
-    output = llm.invoke(CONST_PROMPT)
-    assert isinstance(output, AIMessage)
-    assert output.content == CONST_COMPLETION
+@mock.patch("ads.aqua.client.client.HttpxOCIAuth")
+def test_authentication_with_ads(*args):
+    with mock.patch(
+        "ads.common.auth.default_signer", return_value=dict(signer=mock.MagicMock())
+    ) as ads_auth:
+        llm = ChatOCIModelDeployment(endpoint="<endpoint>", model="my-model")
+        ads_auth.assert_called_once()
+        assert llm.openai_api_base == "<endpoint>/v1"
+        assert llm.model_name == "my-model"
 
 
-@pytest.mark.requires("ads")
-@pytest.mark.requires("langchain_openai")
-@mock.patch("ads.common.auth.default_signer", return_value=dict(signer=None))
-@mock.patch("requests.post", side_effect=mocked_requests_post)
-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().get("route") == CONST_COMPLETION_ROUTE
-    output = llm.invoke(CONST_PROMPT)
-    assert isinstance(output, AIMessage)
-    assert output.content == CONST_COMPLETION
-
-
-@pytest.mark.requires("ads")
-@pytest.mark.requires("langchain_openai")
-@mock.patch("ads.common.auth.default_signer", return_value=dict(signer=None))
-@mock.patch("requests.post", side_effect=mocked_requests_post)
-def test_stream_vllm(*args: Any) -> None:
-    """Tests streaming with vLLM endpoint using OpenAI spec."""
-    llm = ChatOCIModelDeploymentVLLM(
-        endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME, streaming=True
-    )
-    assert llm._headers().get("route") == CONST_COMPLETION_ROUTE
-    output = None
-    count = 0
-    for chunk in llm.stream(CONST_PROMPT):
-        assert isinstance(chunk, AIMessageChunk)
-        if output is None:
-            output = chunk
-        else:
-            output += chunk
-        count += 1
-    assert count == 5
-    assert output is not None
-    if output is not None:
-        assert str(output.content).strip() == CONST_COMPLETION
-
-
-async def mocked_async_streaming_response(
-    *args: Any, **kwargs: Any
-) -> AsyncGenerator[bytes, None]:
-    """Returns mocked response for async streaming."""
-    for item in CONST_ASYNC_STREAM_RESPONSE:
-        yield item
-
-
-@pytest.mark.asyncio
-@pytest.mark.requires("ads")
 @pytest.mark.requires("langchain_openai")
-@mock.patch(
-    "ads.common.auth.default_signer", return_value=dict(signer=mock.MagicMock())
-)
-@mock.patch(
-    "langchain_oci.llms.oci_data_science_model_deployment_endpoint.BaseOCIModelDeployment._arequest",
-    mock.MagicMock(),
-)
-async def test_stream_async(*args: Any) -> None:
-    """Tests async streaming."""
-    llm = ChatOCIModelDeploymentVLLM(
-        endpoint=CONST_ENDPOINT, model=CONST_MODEL_NAME, streaming=True
-    )
-    assert llm._headers().get("route") == CONST_COMPLETION_ROUTE
-    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
+def test_authentication_with_custom_client():
+    http_client = httpx.Client()
+    llm = ChatOCIModelDeployment(base_url="<endpoint>/v1", http_client=http_client)
+    assert llm.http_client == http_client
+    assert llm.openai_api_base == "<endpoint>/v1"
+    assert llm.model_name == "odsc-llm"