Code / services/cj_assessment_service/protocols.py
services/cj_assessment_service/protocols.py
1
"""CJ client and orchestration protocols owned by the assessment service.
2
3
Persistence, completion-event, and pair-selection contracts are defined in
4
their focused modules and imported directly by their consumers.
5
"""
6
7
from __future__ import annotations
8
9
from typing import (
10
TYPE_CHECKING,
11
Awaitable,
12
Callable,
13
Protocol,
14
TypeVar,
15
)
16
from uuid import UUID
17
18
from pydantic import JsonValue
19
20
from common_core import LLMProviderType
21
22
if TYPE_CHECKING:
23
from services.cj_assessment_service.cj_core_logic.batch_config import BatchConfigOverrides
24
from services.cj_assessment_service.cj_core_logic.batch_submission import BatchSubmissionResult
25
from services.cj_assessment_service.models_api import (
26
CJAssessmentRequestData,
27
ComparisonAdmissionResult,
28
ComparisonTask,
29
ComparisonTrackingKey,
30
LLMProviderAdmission,
31
)
32
from services.cj_assessment_service.pair_selection_protocols import PairGenerationMode
33
34
T = TypeVar("T")
35
36
__all__ = [
37
"AtomicRedisClientProtocol",
38
"BatchCompletionCheckerProtocol",
39
"BatchProcessorProtocol",
40
"BatchRetryProcessorProtocol",
41
"ContentClientProtocol",
42
"LLMInteractionProtocol",
43
"LLMProviderProtocol",
44
"RetryManagerProtocol",
45
]
46
47
48
class ContentClientProtocol(Protocol):
49
"""Protocol for fetching spellchecked essay text from content storage API."""
50
51
async def fetch_content(self, storage_id: str, correlation_id: UUID) -> str:
52
"""Fetch essay text content by storage ID.
53
54
Args:
55
storage_id: The storage reference ID for the essay text
56
correlation_id: Request correlation ID for tracing
57
58
Returns:
59
The essay text content
60
61
Raises:
62
HuleEduError: On any failure to fetch content
63
"""
64
...
65
66
async def store_content(self, content: str, content_type: str = "text/plain") -> dict[str, str]:
67
"""Store content in Content storage API.
68
69
Args:
70
content: The text content to store
71
content_type: MIME type of content
72
73
Returns:
74
Dict with 'content_id' key containing the storage ID
75
"""
76
...
77
78
79
class RetryManagerProtocol(Protocol):
80
"""Protocol for managing LLM API retry logic."""
81
82
async def with_retry(
83
self,
84
operation: Callable[[], Awaitable[T]],
85
*,
86
provider_name: str = "unknown",
87
correlation_id: UUID | None = None,
88
) -> T:
89
"""Execute operation with retry logic.
90
91
Args:
92
operation: Async callable to execute with retries
93
*args: Positional arguments for the operation
94
**kwargs: Keyword arguments for the operation
95
96
Returns:
97
The result from the successful operation
98
99
Raises:
100
HuleEduError: On permanent failure after all retries exhausted
101
"""
102
...
103
104
105
class LLMProviderProtocol(Protocol):
106
"""Protocol for individual LLM provider implementations."""
107
108
async def generate_comparison(
109
self,
110
user_prompt: str,
111
correlation_id: UUID,
112
prompt_blocks: list[dict[str, JsonValue]] | None = None,
113
system_prompt_override: str | None = None,
114
model_override: str | None = None,
115
temperature_override: float | None = None,
116
max_tokens_override: int | None = None,
117
provider_override: str | None = None,
118
request_metadata: dict[str, JsonValue] | None = None,
119
) -> "LLMProviderAdmission":
120
"""Generate a comparison assessment using the LLM.
121
122
Args:
123
user_prompt: The user prompt containing essay comparison request
124
correlation_id: Request correlation ID for tracing
125
system_prompt_override: Optional system prompt override
126
model_override: Optional model name override
127
temperature_override: Optional temperature override (0.0-2.0)
128
max_tokens_override: Optional max tokens override
129
provider_override: Optional provider name override
130
request_metadata: Arbitrary metadata to echo back in callbacks
131
132
Returns:
133
The LLM response data containing the comparison result
134
135
Raises:
136
HuleEduError: On any failure to generate comparison
137
"""
138
...
139
140
141
class LLMInteractionProtocol(Protocol):
142
"""Protocol for performing LLM-based essay comparisons."""
143
144
async def perform_comparisons(
145
self,
146
tasks: list["ComparisonTask"],
147
correlation_id: UUID,
148
tracking_map: dict["ComparisonTrackingKey", UUID] | None = None,
149
bos_batch_id: str | None = None,
150
model_override: str | None = None,
151
temperature_override: float | None = None,
152
max_tokens_override: int | None = None,
153
system_prompt_override: str | None = None,
154
provider_override: str | LLMProviderType | None = None,
155
metadata_context: dict[str, JsonValue] | None = None,
156
) -> list["ComparisonAdmissionResult"]:
157
"""Perform comparative judgment on a list of comparison tasks.
158
159
Args:
160
tasks: List of ComparisonTask objects
161
correlation_id: Request correlation ID for tracing
162
tracking_map: Optional per-pair correlation override map
163
bos_batch_id: Optional BOS batch identifier propagated to metadata
164
model_override: Optional model name override
165
temperature_override: Optional temperature override (0.0-2.0)
166
max_tokens_override: Optional max tokens override
167
system_prompt_override: Optional system prompt override
168
provider_override: Optional provider name override forwarded to LPS
169
metadata_context: Optional metadata additions for downstream callbacks
170
171
Returns:
172
One admission outcome for every requested comparison task.
173
"""
174
...
175
176
177
class BatchProcessorProtocol(Protocol):
178
"""Protocol for core batch submission orchestration."""
179
180
async def submit_comparison_batch(
181
self,
182
cj_batch_id: int,
183
comparison_tasks: list["ComparisonTask"],
184
correlation_id: UUID,
185
config_overrides: "BatchConfigOverrides | None" = None,
186
model_override: str | None = None,
187
temperature_override: float | None = None,
188
max_tokens_override: int | None = None,
189
system_prompt_override: str | None = None,
190
provider_override: str | LLMProviderType | None = None,
191
metadata_context: dict[str, JsonValue] | None = None,
192
pair_generation_mode: "PairGenerationMode | None" = None,
193
) -> "BatchSubmissionResult":
194
"""Submit comparison batch with configurable batch size."""
195
...
196
197
async def handle_batch_submission(
198
self,
199
cj_batch_id: int,
200
comparison_tasks: list["ComparisonTask"],
201
correlation_id: UUID,
202
request_data: "CJAssessmentRequestData",
203
) -> "BatchSubmissionResult":
204
"""Handle batch submission with state tracking."""
205
...
206
207
208
class BatchCompletionCheckerProtocol(Protocol):
209
"""Protocol for batch completion evaluation and threshold checking."""
210
211
async def check_batch_completion(
212
self,
213
cj_batch_id: int,
214
correlation_id: UUID,
215
config_overrides: "BatchConfigOverrides | None" = None,
216
) -> bool:
217
"""Check if batch is complete or has reached threshold."""
218
...
219
220
221
class BatchRetryProcessorProtocol(Protocol):
222
"""Protocol for retry batch processing and end-of-batch fairness logic."""
223
224
async def submit_retry_batch(
225
self,
226
cj_batch_id: int,
227
correlation_id: UUID,
228
force_retry_all: bool = False,
229
model_override: str | None = None,
230
temperature_override: float | None = None,
231
max_tokens_override: int | None = None,
232
system_prompt_override: str | None = None,
233
provider_override: str | LLMProviderType | None = None,
234
) -> "BatchSubmissionResult | None":
235
"""Submit retry batch if threshold reached."""
236
...
237
238
async def process_remaining_failed_comparisons(
239
self,
240
cj_batch_id: int,
241
correlation_id: UUID,
242
model_override: str | None = None,
243
temperature_override: float | None = None,
244
max_tokens_override: int | None = None,
245
system_prompt_override: str | None = None,
246
provider_override: str | LLMProviderType | None = None,
247
) -> "BatchSubmissionResult | None":
248
"""Process all remaining failed comparisons at end of batch."""
249
...
250
251
252
class AtomicRedisClientProtocol(Protocol):
253
"""Protocol for Redis client operations."""
254
255
async def lpush(self, key: str, value: str) -> int:
256
"""Push value to Redis list for relay worker notification."""
257
...