mirror of
https://github.com/hoshikawa2/compass_backoffice.git
synced 2026-07-09 22:04:20 +00:00
156 lines
4.3 KiB
Python
156 lines
4.3 KiB
Python
"""Mock LLM implementations for testing without API calls."""
|
|
|
|
from typing import List, Optional, Any
|
|
from langchain_core.messages import AIMessage, BaseMessage
|
|
from langchain_core.language_models import BaseChatModel
|
|
from langchain_core.outputs import ChatGeneration, ChatResult
|
|
import pytest
|
|
|
|
|
|
class MockLLM(BaseChatModel):
|
|
"""
|
|
Mock LLM for testing without making actual API calls.
|
|
|
|
This mock can be configured with predefined responses and tracks
|
|
the number of times it's been called.
|
|
|
|
Example:
|
|
>>> mock = MockLLM(responses=["Hello!", "How can I help?"])
|
|
>>> response = await mock.ainvoke([HumanMessage(content="Hi")])
|
|
>>> assert response.content == "Hello!"
|
|
"""
|
|
|
|
responses: List[str]
|
|
call_count: int = 0
|
|
|
|
def __init__(self, responses: Optional[List[str]] = None, **kwargs):
|
|
"""
|
|
Initialize mock LLM with predefined responses.
|
|
|
|
Args:
|
|
responses: List of responses to return in sequence.
|
|
If None, returns "Mock response" for all calls.
|
|
"""
|
|
super().__init__(**kwargs)
|
|
self.responses = responses or ["Mock response"]
|
|
self.call_count = 0
|
|
|
|
@property
|
|
def _llm_type(self) -> str:
|
|
"""Return identifier for this LLM type."""
|
|
return "mock"
|
|
|
|
def _generate(
|
|
self,
|
|
messages: List[BaseMessage],
|
|
stop: Optional[List[str]] = None,
|
|
**kwargs: Any,
|
|
) -> ChatResult:
|
|
"""Generate a mock response synchronously."""
|
|
response = self.responses[self.call_count % len(self.responses)]
|
|
self.call_count += 1
|
|
|
|
message = AIMessage(content=response)
|
|
generation = ChatGeneration(message=message)
|
|
|
|
return ChatResult(generations=[generation])
|
|
|
|
async def _agenerate(
|
|
self,
|
|
messages: List[BaseMessage],
|
|
stop: Optional[List[str]] = None,
|
|
**kwargs: Any,
|
|
) -> ChatResult:
|
|
"""Generate a mock response asynchronously."""
|
|
return self._generate(messages, stop, **kwargs)
|
|
|
|
def reset(self):
|
|
"""Reset call count for reuse in tests."""
|
|
self.call_count = 0
|
|
|
|
|
|
class MockLLMWithError(BaseChatModel):
|
|
"""
|
|
Mock LLM that raises errors for testing error handling.
|
|
|
|
Example:
|
|
>>> mock = MockLLMWithError(error_message="API timeout")
|
|
>>> with pytest.raises(Exception):
|
|
... await mock.ainvoke([HumanMessage(content="Hi")])
|
|
"""
|
|
|
|
error_message: str
|
|
|
|
def __init__(self, error_message: str = "Mock LLM error", **kwargs):
|
|
"""
|
|
Initialize mock LLM that raises errors.
|
|
|
|
Args:
|
|
error_message: Error message to raise
|
|
"""
|
|
super().__init__(**kwargs)
|
|
self.error_message = error_message
|
|
|
|
@property
|
|
def _llm_type(self) -> str:
|
|
"""Return identifier for this LLM type."""
|
|
return "mock_error"
|
|
|
|
def _generate(
|
|
self,
|
|
messages: List[BaseMessage],
|
|
stop: Optional[List[str]] = None,
|
|
**kwargs: Any,
|
|
) -> ChatResult:
|
|
"""Raise an error."""
|
|
raise Exception(self.error_message)
|
|
|
|
async def _agenerate(
|
|
self,
|
|
messages: List[BaseMessage],
|
|
stop: Optional[List[str]] = None,
|
|
**kwargs: Any,
|
|
) -> ChatResult:
|
|
"""Raise an error asynchronously."""
|
|
raise Exception(self.error_message)
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_llm():
|
|
"""
|
|
Pytest fixture providing a basic mock LLM.
|
|
|
|
Returns:
|
|
MockLLM instance with default response
|
|
"""
|
|
return MockLLM()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_llm_with_responses():
|
|
"""
|
|
Pytest fixture factory for creating mock LLM with custom responses.
|
|
|
|
Returns:
|
|
Function that creates MockLLM with specified responses
|
|
|
|
Example:
|
|
def test_conversation(mock_llm_with_responses):
|
|
llm = mock_llm_with_responses(["Hi!", "Goodbye!"])
|
|
# Use llm in test
|
|
"""
|
|
def _create_mock(responses: List[str]) -> MockLLM:
|
|
return MockLLM(responses=responses)
|
|
return _create_mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_llm_with_error():
|
|
"""
|
|
Pytest fixture providing a mock LLM that raises errors.
|
|
|
|
Returns:
|
|
MockLLMWithError instance
|
|
"""
|
|
return MockLLMWithError()
|