prefactor-core 0.2.6__tar.gz → 0.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/.gitignore +0 -1
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/PKG-INFO +2 -2
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/pyproject.toml +1 -1
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/_version.py +1 -1
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/client.py +49 -1
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/managers/agent_instance.py +57 -2
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/managers/span.py +21 -5
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/models.py +2 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/operations.py +1 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/schema_registry.py +52 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/span_context.py +35 -6
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_agent_instance_finish_status.py +3 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_agent_instance_register.py +3 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_failure_handling.py +31 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_sdk_header.py +8 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/README.md +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/examples/agent_e2e.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/__init__.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/config.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/context_stack.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/exceptions.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/managers/__init__.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/monitoring/__init__.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/monitoring/termination_monitor.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/queue/__init__.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/queue/base.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/queue/executor.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/queue/memory.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/utils.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/monitoring/__init__.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/monitoring/test_termination_monitor.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_client.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_imports.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_queue.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_span_context.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_span_manager.py +0 -0
- {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_utils.py +0 -0
|
@@ -1,11 +1,11 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: prefactor-core
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.0
|
|
4
4
|
Summary: Core Prefactor SDK with async queue-based operations
|
|
5
5
|
Author-email: Prefactor Pty Ltd <josh@prefactor.tech>
|
|
6
6
|
License: MIT
|
|
7
7
|
Requires-Python: <4.0.0,>=3.11.0
|
|
8
|
-
Requires-Dist: prefactor-http>=0.
|
|
8
|
+
Requires-Dist: prefactor-http>=0.2.0
|
|
9
9
|
Requires-Dist: pydantic>=2.0.0
|
|
10
10
|
Description-Content-Type: text/markdown
|
|
11
11
|
|
|
@@ -11,6 +11,7 @@ import asyncio
|
|
|
11
11
|
import logging
|
|
12
12
|
import time
|
|
13
13
|
from contextlib import asynccontextmanager
|
|
14
|
+
from datetime import datetime
|
|
14
15
|
from typing import TYPE_CHECKING, Any
|
|
15
16
|
|
|
16
17
|
from prefactor_http.client import PrefactorHttpClient
|
|
@@ -19,6 +20,7 @@ from prefactor_http.exceptions import (
|
|
|
19
20
|
is_permanent_http_error,
|
|
20
21
|
is_transient_http_error,
|
|
21
22
|
)
|
|
23
|
+
from prefactor_http.models.types import InstancePurpose
|
|
22
24
|
|
|
23
25
|
from ._version import PACKAGE_NAME as CORE_PACKAGE_NAME
|
|
24
26
|
from ._version import PACKAGE_VERSION as CORE_PACKAGE_VERSION
|
|
@@ -142,6 +144,14 @@ class PrefactorCoreClient:
|
|
|
142
144
|
)
|
|
143
145
|
await self._http.__aenter__()
|
|
144
146
|
|
|
147
|
+
try:
|
|
148
|
+
await self._http.validate_token()
|
|
149
|
+
except BaseException as exc:
|
|
150
|
+
http = self._http
|
|
151
|
+
self._http = None
|
|
152
|
+
await http.__aexit__(type(exc), exc, exc.__traceback__)
|
|
153
|
+
raise
|
|
154
|
+
|
|
145
155
|
# Initialize executor
|
|
146
156
|
self._executor = TaskExecutor(
|
|
147
157
|
queue=self._queue,
|
|
@@ -323,6 +333,11 @@ class PrefactorCoreClient:
|
|
|
323
333
|
return
|
|
324
334
|
raise
|
|
325
335
|
|
|
336
|
+
elif operation.type == OperationType.UPDATE_AGENT_INSTANCE:
|
|
337
|
+
await self._http.agent_instances.update(
|
|
338
|
+
agent_instance_id=operation.payload["instance_id"],
|
|
339
|
+
quality_payload=operation.payload.get("quality_payload"),
|
|
340
|
+
)
|
|
326
341
|
elif operation.type == OperationType.CREATE_SPAN:
|
|
327
342
|
await self._http.agent_spans.create(
|
|
328
343
|
agent_instance_id=operation.payload["instance_id"],
|
|
@@ -331,6 +346,7 @@ class PrefactorCoreClient:
|
|
|
331
346
|
id=operation.payload.get("span_id"),
|
|
332
347
|
parent_span_id=operation.payload.get("parent_span_id"),
|
|
333
348
|
payload=operation.payload.get("payload"),
|
|
349
|
+
started_at=operation.payload.get("started_at"),
|
|
334
350
|
control_signal_callback=self._on_control_signal,
|
|
335
351
|
)
|
|
336
352
|
|
|
@@ -393,6 +409,7 @@ class PrefactorCoreClient:
|
|
|
393
409
|
instance_id: str | None = None,
|
|
394
410
|
external_schema_version_id: str | None = None,
|
|
395
411
|
environment_id: str | None = None,
|
|
412
|
+
purpose: InstancePurpose | None = None,
|
|
396
413
|
) -> "AgentInstanceHandle":
|
|
397
414
|
"""Create a new agent instance.
|
|
398
415
|
|
|
@@ -411,6 +428,8 @@ class PrefactorCoreClient:
|
|
|
411
428
|
external_schema_version_id: Optional external identifier for the
|
|
412
429
|
schema version. Defaults to "auto-generated" when using registry.
|
|
413
430
|
environment_id: Optional environment ID used to scope the agent instance.
|
|
431
|
+
purpose: Why this instance ran — ``"live"``, ``"smoke_test"``,
|
|
432
|
+
or ``"eval"``. Omitted (None) lets the API default to ``"live"``.
|
|
414
433
|
|
|
415
434
|
Returns:
|
|
416
435
|
AgentInstanceHandle for the created instance.
|
|
@@ -452,6 +471,7 @@ class PrefactorCoreClient:
|
|
|
452
471
|
agent_schema_version=final_schema_version,
|
|
453
472
|
instance_id=instance_id,
|
|
454
473
|
environment_id=environment_id,
|
|
474
|
+
purpose=purpose,
|
|
455
475
|
)
|
|
456
476
|
|
|
457
477
|
self._set_current_instance(instance_id)
|
|
@@ -467,6 +487,7 @@ class PrefactorCoreClient:
|
|
|
467
487
|
schema_name: str,
|
|
468
488
|
parent_span_id: str | None = None,
|
|
469
489
|
payload: dict[str, Any] | None = None,
|
|
490
|
+
started_at: datetime | None = None,
|
|
470
491
|
) -> str:
|
|
471
492
|
"""Create a span and return its ID without finishing it.
|
|
472
493
|
|
|
@@ -478,6 +499,7 @@ class PrefactorCoreClient:
|
|
|
478
499
|
schema_name: Name of the schema for this span.
|
|
479
500
|
parent_span_id: Optional explicit parent span ID.
|
|
480
501
|
payload: Optional initial payload (params/inputs) stored on creation.
|
|
502
|
+
started_at: Optional ISO 8601 start time (defaults to current time).
|
|
481
503
|
|
|
482
504
|
Returns:
|
|
483
505
|
The span ID.
|
|
@@ -494,23 +516,49 @@ class PrefactorCoreClient:
|
|
|
494
516
|
schema_name=schema_name,
|
|
495
517
|
parent_span_id=parent_span_id,
|
|
496
518
|
payload=payload,
|
|
519
|
+
started_at=started_at,
|
|
497
520
|
)
|
|
498
521
|
|
|
499
522
|
async def finish_span(
|
|
500
523
|
self,
|
|
501
524
|
span_id: str,
|
|
502
525
|
result_payload: dict[str, Any] | None = None,
|
|
526
|
+
timestamp: datetime | None = None,
|
|
503
527
|
) -> None:
|
|
504
528
|
"""Finish a previously created span.
|
|
505
529
|
|
|
506
530
|
Args:
|
|
507
531
|
span_id: The ID of the span to finish.
|
|
508
532
|
result_payload: Optional result data to store on the span.
|
|
533
|
+
timestamp: Optional ISO 8601 finish time (defaults to current time).
|
|
509
534
|
"""
|
|
510
535
|
self._ensure_initialized()
|
|
511
536
|
assert self._span_manager is not None
|
|
512
537
|
|
|
513
|
-
await self._span_manager.finish(
|
|
538
|
+
await self._span_manager.finish(
|
|
539
|
+
span_id,
|
|
540
|
+
result_payload=result_payload,
|
|
541
|
+
timestamp=timestamp,
|
|
542
|
+
)
|
|
543
|
+
|
|
544
|
+
async def update_agent_instance(
|
|
545
|
+
self,
|
|
546
|
+
instance_id: str,
|
|
547
|
+
quality_payload: dict[str, Any] | None = None,
|
|
548
|
+
) -> None:
|
|
549
|
+
"""Update an agent instance (e.g., set quality payload).
|
|
550
|
+
|
|
551
|
+
Args:
|
|
552
|
+
instance_id: The ID of the instance to update.
|
|
553
|
+
quality_payload: Quality evaluation payload (None to clear).
|
|
554
|
+
"""
|
|
555
|
+
self._ensure_initialized()
|
|
556
|
+
assert self._instance_manager is not None
|
|
557
|
+
|
|
558
|
+
await self._instance_manager.update(
|
|
559
|
+
instance_id,
|
|
560
|
+
quality_payload=quality_payload,
|
|
561
|
+
)
|
|
514
562
|
|
|
515
563
|
@asynccontextmanager
|
|
516
564
|
async def span(
|
|
@@ -14,7 +14,7 @@ from ..utils import generate_idempotency_key
|
|
|
14
14
|
|
|
15
15
|
if TYPE_CHECKING:
|
|
16
16
|
from prefactor_http.client import PrefactorHttpClient
|
|
17
|
-
from prefactor_http.models.types import FinishStatus
|
|
17
|
+
from prefactor_http.models.types import FinishStatus, InstancePurpose
|
|
18
18
|
|
|
19
19
|
from ..client import PrefactorCoreClient
|
|
20
20
|
|
|
@@ -64,6 +64,7 @@ class AgentInstanceManager:
|
|
|
64
64
|
agent_id: str | None = None,
|
|
65
65
|
instance_id: str | None = None,
|
|
66
66
|
environment_id: str | None = None,
|
|
67
|
+
purpose: "InstancePurpose | None" = None,
|
|
67
68
|
) -> str:
|
|
68
69
|
"""Register a new agent instance.
|
|
69
70
|
|
|
@@ -79,6 +80,8 @@ class AgentInstanceManager:
|
|
|
79
80
|
the API generates one.
|
|
80
81
|
environment_id: Optional environment ID. Required when using an
|
|
81
82
|
account-scoped token; omit when using a deployment-scoped token.
|
|
83
|
+
purpose: Why this instance ran — ``"live"``, ``"smoke_test"``,
|
|
84
|
+
or ``"eval"``. Omitted (None) lets the API default to ``"live"``.
|
|
82
85
|
|
|
83
86
|
Returns:
|
|
84
87
|
The instance ID (API-generated).
|
|
@@ -90,6 +93,7 @@ class AgentInstanceManager:
|
|
|
90
93
|
environment_id=environment_id,
|
|
91
94
|
id=instance_id,
|
|
92
95
|
idempotency_key=generate_idempotency_key(),
|
|
96
|
+
purpose=purpose,
|
|
93
97
|
)
|
|
94
98
|
return result.id
|
|
95
99
|
|
|
@@ -158,6 +162,30 @@ class AgentInstanceManager:
|
|
|
158
162
|
|
|
159
163
|
await self._enqueue(operation)
|
|
160
164
|
|
|
165
|
+
async def update(
|
|
166
|
+
self,
|
|
167
|
+
instance_id: str,
|
|
168
|
+
quality_payload: dict[str, Any] | None = None,
|
|
169
|
+
) -> None:
|
|
170
|
+
"""Update an agent instance (e.g., set quality payload).
|
|
171
|
+
|
|
172
|
+
Queues an update operation for the instance.
|
|
173
|
+
|
|
174
|
+
Args:
|
|
175
|
+
instance_id: The ID of the instance to update.
|
|
176
|
+
quality_payload: Quality evaluation payload (None to clear).
|
|
177
|
+
"""
|
|
178
|
+
operation = Operation(
|
|
179
|
+
type=OperationType.UPDATE_AGENT_INSTANCE,
|
|
180
|
+
payload={
|
|
181
|
+
"instance_id": instance_id,
|
|
182
|
+
"quality_payload": quality_payload,
|
|
183
|
+
},
|
|
184
|
+
timestamp=datetime.now(timezone.utc),
|
|
185
|
+
)
|
|
186
|
+
|
|
187
|
+
await self._enqueue(operation)
|
|
188
|
+
|
|
161
189
|
|
|
162
190
|
class AgentInstanceHandle:
|
|
163
191
|
"""Handle to an agent instance with convenience methods.
|
|
@@ -238,11 +266,30 @@ class AgentInstanceHandle:
|
|
|
238
266
|
status=status,
|
|
239
267
|
)
|
|
240
268
|
|
|
269
|
+
async def update(
|
|
270
|
+
self,
|
|
271
|
+
quality_payload: dict[str, Any] | None = None,
|
|
272
|
+
) -> None:
|
|
273
|
+
"""Update the instance (e.g., set quality payload).
|
|
274
|
+
|
|
275
|
+
This queues an update operation for the instance.
|
|
276
|
+
|
|
277
|
+
Args:
|
|
278
|
+
quality_payload: Quality evaluation payload (None to clear).
|
|
279
|
+
"""
|
|
280
|
+
manager = self._client.instance_manager
|
|
281
|
+
assert manager is not None
|
|
282
|
+
await manager.update(
|
|
283
|
+
self._instance_id,
|
|
284
|
+
quality_payload=quality_payload,
|
|
285
|
+
)
|
|
286
|
+
|
|
241
287
|
async def create_span(
|
|
242
288
|
self,
|
|
243
289
|
schema_name: str,
|
|
244
290
|
parent_span_id: str | None = None,
|
|
245
291
|
payload: dict[str, Any] | None = None,
|
|
292
|
+
started_at: datetime | None = None,
|
|
246
293
|
) -> str:
|
|
247
294
|
"""Create a span within this instance and return its ID.
|
|
248
295
|
|
|
@@ -252,6 +299,7 @@ class AgentInstanceHandle:
|
|
|
252
299
|
schema_name: Name of the schema for this span.
|
|
253
300
|
parent_span_id: Optional explicit parent span ID.
|
|
254
301
|
payload: Optional initial payload (params/inputs) stored on creation.
|
|
302
|
+
started_at: Optional ISO 8601 start time (defaults to current time).
|
|
255
303
|
|
|
256
304
|
Returns:
|
|
257
305
|
The span ID.
|
|
@@ -262,20 +310,27 @@ class AgentInstanceHandle:
|
|
|
262
310
|
schema_name=schema_name,
|
|
263
311
|
parent_span_id=parent_span_id,
|
|
264
312
|
payload=payload,
|
|
313
|
+
started_at=started_at,
|
|
265
314
|
)
|
|
266
315
|
|
|
267
316
|
async def finish_span(
|
|
268
317
|
self,
|
|
269
318
|
span_id: str,
|
|
270
319
|
result_payload: dict[str, Any] | None = None,
|
|
320
|
+
timestamp: datetime | None = None,
|
|
271
321
|
) -> None:
|
|
272
322
|
"""Finish a previously created span.
|
|
273
323
|
|
|
274
324
|
Args:
|
|
275
325
|
span_id: The ID of the span to finish.
|
|
276
326
|
result_payload: Optional result data to store on the span.
|
|
327
|
+
timestamp: Optional ISO 8601 finish time (defaults to current time).
|
|
277
328
|
"""
|
|
278
|
-
await self._client.finish_span(
|
|
329
|
+
await self._client.finish_span(
|
|
330
|
+
span_id,
|
|
331
|
+
result_payload=result_payload,
|
|
332
|
+
timestamp=timestamp,
|
|
333
|
+
)
|
|
279
334
|
|
|
280
335
|
@asynccontextmanager
|
|
281
336
|
async def span(
|
|
@@ -109,6 +109,7 @@ class SpanManager:
|
|
|
109
109
|
self,
|
|
110
110
|
temp_id: str,
|
|
111
111
|
payload: dict[str, Any] | None = None,
|
|
112
|
+
started_at: datetime | None = None,
|
|
112
113
|
) -> str:
|
|
113
114
|
"""Post the span to the API as ``active`` and return the API-generated ID.
|
|
114
115
|
|
|
@@ -121,6 +122,7 @@ class SpanManager:
|
|
|
121
122
|
Args:
|
|
122
123
|
temp_id: The temporary ID returned by ``prepare()``.
|
|
123
124
|
payload: Optional params/inputs to send with the span.
|
|
125
|
+
started_at: Optional ISO 8601 start time (defaults to current time).
|
|
124
126
|
|
|
125
127
|
Returns:
|
|
126
128
|
The API-generated span ID.
|
|
@@ -139,6 +141,7 @@ class SpanManager:
|
|
|
139
141
|
status="active",
|
|
140
142
|
payload=payload or {},
|
|
141
143
|
parent_span_id=span.parent_span_id,
|
|
144
|
+
started_at=started_at,
|
|
142
145
|
idempotency_key=generate_idempotency_key(),
|
|
143
146
|
)
|
|
144
147
|
|
|
@@ -148,6 +151,7 @@ class SpanManager:
|
|
|
148
151
|
span.id = api_id
|
|
149
152
|
span.payload = payload or {}
|
|
150
153
|
span.status = "active"
|
|
154
|
+
span.started_at = started_at or datetime.now(timezone.utc)
|
|
151
155
|
del self._spans[temp_id]
|
|
152
156
|
self._spans[api_id] = span
|
|
153
157
|
|
|
@@ -165,7 +169,11 @@ class SpanManager:
|
|
|
165
169
|
|
|
166
170
|
return api_id
|
|
167
171
|
|
|
168
|
-
async def cancel_unstarted(
|
|
172
|
+
async def cancel_unstarted(
|
|
173
|
+
self,
|
|
174
|
+
temp_id: str,
|
|
175
|
+
timestamp: datetime | None = None,
|
|
176
|
+
) -> None:
|
|
169
177
|
"""Cancel a span that was never started.
|
|
170
178
|
|
|
171
179
|
When ``cancel()`` is called before ``start()``, the span has not yet
|
|
@@ -175,6 +183,7 @@ class SpanManager:
|
|
|
175
183
|
|
|
176
184
|
Args:
|
|
177
185
|
temp_id: The temporary ID returned by ``prepare()``.
|
|
186
|
+
timestamp: Optional ISO 8601 finish time (defaults to current time).
|
|
178
187
|
|
|
179
188
|
Raises:
|
|
180
189
|
KeyError: If temp_id is not a known pending span.
|
|
@@ -197,11 +206,13 @@ class SpanManager:
|
|
|
197
206
|
await self._http.agent_spans.finish(
|
|
198
207
|
agent_span_id=api_id,
|
|
199
208
|
status="cancelled",
|
|
209
|
+
timestamp=timestamp,
|
|
200
210
|
idempotency_key=generate_idempotency_key(),
|
|
201
211
|
)
|
|
202
212
|
|
|
213
|
+
effective_timestamp = timestamp or datetime.now(timezone.utc)
|
|
203
214
|
span.status = "cancelled"
|
|
204
|
-
span.finished_at =
|
|
215
|
+
span.finished_at = effective_timestamp
|
|
205
216
|
|
|
206
217
|
stack = SpanContextStack.get_stack()
|
|
207
218
|
if temp_id in stack:
|
|
@@ -218,6 +229,7 @@ class SpanManager:
|
|
|
218
229
|
parent_span_id: str | None = None,
|
|
219
230
|
payload: dict[str, Any] | None = None,
|
|
220
231
|
span_id: str | None = None,
|
|
232
|
+
started_at: datetime | None = None,
|
|
221
233
|
) -> str:
|
|
222
234
|
"""Create a span in one step (prepare + start).
|
|
223
235
|
|
|
@@ -230,6 +242,7 @@ class SpanManager:
|
|
|
230
242
|
parent_span_id: Optional parent span ID (auto-detected if None).
|
|
231
243
|
payload: Optional initial payload data.
|
|
232
244
|
span_id: Ignored (API generates IDs).
|
|
245
|
+
started_at: Optional ISO 8601 start time (defaults to current time).
|
|
233
246
|
|
|
234
247
|
Returns:
|
|
235
248
|
The API-generated span ID.
|
|
@@ -239,7 +252,7 @@ class SpanManager:
|
|
|
239
252
|
schema_name=schema_name,
|
|
240
253
|
parent_span_id=parent_span_id,
|
|
241
254
|
)
|
|
242
|
-
return await self.start(temp_id, payload=payload)
|
|
255
|
+
return await self.start(temp_id, payload=payload, started_at=started_at)
|
|
243
256
|
|
|
244
257
|
async def finish(
|
|
245
258
|
self,
|
|
@@ -247,6 +260,7 @@ class SpanManager:
|
|
|
247
260
|
result_payload: dict[str, Any] | None = None,
|
|
248
261
|
status: "FinishStatus" = "complete",
|
|
249
262
|
idempotency_key: str | None = None,
|
|
263
|
+
timestamp: datetime | None = None,
|
|
250
264
|
) -> None:
|
|
251
265
|
"""Mark a span as finished.
|
|
252
266
|
|
|
@@ -261,6 +275,7 @@ class SpanManager:
|
|
|
261
275
|
to cancel a span that was never started.
|
|
262
276
|
idempotency_key: Optional key to make repeated finish requests
|
|
263
277
|
duplicate-safe. When omitted, a new key is generated.
|
|
278
|
+
timestamp: Optional ISO 8601 finish time (defaults to current time).
|
|
264
279
|
|
|
265
280
|
Raises:
|
|
266
281
|
KeyError: If the span ID is not known.
|
|
@@ -268,8 +283,9 @@ class SpanManager:
|
|
|
268
283
|
if span_id not in self._spans:
|
|
269
284
|
raise KeyError(f"Unknown span: {span_id}")
|
|
270
285
|
|
|
286
|
+
effective_timestamp = timestamp or datetime.now(timezone.utc)
|
|
271
287
|
self._spans[span_id].status = status
|
|
272
|
-
self._spans[span_id].finished_at =
|
|
288
|
+
self._spans[span_id].finished_at = effective_timestamp
|
|
273
289
|
|
|
274
290
|
stack = SpanContextStack.get_stack()
|
|
275
291
|
if span_id in stack:
|
|
@@ -288,7 +304,7 @@ class SpanManager:
|
|
|
288
304
|
operation = Operation(
|
|
289
305
|
type=OperationType.FINISH_SPAN,
|
|
290
306
|
payload=op_payload,
|
|
291
|
-
timestamp=
|
|
307
|
+
timestamp=effective_timestamp,
|
|
292
308
|
)
|
|
293
309
|
|
|
294
310
|
await self._enqueue(operation)
|
|
@@ -28,10 +28,12 @@ class AgentInstance:
|
|
|
28
28
|
id: str
|
|
29
29
|
agent_id: str
|
|
30
30
|
status: str = "pending"
|
|
31
|
+
purpose: str | None = None
|
|
31
32
|
created_at: datetime = field(default_factory=datetime.now)
|
|
32
33
|
started_at: datetime | None = None
|
|
33
34
|
finished_at: datetime | None = None
|
|
34
35
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
36
|
+
quality_payload: dict[str, Any] | None = None
|
|
35
37
|
|
|
36
38
|
|
|
37
39
|
@dataclass
|
|
@@ -67,6 +67,8 @@ class SchemaRegistry:
|
|
|
67
67
|
self._span_result_schemas: dict[str, dict[str, Any]] = {}
|
|
68
68
|
# span_type_schemas: name → full structured entry
|
|
69
69
|
self._span_type_schemas: dict[str, dict[str, Any]] = {}
|
|
70
|
+
# quality_schema: optional quality schema for instance evaluations
|
|
71
|
+
self._quality_schema: dict[str, Any] | None = None
|
|
70
72
|
|
|
71
73
|
def register(
|
|
72
74
|
self,
|
|
@@ -188,6 +190,46 @@ class SchemaRegistry:
|
|
|
188
190
|
|
|
189
191
|
self._span_type_schemas[name] = entry
|
|
190
192
|
|
|
193
|
+
def register_quality_schema(
|
|
194
|
+
self,
|
|
195
|
+
schema: dict[str, Any],
|
|
196
|
+
title: str | None = None,
|
|
197
|
+
description: str | None = None,
|
|
198
|
+
template: str | None = None,
|
|
199
|
+
data_risk: dict[str, Any] | None = None,
|
|
200
|
+
) -> None:
|
|
201
|
+
"""Register a quality schema for instance evaluations.
|
|
202
|
+
|
|
203
|
+
The quality schema defines the shape of quality payloads that can be
|
|
204
|
+
set on agent instances. It uses the same ``title``, ``description``,
|
|
205
|
+
``template``, and ``data_risk`` fields as span type schemas.
|
|
206
|
+
|
|
207
|
+
Args:
|
|
208
|
+
schema: JSON Schema dict defining the quality payload structure.
|
|
209
|
+
title: Optional human-readable title (defaults to "quality" on API).
|
|
210
|
+
description: Optional description of the quality evaluation.
|
|
211
|
+
template: Optional display template using ``{{field}}`` interpolation.
|
|
212
|
+
data_risk: Optional data risk classification dict (same structure
|
|
213
|
+
as span type data_risk).
|
|
214
|
+
|
|
215
|
+
Raises:
|
|
216
|
+
ValueError: If a quality schema is already registered.
|
|
217
|
+
"""
|
|
218
|
+
if self._quality_schema is not None:
|
|
219
|
+
raise ValueError("Quality schema is already registered")
|
|
220
|
+
|
|
221
|
+
entry: dict[str, Any] = {"schema": schema}
|
|
222
|
+
if title is not None:
|
|
223
|
+
entry["title"] = title
|
|
224
|
+
if description is not None:
|
|
225
|
+
entry["description"] = description
|
|
226
|
+
if template is not None:
|
|
227
|
+
entry["template"] = template
|
|
228
|
+
if data_risk is not None:
|
|
229
|
+
entry["data_risk"] = data_risk
|
|
230
|
+
|
|
231
|
+
self._quality_schema = entry
|
|
232
|
+
|
|
191
233
|
def get(self, schema_name: str) -> dict[str, Any] | None:
|
|
192
234
|
"""Get a params schema by name.
|
|
193
235
|
|
|
@@ -244,6 +286,9 @@ class SchemaRegistry:
|
|
|
244
286
|
if self._span_type_schemas:
|
|
245
287
|
result["span_type_schemas"] = list(self._span_type_schemas.values())
|
|
246
288
|
|
|
289
|
+
if self._quality_schema is not None:
|
|
290
|
+
result["quality_schema"] = dict(self._quality_schema)
|
|
291
|
+
|
|
247
292
|
return result
|
|
248
293
|
|
|
249
294
|
def merge(self, other: "SchemaRegistry") -> None:
|
|
@@ -270,6 +315,10 @@ class SchemaRegistry:
|
|
|
270
315
|
if name in self._span_type_schemas:
|
|
271
316
|
conflicts.append(f"span_type_schemas/{name}")
|
|
272
317
|
|
|
318
|
+
if other._quality_schema is not None:
|
|
319
|
+
if self._quality_schema is not None:
|
|
320
|
+
conflicts.append("quality_schema")
|
|
321
|
+
|
|
273
322
|
if conflicts:
|
|
274
323
|
msg = f"Cannot merge registries - conflicting schemas: {conflicts}"
|
|
275
324
|
raise ValueError(msg)
|
|
@@ -283,5 +332,8 @@ class SchemaRegistry:
|
|
|
283
332
|
for name, entry in other._span_type_schemas.items():
|
|
284
333
|
self._span_type_schemas[name] = entry.copy()
|
|
285
334
|
|
|
335
|
+
if other._quality_schema is not None:
|
|
336
|
+
self._quality_schema = other._quality_schema.copy()
|
|
337
|
+
|
|
286
338
|
|
|
287
339
|
__all__ = ["SchemaRegistry"]
|
|
@@ -6,6 +6,7 @@ and ensures proper cleanup when the span completes.
|
|
|
6
6
|
|
|
7
7
|
from __future__ import annotations
|
|
8
8
|
|
|
9
|
+
from datetime import datetime
|
|
9
10
|
from typing import TYPE_CHECKING, Any
|
|
10
11
|
|
|
11
12
|
from .utils import generate_idempotency_key
|
|
@@ -80,6 +81,7 @@ class SpanContext:
|
|
|
80
81
|
tuple[FinishStatus, tuple[tuple[str, Any], ...]] | None
|
|
81
82
|
) = None
|
|
82
83
|
self._finish_idempotency_key: str | None = None
|
|
84
|
+
self._finish_timestamp: datetime | None = None
|
|
83
85
|
|
|
84
86
|
@property
|
|
85
87
|
def id(self) -> str:
|
|
@@ -93,7 +95,11 @@ class SpanContext:
|
|
|
93
95
|
"""
|
|
94
96
|
return self._span_id
|
|
95
97
|
|
|
96
|
-
async def start(
|
|
98
|
+
async def start(
|
|
99
|
+
self,
|
|
100
|
+
payload: dict[str, Any] | None = None,
|
|
101
|
+
started_at: datetime | None = None,
|
|
102
|
+
) -> None:
|
|
97
103
|
"""Post the span to the API as ``active`` with the given params payload.
|
|
98
104
|
|
|
99
105
|
This triggers ``POST /api/v1/agent_spans``. The span is created as
|
|
@@ -105,11 +111,14 @@ class SpanContext:
|
|
|
105
111
|
payload: Optional params/inputs for the span (e.g. model name,
|
|
106
112
|
prompt text, tool input). Stored as the span's ``payload``
|
|
107
113
|
field in the API.
|
|
114
|
+
started_at: Optional ISO 8601 start time (defaults to current time).
|
|
108
115
|
"""
|
|
109
116
|
if self._started:
|
|
110
117
|
return
|
|
111
118
|
|
|
112
|
-
api_id = await self._span_manager.start(
|
|
119
|
+
api_id = await self._span_manager.start(
|
|
120
|
+
self._span_id, payload=payload, started_at=started_at
|
|
121
|
+
)
|
|
113
122
|
self._span_id = api_id
|
|
114
123
|
self._started = True
|
|
115
124
|
|
|
@@ -124,37 +133,53 @@ class SpanContext:
|
|
|
124
133
|
"""
|
|
125
134
|
self._result_payload.update(data)
|
|
126
135
|
|
|
127
|
-
async def complete(
|
|
136
|
+
async def complete(
|
|
137
|
+
self,
|
|
138
|
+
result: dict[str, Any] | None = None,
|
|
139
|
+
timestamp: datetime | None = None,
|
|
140
|
+
) -> None:
|
|
128
141
|
"""Finish the span with ``complete`` status.
|
|
129
142
|
|
|
130
143
|
Args:
|
|
131
144
|
result: Optional result payload to attach to the span.
|
|
145
|
+
timestamp: Optional ISO 8601 finish time (defaults to current time).
|
|
132
146
|
"""
|
|
133
147
|
if result:
|
|
134
148
|
self.set_result(result)
|
|
135
149
|
self._finish_status = "complete"
|
|
150
|
+
self._finish_timestamp = timestamp
|
|
136
151
|
await self._finish()
|
|
137
152
|
|
|
138
|
-
async def fail(
|
|
153
|
+
async def fail(
|
|
154
|
+
self,
|
|
155
|
+
result: dict[str, Any] | None = None,
|
|
156
|
+
timestamp: datetime | None = None,
|
|
157
|
+
) -> None:
|
|
139
158
|
"""Finish the span with ``failed`` status.
|
|
140
159
|
|
|
141
160
|
Args:
|
|
142
161
|
result: Optional result payload (e.g. error details).
|
|
162
|
+
timestamp: Optional ISO 8601 finish time (defaults to current time).
|
|
143
163
|
"""
|
|
144
164
|
if result:
|
|
145
165
|
self.set_result(result)
|
|
146
166
|
self._finish_status = "failed"
|
|
167
|
+
self._finish_timestamp = timestamp
|
|
147
168
|
await self._finish()
|
|
148
169
|
|
|
149
|
-
async def cancel(self) -> None:
|
|
170
|
+
async def cancel(self, timestamp: datetime | None = None) -> None:
|
|
150
171
|
"""Finish the span with ``cancelled`` status.
|
|
151
172
|
|
|
152
173
|
Can be called before or after ``start()``. If ``start()`` has not
|
|
153
174
|
been called yet, the span is posted as ``pending`` then immediately
|
|
154
175
|
cancelled — the API only accepts cancellation from the ``pending``
|
|
155
176
|
state, so this is always a valid sequence.
|
|
177
|
+
|
|
178
|
+
Args:
|
|
179
|
+
timestamp: Optional ISO 8601 finish time (defaults to current time).
|
|
156
180
|
"""
|
|
157
181
|
self._finish_status = "cancelled"
|
|
182
|
+
self._finish_timestamp = timestamp
|
|
158
183
|
await self._finish()
|
|
159
184
|
|
|
160
185
|
async def finish(self) -> None:
|
|
@@ -192,7 +217,10 @@ class SpanContext:
|
|
|
192
217
|
|
|
193
218
|
try:
|
|
194
219
|
if not self._started and self._finish_status == "cancelled":
|
|
195
|
-
await self._span_manager.cancel_unstarted(
|
|
220
|
+
await self._span_manager.cancel_unstarted(
|
|
221
|
+
self._span_id,
|
|
222
|
+
timestamp=self._finish_timestamp,
|
|
223
|
+
)
|
|
196
224
|
self._finished = True
|
|
197
225
|
return
|
|
198
226
|
|
|
@@ -206,6 +234,7 @@ class SpanContext:
|
|
|
206
234
|
result_payload=self._result_payload or None,
|
|
207
235
|
status=self._finish_status,
|
|
208
236
|
idempotency_key=self._finish_idempotency_key,
|
|
237
|
+
timestamp=self._finish_timestamp,
|
|
209
238
|
)
|
|
210
239
|
self._finished = True
|
|
211
240
|
except Exception:
|
|
@@ -56,6 +56,9 @@ class _StubHttpClient:
|
|
|
56
56
|
async def __aexit__(self, exc_type, exc, tb):
|
|
57
57
|
return None
|
|
58
58
|
|
|
59
|
+
async def validate_token(self):
|
|
60
|
+
return {"status": "success"}
|
|
61
|
+
|
|
59
62
|
|
|
60
63
|
@pytest.mark.asyncio
|
|
61
64
|
async def test_create_agent_instance_passes_environment_id_to_register():
|
|
@@ -68,17 +68,28 @@ class _StubHttpClient:
|
|
|
68
68
|
*_args,
|
|
69
69
|
agent_instances: _StubAgentInstances | None = None,
|
|
70
70
|
agent_spans: _StubAgentSpans | None = None,
|
|
71
|
+
validate_token_error: PrefactorAuthError | None = None,
|
|
71
72
|
**_kwargs,
|
|
72
73
|
) -> None:
|
|
73
74
|
self.agent_instances = agent_instances or _StubAgentInstances()
|
|
74
75
|
self.agent_spans = agent_spans or _StubAgentSpans()
|
|
76
|
+
self._validate_token_error = validate_token_error
|
|
77
|
+
self.enter_calls = 0
|
|
78
|
+
self.exit_calls = 0
|
|
75
79
|
|
|
76
80
|
async def __aenter__(self):
|
|
81
|
+
self.enter_calls += 1
|
|
77
82
|
return self
|
|
78
83
|
|
|
79
84
|
async def __aexit__(self, exc_type, exc, tb):
|
|
85
|
+
self.exit_calls += 1
|
|
80
86
|
return None
|
|
81
87
|
|
|
88
|
+
async def validate_token(self):
|
|
89
|
+
if self._validate_token_error is not None:
|
|
90
|
+
raise self._validate_token_error
|
|
91
|
+
return {"status": "success"}
|
|
92
|
+
|
|
82
93
|
|
|
83
94
|
def _make_client_config(max_retries: int = 0) -> PrefactorCoreConfig:
|
|
84
95
|
return PrefactorCoreConfig(
|
|
@@ -375,3 +386,23 @@ async def test_latched_failure_drops_already_queued_backlog():
|
|
|
375
386
|
await client.close()
|
|
376
387
|
|
|
377
388
|
assert exc_info.value.dropped_operations == 1
|
|
389
|
+
|
|
390
|
+
|
|
391
|
+
@pytest.mark.asyncio
|
|
392
|
+
async def test_initialize_raises_auth_error_when_token_validation_fails():
|
|
393
|
+
"""Invalid tokens should fail during initialize() before workers start."""
|
|
394
|
+
stub_http = _StubHttpClient(
|
|
395
|
+
validate_token_error=PrefactorAuthError("Token expired", "bad_authtoken", 401)
|
|
396
|
+
)
|
|
397
|
+
|
|
398
|
+
with patch("prefactor_core.client.PrefactorHttpClient", return_value=stub_http):
|
|
399
|
+
client = PrefactorCoreClient(_make_client_config())
|
|
400
|
+
with pytest.raises(PrefactorAuthError) as exc_info:
|
|
401
|
+
await client.initialize()
|
|
402
|
+
|
|
403
|
+
assert exc_info.value.code == "bad_authtoken"
|
|
404
|
+
assert stub_http.enter_calls == 1
|
|
405
|
+
assert stub_http.exit_calls == 1
|
|
406
|
+
assert client._http is None
|
|
407
|
+
assert client._executor is None
|
|
408
|
+
await client.close()
|
|
@@ -37,6 +37,10 @@ class TestPrefactorCoreSdkHeader:
|
|
|
37
37
|
"prefactor_http.client.PrefactorHttpClient.__aexit__",
|
|
38
38
|
AsyncMock(return_value=None),
|
|
39
39
|
),
|
|
40
|
+
patch(
|
|
41
|
+
"prefactor_http.client.PrefactorHttpClient.validate_token",
|
|
42
|
+
AsyncMock(return_value={"status": "success"}),
|
|
43
|
+
),
|
|
40
44
|
):
|
|
41
45
|
await client.initialize()
|
|
42
46
|
assert client._http is not None
|
|
@@ -59,6 +63,10 @@ class TestPrefactorCoreSdkHeader:
|
|
|
59
63
|
"prefactor_http.client.PrefactorHttpClient.__aexit__",
|
|
60
64
|
AsyncMock(return_value=None),
|
|
61
65
|
),
|
|
66
|
+
patch(
|
|
67
|
+
"prefactor_http.client.PrefactorHttpClient.validate_token",
|
|
68
|
+
AsyncMock(return_value={"status": "success"}),
|
|
69
|
+
),
|
|
62
70
|
):
|
|
63
71
|
await client.initialize()
|
|
64
72
|
assert client._http is not None
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/monitoring/termination_monitor.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|