Code / services/llm_provider_service/internal_models.py
services/llm_provider_service/internal_models.py
1
"""Internal Pydantic models for LLM Provider Service."""
2
3
from datetime import datetime, timezone
4
from enum import StrEnum
5
from typing import Any, Dict, Optional
6
from uuid import UUID
7
8
from pydantic import BaseModel, Field, model_validator
9
10
from common_core import EssayComparisonWinner, LLMProviderType
11
from common_core.error_enums import ErrorCode
12
from common_core.events.llm_provider_events import ProviderReportedUsageV1
13
from services.llm_provider_service.error_types import LLMProviderErrorType
14
15
16
class LLMProviderResponse(BaseModel):
17
"""Internal model for LLM provider responses."""
18
19
# Core comparison result
20
winner: EssayComparisonWinner = Field(description="Selected essay")
21
justification: str | None = Field(
22
default=None, description="Optional justification for the choice"
23
)
24
confidence: float = Field(ge=0.0, le=1.0, description="Confidence score")
25
26
# Provider metadata
27
provider: LLMProviderType = Field(description="Provider that generated the response")
28
model: str = Field(description="Model used for generation")
29
30
# Performance metrics
31
prompt_tokens: int = Field(ge=0, description="Provider-reported prompt tokens")
32
completion_tokens: int = Field(ge=0, description="Provider-reported completion tokens")
33
total_tokens: int = Field(ge=0, description="Provider-reported total tokens")
34
provider_reported_usage: ProviderReportedUsageV1 = Field(
35
default_factory=ProviderReportedUsageV1,
36
description="Actual usage and cost fields supplied by the provider",
37
)
38
# Additional metadata
39
raw_response: Dict[str, Any] = Field(default_factory=dict, description="Raw provider response")
40
metadata: Dict[str, Any] = Field(
41
default_factory=dict, description="Additional provider metadata"
42
)
43
44
@model_validator(mode="after")
45
def validate_provider_reported_usage(self) -> "LLMProviderResponse":
46
"""Require completed usage evidence from non-mock providers."""
47
if self.provider is LLMProviderType.MOCK:
48
return self
49
50
reported_usage = self.provider_reported_usage
51
if (
52
reported_usage.prompt_tokens is None
53
or reported_usage.completion_tokens is None
54
or reported_usage.total_tokens is None
55
):
56
raise ValueError("Non-mock provider responses require complete reported token usage")
57
58
if (
59
reported_usage.prompt_tokens,
60
reported_usage.completion_tokens,
61
reported_usage.total_tokens,
62
) != (
63
self.prompt_tokens,
64
self.completion_tokens,
65
self.total_tokens,
66
):
67
raise ValueError("Provider-reported token usage must match response token counts")
68
69
return self
70
71
72
class LLMOrchestratorResponse(BaseModel):
73
"""Internal model for orchestrator responses."""
74
75
# Core result (from provider)
76
winner: EssayComparisonWinner = Field(description="Selected essay")
77
justification: str | None = Field(
78
default=None, description="Optional justification for the choice"
79
)
80
confidence: float = Field(ge=0.0, le=1.0, description="Confidence score")
81
82
# Provider information
83
provider: LLMProviderType = Field(description="Actual provider used")
84
model: str = Field(description="Actual model used")
85
86
# Performance metrics
87
response_time_ms: int = Field(ge=0, description="Total response time in milliseconds")
88
89
# Token usage and cost
90
token_usage: Dict[str, int] = Field(
91
description="Provider-reported token usage",
92
)
93
cost_estimate: float = Field(
94
ge=0.0,
95
description="Actual provider charge or returned-usage manifest cost in USD",
96
)
97
provider_reported_usage: ProviderReportedUsageV1 = Field(
98
default_factory=ProviderReportedUsageV1,
99
description="Actual usage and cost fields supplied by the provider",
100
)
101
102
# Tracing
103
correlation_id: UUID = Field(description="Request correlation ID")
104
trace_id: str | None = Field(default=None, description="Distributed trace ID")
105
106
# Additional metadata
107
metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
108
109
110
class LLMProviderError(BaseModel):
111
"""Model for provider error responses."""
112
113
error_type: ErrorCode | LLMProviderErrorType = Field(description="Type of error")
114
message: str = Field(description="Error message")
115
provider: LLMProviderType = Field(description="Provider that failed")
116
correlation_id: UUID = Field(description="Request correlation ID")
117
retry_after: int | None = Field(default=None, description="Seconds to wait before retry")
118
is_retryable: bool = Field(default=True, description="Whether error is retryable")
119
120
121
class LLMQueuedResult(BaseModel):
122
"""Model for queued LLM request result."""
123
124
queue_id: UUID = Field(description="Unique queue identifier")
125
correlation_id: UUID = Field(description="Request correlation ID")
126
provider: LLMProviderType = Field(description="Requested provider")
127
status: str = Field(default="queued", description="Request status")
128
estimated_wait_minutes: int | None = Field(default=None, description="Estimated wait time")
129
priority: int = Field(description="Request priority")
130
queued_at: str = Field(description="ISO timestamp when queued")
131
132
133
class BatchComparisonItem(BaseModel):
134
"""Container for batch comparison processing input."""
135
136
provider: LLMProviderType = Field(description="Provider requested for each comparison")
137
user_prompt: str = Field(description="Fully rendered comparison prompt")
138
prompt_blocks: list[dict[str, Any]] | None = Field(
139
default=None,
140
description="Structured prompt blocks (preferred for caching-capable providers)",
141
)
142
correlation_id: UUID = Field(description="Per-comparison correlation identifier")
143
overrides: Optional[Dict[str, Any]] = Field(
144
default=None,
145
description="Optional override kwargs forwarded to process_comparison",
146
)
147
prefer_null_justification: bool = Field(
148
default=False,
149
description="Whether this bundle should ask the provider to return null justification",
150
)
151
152
153
class BatchJobStatus(StrEnum):
154
"""Lifecycle status for provider batch jobs."""
155
156
PENDING = "pending"
157
SCHEDULED = "scheduled"
158
RUNNING = "running"
159
COMPLETED = "completed"
160
FAILED = "failed"
161
CANCELLED = "cancelled"
162
163
164
class BatchJobItemStatus(StrEnum):
165
"""Outcome status for individual batch job items."""
166
167
SUCCESS = "success"
168
FAILED = "failed"
169
170
171
class BatchJobItem(BaseModel):
172
"""Link a queued request to a provider-native batch job item."""
173
174
queue_id: UUID = Field(description="Queue identifier for the originating request")
175
provider: LLMProviderType = Field(description="Provider responsible for this item")
176
model: str = Field(description="Model identifier used for this item")
177
custom_id: str = Field(
178
description=(
179
"Provider-facing custom identifier used to correlate job results back to queue items"
180
)
181
)
182
user_prompt: str = Field(description="Rendered comparison prompt for this item")
183
prompt_blocks: list[dict[str, Any]] | None = Field(
184
default=None,
185
description="Structured prompt blocks used when constructing provider payloads",
186
)
187
correlation_id: UUID = Field(
188
description="Correlation identifier propagated from the original comparison request"
189
)
190
overrides: Dict[str, Any] = Field(
191
default_factory=dict,
192
description="Resolved override kwargs (model, temperature, tokens, etc.)",
193
)
194
metadata: Dict[str, Any] = Field(
195
default_factory=dict,
196
description="Caller-supplied metadata to preserve through provider batch APIs",
197
)
198
199
200
class BatchJobResult(BaseModel):
201
"""Per-item outcome for a provider batch job."""
202
203
queue_id: UUID = Field(description="Queue identifier for the originating request")
204
provider: LLMProviderType = Field(description="Provider that executed the item")
205
model: str = Field(description="Model used for the item")
206
status: BatchJobItemStatus = Field(description="Item-level outcome status")
207
response: Optional["LLMOrchestratorResponse"] = Field(
208
default=None,
209
description="Normalized LLM response when the item succeeds",
210
)
211
error_code: Optional[ErrorCode] = Field(
212
default=None,
213
description="Optional error code when the item fails",
214
)
215
error_message: Optional[str] = Field(
216
default=None,
217
description="Human-readable error message when the item fails",
218
)
219
raw_error: Dict[str, Any] = Field(
220
default_factory=dict,
221
description="Provider- or implementation-specific error payload for diagnostics",
222
)
223
224
225
class BatchJobRef(BaseModel):
226
"""Handle for tracking provider batch job lifecycle."""
227
228
job_id: UUID = Field(description="Internal identifier for the batch job")
229
provider: LLMProviderType = Field(description="Provider responsible for the job")
230
model: str = Field(description="Model used for the job")
231
status: BatchJobStatus = Field(description="Current lifecycle status for the job")
232
provider_job_id: Optional[str] = Field(
233
default=None,
234
description="Identifier returned by the provider for this job, if applicable",
235
)
236
created_at: datetime = Field(
237
default_factory=lambda: datetime.now(timezone.utc),
238
description="Timestamp when the job was created",
239
)
240
completed_at: Optional[datetime] = Field(
241
default=None,
242
description="Timestamp when the job reached a terminal state",
243
)