From 3121fd7dc0106f6c3896769395b3e6a74c7e6979 Mon Sep 17 00:00:00 2001 From: Chester Curme Date: Mon, 6 Jan 2025 11:51:37 -0500 Subject: [PATCH] test azure --- .../chat_models/test_azure.py | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/libs/partners/openai/tests/integration_tests/chat_models/test_azure.py b/libs/partners/openai/tests/integration_tests/chat_models/test_azure.py index 4ed531ad119ff..e9228df0730c5 100644 --- a/libs/partners/openai/tests/integration_tests/chat_models/test_azure.py +++ b/libs/partners/openai/tests/integration_tests/chat_models/test_azure.py @@ -13,6 +13,7 @@ HumanMessage, ) from langchain_core.outputs import ChatGeneration, ChatResult, LLMResult +from pydantic import BaseModel from langchain_openai import AzureChatOpenAI from tests.unit_tests.fake.callbacks import FakeCallbackHandler @@ -262,3 +263,37 @@ async def test_json_mode_async(llm: AzureChatOpenAI) -> None: assert isinstance(full, AIMessageChunk) assert isinstance(full.content, str) assert json.loads(full.content) == {"a": 1} + + +class Foo(BaseModel): + response: str + + +def test_stream_response_format(llm: AzureChatOpenAI) -> None: + full: Optional[BaseMessageChunk] = None + chunks = [] + for chunk in llm.stream("how are ya", response_format=Foo): + chunks.append(chunk) + full = chunk if full is None else full + chunk + assert len(chunks) > 1 + assert isinstance(full, AIMessageChunk) + parsed = full.additional_kwargs["parsed"] + assert isinstance(parsed, Foo) + assert isinstance(full.content, str) + parsed_content = json.loads(full.content) + assert parsed.response == parsed_content["response"] + + +async def test_astream_response_format(llm: AzureChatOpenAI) -> None: + full: Optional[BaseMessageChunk] = None + chunks = [] + async for chunk in llm.astream("how are ya", response_format=Foo): + chunks.append(chunk) + full = chunk if full is None else full + chunk + assert len(chunks) > 1 + assert isinstance(full, AIMessageChunk) + parsed = full.additional_kwargs["parsed"] + assert isinstance(parsed, Foo) + assert isinstance(full.content, str) + parsed_content = json.loads(full.content) + assert parsed.response == parsed_content["response"]