prefactor-core 0.2.7__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.
Files changed (37) hide show
  1. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/PKG-INFO +2 -2
  2. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/pyproject.toml +1 -1
  3. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/_version.py +1 -1
  4. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/client.py +41 -1
  5. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/managers/agent_instance.py +57 -2
  6. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/managers/span.py +21 -5
  7. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/models.py +2 -0
  8. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/operations.py +1 -0
  9. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/schema_registry.py +52 -0
  10. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/span_context.py +35 -6
  11. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/.gitignore +0 -0
  12. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/README.md +0 -0
  13. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/examples/agent_e2e.py +0 -0
  14. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/__init__.py +0 -0
  15. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/config.py +0 -0
  16. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/context_stack.py +0 -0
  17. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/exceptions.py +0 -0
  18. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/managers/__init__.py +0 -0
  19. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/monitoring/__init__.py +0 -0
  20. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/monitoring/termination_monitor.py +0 -0
  21. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/queue/__init__.py +0 -0
  22. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/queue/base.py +0 -0
  23. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/queue/executor.py +0 -0
  24. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/queue/memory.py +0 -0
  25. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/src/prefactor_core/utils.py +0 -0
  26. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/monitoring/__init__.py +0 -0
  27. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/monitoring/test_termination_monitor.py +0 -0
  28. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/test_agent_instance_finish_status.py +0 -0
  29. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/test_agent_instance_register.py +0 -0
  30. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/test_client.py +0 -0
  31. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/test_failure_handling.py +0 -0
  32. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/test_imports.py +0 -0
  33. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/test_queue.py +0 -0
  34. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/test_sdk_header.py +0 -0
  35. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/test_span_context.py +0 -0
  36. {prefactor_core-0.2.7 → prefactor_core-0.3.0}/tests/test_span_manager.py +0 -0
  37. {prefactor_core-0.2.7 → 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.2.7
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.1.5
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
 
@@ -9,7 +9,7 @@ authors = [
9
9
  ]
10
10
  requires-python = ">=3.11.0, <4.0.0"
11
11
  dependencies = [
12
- "prefactor-http>=0.1.5",
12
+ "prefactor-http>=0.2.0",
13
13
  "pydantic>=2.0.0",
14
14
  ]
15
15
 
@@ -3,5 +3,5 @@
3
3
  from __future__ import annotations
4
4
 
5
5
  PACKAGE_NAME = "prefactor-core"
6
- __version__ = "0.2.7"
6
+ __version__ = "0.3.0"
7
7
  PACKAGE_VERSION = __version__
@@ -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
@@ -331,6 +333,11 @@ class PrefactorCoreClient:
331
333
  return
332
334
  raise
333
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
+ )
334
341
  elif operation.type == OperationType.CREATE_SPAN:
335
342
  await self._http.agent_spans.create(
336
343
  agent_instance_id=operation.payload["instance_id"],
@@ -339,6 +346,7 @@ class PrefactorCoreClient:
339
346
  id=operation.payload.get("span_id"),
340
347
  parent_span_id=operation.payload.get("parent_span_id"),
341
348
  payload=operation.payload.get("payload"),
349
+ started_at=operation.payload.get("started_at"),
342
350
  control_signal_callback=self._on_control_signal,
343
351
  )
344
352
 
@@ -401,6 +409,7 @@ class PrefactorCoreClient:
401
409
  instance_id: str | None = None,
402
410
  external_schema_version_id: str | None = None,
403
411
  environment_id: str | None = None,
412
+ purpose: InstancePurpose | None = None,
404
413
  ) -> "AgentInstanceHandle":
405
414
  """Create a new agent instance.
406
415
 
@@ -419,6 +428,8 @@ class PrefactorCoreClient:
419
428
  external_schema_version_id: Optional external identifier for the
420
429
  schema version. Defaults to "auto-generated" when using registry.
421
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"``.
422
433
 
423
434
  Returns:
424
435
  AgentInstanceHandle for the created instance.
@@ -460,6 +471,7 @@ class PrefactorCoreClient:
460
471
  agent_schema_version=final_schema_version,
461
472
  instance_id=instance_id,
462
473
  environment_id=environment_id,
474
+ purpose=purpose,
463
475
  )
464
476
 
465
477
  self._set_current_instance(instance_id)
@@ -475,6 +487,7 @@ class PrefactorCoreClient:
475
487
  schema_name: str,
476
488
  parent_span_id: str | None = None,
477
489
  payload: dict[str, Any] | None = None,
490
+ started_at: datetime | None = None,
478
491
  ) -> str:
479
492
  """Create a span and return its ID without finishing it.
480
493
 
@@ -486,6 +499,7 @@ class PrefactorCoreClient:
486
499
  schema_name: Name of the schema for this span.
487
500
  parent_span_id: Optional explicit parent span ID.
488
501
  payload: Optional initial payload (params/inputs) stored on creation.
502
+ started_at: Optional ISO 8601 start time (defaults to current time).
489
503
 
490
504
  Returns:
491
505
  The span ID.
@@ -502,23 +516,49 @@ class PrefactorCoreClient:
502
516
  schema_name=schema_name,
503
517
  parent_span_id=parent_span_id,
504
518
  payload=payload,
519
+ started_at=started_at,
505
520
  )
506
521
 
507
522
  async def finish_span(
508
523
  self,
509
524
  span_id: str,
510
525
  result_payload: dict[str, Any] | None = None,
526
+ timestamp: datetime | None = None,
511
527
  ) -> None:
512
528
  """Finish a previously created span.
513
529
 
514
530
  Args:
515
531
  span_id: The ID of the span to finish.
516
532
  result_payload: Optional result data to store on the span.
533
+ timestamp: Optional ISO 8601 finish time (defaults to current time).
517
534
  """
518
535
  self._ensure_initialized()
519
536
  assert self._span_manager is not None
520
537
 
521
- await self._span_manager.finish(span_id, result_payload=result_payload)
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
+ )
522
562
 
523
563
  @asynccontextmanager
524
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(span_id, result_payload=result_payload)
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(self, temp_id: str) -> None:
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 = datetime.now(timezone.utc)
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 = datetime.now(timezone.utc)
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=datetime.now(timezone.utc),
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
@@ -19,6 +19,7 @@ class OperationType(Enum):
19
19
  REGISTER_AGENT_INSTANCE = auto()
20
20
  START_AGENT_INSTANCE = auto()
21
21
  FINISH_AGENT_INSTANCE = auto()
22
+ UPDATE_AGENT_INSTANCE = auto()
22
23
  CREATE_SPAN = auto()
23
24
  FINISH_SPAN = auto()
24
25
 
@@ -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(self, payload: dict[str, Any] | None = None) -> None:
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(self._span_id, payload=payload)
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(self, result: dict[str, Any] | None = None) -> None:
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(self, result: dict[str, Any] | None = None) -> None:
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(self._span_id)
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:
File without changes