Spaces:
Running on Zero
Running on Zero
File size: 4,197 Bytes
d840c10 1163af8 d840c10 1163af8 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """Provider-neutral model interface and retry helpers."""
from __future__ import annotations
import time
from collections.abc import Mapping
from dataclasses import dataclass, field, replace
from enum import Enum
from typing import Any, Generic, Protocol, TypeVar
from pydantic import BaseModel
from gcmd_classifier.config import ModelSettings
from gcmd_classifier.errors import ModelRetriesExhaustedError, RetryableModelError
StructuredResponseT = TypeVar("StructuredResponseT", bound=BaseModel)
class ModelStage(str, Enum): # noqa: UP042
"""Supported model-call stages for the MVP classifier."""
TOPIC = "topic"
TERM = "term"
VARIABLE = "variable"
@dataclass(frozen=True)
class TokenUsage:
"""Token usage metadata returned by a provider when available."""
input_tokens: int | None = None
output_tokens: int | None = None
total_tokens: int | None = None
@dataclass(frozen=True)
class ModelRequest(Generic[StructuredResponseT]):
"""Provider-neutral structured model request."""
stage: ModelStage
prompt: str
response_schema: type[StructuredResponseT]
provider: str
model_name: str
prompt_version: str
temperature: float
timeout_seconds: float
metadata: Mapping[str, Any] = field(default_factory=dict)
@classmethod
def from_settings(
cls,
*,
stage: ModelStage,
prompt: str,
response_schema: type[StructuredResponseT],
settings: ModelSettings,
metadata: Mapping[str, Any] | None = None,
) -> ModelRequest[StructuredResponseT]:
"""Create a request from model settings for a specific stage."""
return cls(
stage=stage,
prompt=prompt,
response_schema=response_schema,
provider=settings.provider,
model_name=settings.model_name,
prompt_version=settings.prompt_version_for_stage(stage.value),
temperature=settings.temperature,
timeout_seconds=settings.timeout_seconds,
metadata={} if metadata is None else dict(metadata),
)
@dataclass(frozen=True)
class ModelResponse(Generic[StructuredResponseT]):
"""Provider-neutral structured model response."""
parsed: StructuredResponseT
provider: str
model_name: str
prompt_version: str
retry_count: int = 0
duration_seconds: float | None = None
token_usage: TokenUsage | None = None
estimated_cost: float | None = None
raw_response: str | None = None
warnings: tuple[str, ...] = ()
diagnostics: Mapping[str, Any] = field(default_factory=dict)
class ModelClient(Protocol):
"""Provider-neutral interface used by future classification components."""
def generate_structured(
self,
request: ModelRequest[StructuredResponseT],
) -> ModelResponse[StructuredResponseT]:
"""Return a typed structured response for a model request."""
@dataclass(frozen=True)
class RetryPolicy:
"""Retry settings for temporary model failures."""
max_retries: int = 2
@classmethod
def from_settings(cls, settings: ModelSettings) -> RetryPolicy:
"""Create a retry policy from model settings."""
return cls(max_retries=settings.max_retries)
def generate_with_retries(
client: ModelClient,
request: ModelRequest[StructuredResponseT],
policy: RetryPolicy,
) -> ModelResponse[StructuredResponseT]:
"""Call a model client, retrying only explicitly retryable failures."""
retry_count = 0
started = time.perf_counter()
while True:
try:
response = client.generate_structured(request)
duration = response.duration_seconds
if duration is None:
duration = time.perf_counter() - started
return replace(response, retry_count=retry_count, duration_seconds=duration)
except RetryableModelError as exc:
if retry_count >= policy.max_retries:
raise ModelRetriesExhaustedError(
"Retryable model failure exceeded configured max retries.",
retry_count=retry_count,
) from exc
retry_count += 1
|