test_chat_models.py — langchain Source File

Source Code

"""Ollama specific chat model integration tests"""

from __future__ import annotations

from typing import Annotated
from unittest.mock import MagicMock, patch

import pytest
from httpx import ConnectError
from langchain_core.messages.ai import AIMessage, AIMessageChunk
from langchain_core.messages.human import HumanMessage
from langchain_core.messages.tool import ToolCallChunk, ToolMessage
from langchain_core.tools import tool
from ollama import ResponseError
from pydantic import BaseModel, Field, ValidationError
from typing_extensions import TypedDict

from langchain_ollama import ChatOllama

DEFAULT_MODEL_NAME = "llama3.1"
REASONING_MODEL_NAME = "gpt-oss:20b"


@tool
def get_current_weather(location: str) -> dict:
    """Gets the current weather in a given location."""
    if "boston" in location.lower():
        return {"temperature": "15°F", "conditions": "snow"}
    return {"temperature": "unknown", "conditions": "unknown"}


@patch("langchain_ollama.chat_models.Client.list")
def test_init_model_not_found(mock_list: MagicMock) -> None:
    """Test that a ValueError is raised when the model is not found."""
    mock_list.side_effect = ValueError("Test model not found")
    with pytest.raises(ValueError) as excinfo:
        ChatOllama(model="non-existent-model", validate_model_on_init=True)
    assert "Test model not found" in str(excinfo.value)


@patch("langchain_ollama.chat_models.Client.list")
def test_init_connection_error(mock_list: MagicMock) -> None:
    """Test that a `ValidationError` is raised on connect failure during init."""
    mock_list.side_effect = ConnectError("Test connection error")

    with pytest.raises(ValidationError) as excinfo:
        ChatOllama(model="any-model", validate_model_on_init=True)
    assert "Failed to connect to Ollama" in str(excinfo.value)


@patch("langchain_ollama.chat_models.Client.list")
def test_init_response_error(mock_list: MagicMock) -> None:
    """Test that a ResponseError is raised."""
    mock_list.side_effect = ResponseError("Test response error")

    with pytest.raises(ValidationError) as excinfo:
        ChatOllama(model="any-model", validate_model_on_init=True)
    assert "Received an error from the Ollama API" in str(excinfo.value)


// ... (236 more lines)