diff --git a/.github/workflows/scheduled_test.yml b/.github/workflows/scheduled_test.yml index d59ad3c2e3..44de47155c 100644 --- a/.github/workflows/scheduled_test.yml +++ b/.github/workflows/scheduled_test.yml @@ -61,6 +61,10 @@ jobs: env: OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} + AZURE_OPENAI_API_VERSION: ${{ secrets.AZURE_OPENAI_API_VERSION }} + AZURE_OPENAI_API_BASE: ${{ secrets.AZURE_OPENAI_API_BASE }} + AZURE_OPENAI_API_KEY: ${{ secrets.AZURE_OPENAI_API_KEY }} + AZURE_OPENAI_DEPLOYMENT_NAME: ${{ secrets.AZURE_OPENAI_DEPLOYMENT_NAME }} run: | make scheduled_tests diff --git a/libs/langchain/tests/integration_tests/chat_models/test_azure_openai.py b/libs/langchain/tests/integration_tests/chat_models/test_azure_openai.py new file mode 100644 index 0000000000..137c35342b --- /dev/null +++ b/libs/langchain/tests/integration_tests/chat_models/test_azure_openai.py @@ -0,0 +1,228 @@ +"""Test AzureChatOpenAI wrapper.""" +import os +from typing import Any + +import pytest + +from langchain.callbacks.manager import CallbackManager +from langchain.chat_models import AzureChatOpenAI +from langchain.schema import ( + ChatGeneration, + ChatResult, + LLMResult, +) +from langchain.schema.messages import BaseMessage, HumanMessage +from tests.unit_tests.callbacks.fake_callback_handler import FakeCallbackHandler + +OPENAI_API_VERSION = os.environ.get("AZURE_OPENAI_API_VERSION", "") +OPENAI_API_BASE = os.environ.get("AZURE_OPENAI_API_BASE", "") +OPENAI_API_KEY = os.environ.get("AZURE_OPENAI_API_KEY", "") +DEPLOYMENT_NAME = os.environ.get("AZURE_OPENAI_DEPLOYMENT_NAME", "") + + +def _get_llm(**kwargs: Any) -> AzureChatOpenAI: + return AzureChatOpenAI( + deployment_name=DEPLOYMENT_NAME, + openai_api_version=OPENAI_API_VERSION, + openai_api_base=OPENAI_API_BASE, + openai_api_key=OPENAI_API_KEY, + **kwargs, + ) + + +@pytest.mark.scheduled +@pytest.fixture +def llm() -> AzureChatOpenAI: + return _get_llm( + max_tokens=10, + ) + + +def test_chat_openai(llm: AzureChatOpenAI) -> None: + """Test AzureChatOpenAI wrapper.""" + message = HumanMessage(content="Hello") + response = llm([message]) + assert isinstance(response, BaseMessage) + assert isinstance(response.content, str) + + +@pytest.mark.scheduled +def test_chat_openai_generate() -> None: + """Test AzureChatOpenAI wrapper with generate.""" + chat = _get_llm(max_tokens=10, n=2) + message = HumanMessage(content="Hello") + response = chat.generate([[message], [message]]) + assert isinstance(response, LLMResult) + assert len(response.generations) == 2 + for generations in response.generations: + assert len(generations) == 2 + for generation in generations: + assert isinstance(generation, ChatGeneration) + assert isinstance(generation.text, str) + assert generation.text == generation.message.content + + +@pytest.mark.scheduled +def test_chat_openai_multiple_completions() -> None: + """Test AzureChatOpenAI wrapper with multiple completions.""" + chat = _get_llm(max_tokens=10, n=5) + message = HumanMessage(content="Hello") + response = chat._generate([message]) + assert isinstance(response, ChatResult) + assert len(response.generations) == 5 + for generation in response.generations: + assert isinstance(generation.message, BaseMessage) + assert isinstance(generation.message.content, str) + + +@pytest.mark.scheduled +def test_chat_openai_streaming() -> None: + """Test that streaming correctly invokes on_llm_new_token callback.""" + callback_handler = FakeCallbackHandler() + callback_manager = CallbackManager([callback_handler]) + chat = _get_llm( + max_tokens=10, + streaming=True, + temperature=0, + callback_manager=callback_manager, + verbose=True, + ) + message = HumanMessage(content="Hello") + response = chat([message]) + assert callback_handler.llm_streams > 0 + assert isinstance(response, BaseMessage) + + +@pytest.mark.scheduled +def test_chat_openai_streaming_generation_info() -> None: + """Test that generation info is preserved when streaming.""" + + class _FakeCallback(FakeCallbackHandler): + saved_things: dict = {} + + def on_llm_end( + self, + *args: Any, + **kwargs: Any, + ) -> Any: + # Save the generation + self.saved_things["generation"] = args[0] + + callback = _FakeCallback() + callback_manager = CallbackManager([callback]) + chat = _get_llm( + max_tokens=2, + temperature=0, + callback_manager=callback_manager, + ) + list(chat.stream("hi")) + generation = callback.saved_things["generation"] + # `Hello!` is two tokens, assert that that is what is returned + assert generation.generations[0][0].text == "Hello!" + + +@pytest.mark.scheduled +@pytest.mark.asyncio +async def test_async_chat_openai() -> None: + """Test async generation.""" + chat = _get_llm(max_tokens=10, n=2) + message = HumanMessage(content="Hello") + response = await chat.agenerate([[message], [message]]) + assert isinstance(response, LLMResult) + assert len(response.generations) == 2 + for generations in response.generations: + assert len(generations) == 2 + for generation in generations: + assert isinstance(generation, ChatGeneration) + assert isinstance(generation.text, str) + assert generation.text == generation.message.content + + +@pytest.mark.scheduled +@pytest.mark.asyncio +async def test_async_chat_openai_streaming() -> None: + """Test that streaming correctly invokes on_llm_new_token callback.""" + callback_handler = FakeCallbackHandler() + callback_manager = CallbackManager([callback_handler]) + chat = _get_llm( + max_tokens=10, + streaming=True, + temperature=0, + callback_manager=callback_manager, + verbose=True, + ) + message = HumanMessage(content="Hello") + response = await chat.agenerate([[message], [message]]) + assert callback_handler.llm_streams > 0 + assert isinstance(response, LLMResult) + assert len(response.generations) == 2 + for generations in response.generations: + assert len(generations) == 1 + for generation in generations: + assert isinstance(generation, ChatGeneration) + assert isinstance(generation.text, str) + assert generation.text == generation.message.content + + +@pytest.mark.scheduled +def test_openai_streaming(llm: AzureChatOpenAI) -> None: + """Test streaming tokens from OpenAI.""" + + for token in llm.stream("I'm Pickle Rick"): + assert isinstance(token.content, str) + + +@pytest.mark.scheduled +@pytest.mark.asyncio +async def test_openai_astream(llm: AzureChatOpenAI) -> None: + """Test streaming tokens from OpenAI.""" + async for token in llm.astream("I'm Pickle Rick"): + assert isinstance(token.content, str) + + +@pytest.mark.scheduled +@pytest.mark.asyncio +async def test_openai_abatch(llm: AzureChatOpenAI) -> None: + """Test streaming tokens from AzureChatOpenAI.""" + + result = await llm.abatch(["I'm Pickle Rick", "I'm not Pickle Rick"]) + for token in result: + assert isinstance(token.content, str) + + +@pytest.mark.scheduled +@pytest.mark.asyncio +async def test_openai_abatch_tags(llm: AzureChatOpenAI) -> None: + """Test batch tokens from AzureChatOpenAI.""" + + result = await llm.abatch( + ["I'm Pickle Rick", "I'm not Pickle Rick"], config={"tags": ["foo"]} + ) + for token in result: + assert isinstance(token.content, str) + + +@pytest.mark.scheduled +def test_openai_batch(llm: AzureChatOpenAI) -> None: + """Test batch tokens from AzureChatOpenAI.""" + + result = llm.batch(["I'm Pickle Rick", "I'm not Pickle Rick"]) + for token in result: + assert isinstance(token.content, str) + + +@pytest.mark.scheduled +@pytest.mark.asyncio +async def test_openai_ainvoke(llm: AzureChatOpenAI) -> None: + """Test invoke tokens from AzureChatOpenAI.""" + + result = await llm.ainvoke("I'm Pickle Rick", config={"tags": ["foo"]}) + assert isinstance(result.content, str) + + +@pytest.mark.scheduled +def test_openai_invoke(llm: AzureChatOpenAI) -> None: + """Test invoke tokens from AzureChatOpenAI.""" + + result = llm.invoke("I'm Pickle Rick", config=dict(tags=["foo"])) + assert isinstance(result.content, str)