68 lines
2.0 KiB
Python
68 lines
2.0 KiB
Python
"""Stub LLM adapter for tests; returns fixed responses."""
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
from fusionagi.adapters.base import LLMAdapter
|
|
|
|
|
|
class StubAdapter(LLMAdapter):
|
|
"""
|
|
Returns configurable fixed responses; no API calls.
|
|
|
|
Useful for testing without making actual LLM API calls.
|
|
Supports both text and structured (JSON) responses.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
response: str = "Stub response",
|
|
structured_response: dict[str, Any] | list[Any] | None = None,
|
|
) -> None:
|
|
"""
|
|
Initialize the stub adapter.
|
|
|
|
Args:
|
|
response: Fixed text response for complete().
|
|
structured_response: Fixed structured response for complete_structured().
|
|
"""
|
|
self._response = response
|
|
self._structured_response = structured_response
|
|
|
|
def complete(
|
|
self,
|
|
messages: list[dict[str, str]],
|
|
**kwargs: Any,
|
|
) -> str:
|
|
"""Return the configured stub response."""
|
|
return self._response
|
|
|
|
def complete_structured(
|
|
self,
|
|
messages: list[dict[str, str]],
|
|
schema: dict[str, Any] | None = None,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
"""
|
|
Return the configured structured response.
|
|
|
|
If no structured_response was configured, attempts to parse
|
|
the text response as JSON, or returns None.
|
|
"""
|
|
if self._structured_response is not None:
|
|
return self._structured_response
|
|
|
|
# Try to parse text response as JSON
|
|
try:
|
|
return json.loads(self._response)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
|
|
def set_response(self, response: str) -> None:
|
|
"""Update the text response (useful for test scenarios)."""
|
|
self._response = response
|
|
|
|
def set_structured_response(self, response: dict[str, Any] | list[Any] | None) -> None:
|
|
"""Update the structured response (useful for test scenarios)."""
|
|
self._structured_response = response
|