test_azure.py — langchain Source File

Source Code

"""Test AzureOpenAI wrapper."""

import os
from collections.abc import Generator
from typing import Any

import pytest
from langchain_core.callbacks import CallbackManager
from langchain_core.outputs import LLMResult

from langchain_openai import AzureOpenAI
from tests.unit_tests.fake.callbacks 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",
    os.environ.get("AZURE_OPENAI_LLM_DEPLOYMENT_NAME", ""),
)

pytestmark = pytest.mark.skipif(
    True,
    reason=(
        "This entire module is skipped as all Azure OpenAI models supporting text "
        "completions are retired. See: "
        "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/legacy-models"
    ),
)


def _get_llm(**kwargs: Any) -> AzureOpenAI:
    return AzureOpenAI(  # type: ignore[call-arg, call-arg, call-arg]
        deployment_name=DEPLOYMENT_NAME,
        openai_api_version=OPENAI_API_VERSION,
        azure_endpoint=OPENAI_API_BASE,
        openai_api_key=OPENAI_API_KEY,
        **kwargs,
    )


@pytest.fixture
def llm() -> AzureOpenAI:
    return _get_llm(max_tokens=10)


@pytest.mark.scheduled
def test_openai_call(llm: AzureOpenAI) -> None:
    """Test valid call to openai."""
    output = llm.invoke("Say something nice:")
    assert isinstance(output, str)


@pytest.mark.scheduled
def test_openai_streaming(llm: AzureOpenAI) -> None:
    """Test streaming tokens from AzureOpenAI."""
    generator = llm.stream("I'm Pickle Rick")

    assert isinstance(generator, Generator)

// ... (126 more lines)