File size: 2,028 Bytes
d840c10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """Deterministic fake model client for unit tests."""
from __future__ import annotations
from collections.abc import Iterable
from typing import TypeAlias
from pydantic import BaseModel, ValidationError
from gcmd_classifier.errors import NonRetryableModelError, StructuredModelResponseError
from gcmd_classifier.llm.base import ModelRequest, ModelResponse, StructuredResponseT
ScriptedAction: TypeAlias = BaseModel | dict | Exception
class FakeModelClient:
"""Scripted model client that never imports or calls a live provider."""
def __init__(self, actions: Iterable[ScriptedAction]) -> None:
self._actions = list(actions)
self.requests: list[ModelRequest] = []
def generate_structured(
self,
request: ModelRequest[StructuredResponseT],
) -> ModelResponse[StructuredResponseT]:
"""Return the next scripted response or raise the next scripted exception."""
self.requests.append(request)
if not self._actions:
raise NonRetryableModelError("Fake model script is exhausted.")
action = self._actions.pop(0)
if isinstance(action, Exception):
raise action
try:
parsed = _parse_action(action, request.response_schema)
except ValidationError as exc:
raise StructuredModelResponseError(
"Fake model response failed schema validation."
) from exc
return ModelResponse(
parsed=parsed,
provider=request.provider,
model_name=request.model_name,
prompt_version=request.prompt_version,
)
@property
def remaining_actions(self) -> int:
"""Number of scripted actions that have not been consumed."""
return len(self._actions)
def _parse_action(
action: BaseModel | dict,
response_schema: type[StructuredResponseT],
) -> StructuredResponseT:
if isinstance(action, response_schema):
return action
return response_schema.model_validate(action)
|