hedit 0.7.5.dev1__tar.gz → 0.7.5.dev2__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 (81) hide show
  1. {hedit-0.7.5.dev1/hedit.egg-info → hedit-0.7.5.dev2}/PKG-INFO +1 -1
  2. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2/hedit.egg-info}/PKG-INFO +1 -1
  3. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/pyproject.toml +1 -1
  4. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/api/main.py +167 -1
  5. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/version.py +1 -1
  6. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_api_endpoints.py +431 -0
  7. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/LICENSE +0 -0
  8. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/PKG_README.md +0 -0
  9. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/README.md +0 -0
  10. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/hedit.egg-info/SOURCES.txt +0 -0
  11. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/hedit.egg-info/dependency_links.txt +0 -0
  12. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/hedit.egg-info/entry_points.txt +0 -0
  13. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/hedit.egg-info/requires.txt +0 -0
  14. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/hedit.egg-info/top_level.txt +0 -0
  15. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/setup.cfg +0 -0
  16. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/__init__.py +0 -0
  17. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/__init__.py +0 -0
  18. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/annotation_agent.py +0 -0
  19. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/assessment_agent.py +0 -0
  20. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/evaluation_agent.py +0 -0
  21. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/feedback_summarizer.py +0 -0
  22. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/feedback_triage_agent.py +0 -0
  23. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/state.py +0 -0
  24. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/validation_agent.py +0 -0
  25. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/vision_agent.py +0 -0
  26. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/agents/workflow.py +0 -0
  27. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/api/__init__.py +0 -0
  28. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/api/models.py +0 -0
  29. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/api/security.py +0 -0
  30. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/cli/__init__.py +0 -0
  31. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/cli/api_executor.py +0 -0
  32. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/cli/client.py +0 -0
  33. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/cli/commands/__init__.py +0 -0
  34. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/cli/config.py +0 -0
  35. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/cli/executor.py +0 -0
  36. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/cli/local_executor.py +0 -0
  37. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/cli/main.py +0 -0
  38. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/cli/output.py +0 -0
  39. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/scripts/__init__.py +0 -0
  40. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/scripts/process_feedback.py +0 -0
  41. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/telemetry/__init__.py +0 -0
  42. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/telemetry/collector.py +0 -0
  43. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/telemetry/schema.py +0 -0
  44. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/telemetry/storage.py +0 -0
  45. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/__init__.py +0 -0
  46. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/error_remediation.py +0 -0
  47. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/github_client.py +0 -0
  48. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/hed_comprehensive_guide.py +0 -0
  49. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/hed_rules.py +0 -0
  50. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/image_processing.py +0 -0
  51. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/json_schema_loader.py +0 -0
  52. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/litellm_llm.py +0 -0
  53. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/openrouter_llm.py +0 -0
  54. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/utils/schema_loader.py +0 -0
  55. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/validation/__init__.py +0 -0
  56. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/validation/hed_lsp.py +0 -0
  57. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/src/validation/hed_validator.py +0 -0
  58. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_annotation_agent.py +0 -0
  59. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_cli_client.py +0 -0
  60. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_cli_config.py +0 -0
  61. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_cli_integration.py +0 -0
  62. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_cli_main.py +0 -0
  63. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_comprehensive_guide.py +0 -0
  64. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_error_remediation.py +0 -0
  65. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_feedback_integration.py +0 -0
  66. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_feedback_triage.py +0 -0
  67. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_github_client.py +0 -0
  68. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_hed_lsp.py +0 -0
  69. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_integration_openrouter.py +0 -0
  70. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_json_schema_loader.py +0 -0
  71. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_keyword_extraction.py +0 -0
  72. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_litellm_llm.py +0 -0
  73. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_no_extend_propagation.py +0 -0
  74. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_openrouter_llm.py +0 -0
  75. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_schema_loader.py +0 -0
  76. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_security.py +0 -0
  77. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_state.py +0 -0
  78. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_telemetry.py +0 -0
  79. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_validation.py +0 -0
  80. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_validation_agent.py +0 -0
  81. {hedit-0.7.5.dev1 → hedit-0.7.5.dev2}/tests/test_version.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: hedit
3
- Version: 0.7.5.dev1
3
+ Version: 0.7.5.dev2
4
4
  Summary: Multi-agent system for HED annotation generation and validation
5
5
  Author-email: Annotation Garden Initiative <info@annotation.garden>
6
6
  License-Expression: MIT
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: hedit
3
- Version: 0.7.5.dev1
3
+ Version: 0.7.5.dev2
4
4
  Summary: Multi-agent system for HED annotation generation and validation
5
5
  Author-email: Annotation Garden Initiative <info@annotation.garden>
6
6
  License-Expression: MIT
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "hedit"
7
- version = "0.7.5.dev1"
7
+ version = "0.7.5.dev2"
8
8
  description = "Multi-agent system for HED annotation generation and validation"
9
9
  readme = "PKG_README.md"
10
10
  requires-python = ">=3.12"
@@ -927,6 +927,64 @@ async def annotate_from_image(
927
927
  ) from e
928
928
 
929
929
 
930
+ async def _collect_stream_telemetry(
931
+ request: AnnotationRequest | ImageAnnotationRequest,
932
+ req: Request,
933
+ current_state: dict,
934
+ start_time: float,
935
+ source: str,
936
+ description: str,
937
+ ) -> None:
938
+ """Collect telemetry for streaming endpoints.
939
+
940
+ Shared helper used by both /annotate/stream and /annotate-from-image/stream.
941
+ Silently returns if telemetry is disabled or collector is not initialized.
942
+
943
+ Args:
944
+ request: The annotation request (text or image)
945
+ req: FastAPI request for header extraction
946
+ current_state: Current workflow state dict
947
+ start_time: Workflow start time (from time.time())
948
+ source: Telemetry source identifier (e.g., "api-stream", "api-image-stream")
949
+ description: Input description text (or image description for image endpoints)
950
+ """
951
+ if not request.telemetry_enabled or not telemetry_collector:
952
+ return
953
+
954
+ latency_ms = int((time.time() - start_time) * 1000)
955
+
956
+ # Get model info from request body, BYOK headers, or server config
957
+ model_name = (
958
+ request.model
959
+ or req.headers.get("x-openrouter-model")
960
+ or os.getenv("ANNOTATION_MODEL", "openai/gpt-oss-120b")
961
+ )
962
+ temperature = request.temperature
963
+ if temperature is None:
964
+ temp_header = req.headers.get("x-openrouter-temperature")
965
+ if temp_header is not None:
966
+ try:
967
+ temperature = float(temp_header)
968
+ except ValueError:
969
+ temperature = None
970
+ if temperature is None:
971
+ temperature = _byok_config.get("temperature", 0.1)
972
+
973
+ event = TelemetryEvent.create(
974
+ description=description,
975
+ schema_version=request.schema_version,
976
+ hed_string=current_state.get("current_annotation", ""),
977
+ iterations=current_state.get("validation_attempts", 0),
978
+ validation_errors=current_state.get("validation_errors", []),
979
+ model=model_name,
980
+ provider=request.provider or req.headers.get("x-openrouter-provider"),
981
+ temperature=temperature,
982
+ latency_ms=latency_ms,
983
+ source=source,
984
+ )
985
+ await telemetry_collector.collect(event)
986
+
987
+
930
988
  @app.post("/annotate/stream")
931
989
  async def annotate_stream(
932
990
  request: AnnotationRequest,
@@ -1042,6 +1100,9 @@ async def annotate_stream(
1042
1100
  # SSE padding comment to force Safari to open the stream
1043
1101
  yield ": stream opened\n\n"
1044
1102
 
1103
+ start_time = time.time()
1104
+ current_state = initial_state.copy()
1105
+
1045
1106
  try:
1046
1107
  # Send initial start event
1047
1108
  yield send_event(
@@ -1049,7 +1110,6 @@ async def annotate_stream(
1049
1110
  )
1050
1111
 
1051
1112
  # Track state and progress
1052
- current_state = initial_state.copy()
1053
1113
  last_stage = None
1054
1114
  validation_attempt = 0
1055
1115
 
@@ -1132,6 +1192,20 @@ async def annotate_stream(
1132
1192
  }
1133
1193
 
1134
1194
  yield send_event("result", result)
1195
+
1196
+ # Collect telemetry after sending result but before done event
1197
+ try:
1198
+ await _collect_stream_telemetry(
1199
+ request=request,
1200
+ req=req,
1201
+ current_state=current_state,
1202
+ start_time=start_time,
1203
+ source="api-stream",
1204
+ description=request.description,
1205
+ )
1206
+ except Exception:
1207
+ logging.debug("Telemetry collection failed for streaming request", exc_info=True)
1208
+
1135
1209
  yield send_event("done", {"message": "Workflow completed"})
1136
1210
 
1137
1211
  except asyncio.CancelledError:
@@ -1145,6 +1219,18 @@ async def annotate_stream(
1145
1219
  "error_type": "timeout",
1146
1220
  },
1147
1221
  )
1222
+ # Collect telemetry on error
1223
+ try:
1224
+ await _collect_stream_telemetry(
1225
+ request=request,
1226
+ req=req,
1227
+ current_state=current_state,
1228
+ start_time=start_time,
1229
+ source="api-stream",
1230
+ description=request.description,
1231
+ )
1232
+ except Exception:
1233
+ logging.debug("Telemetry collection failed on timeout", exc_info=True)
1148
1234
  yield send_event("done", {"message": "Workflow ended with error"})
1149
1235
  except RateLimitError:
1150
1236
  logging.exception("Streaming workflow rate limit")
@@ -1155,6 +1241,18 @@ async def annotate_stream(
1155
1241
  "error_type": "rate_limit",
1156
1242
  },
1157
1243
  )
1244
+ # Collect telemetry on error
1245
+ try:
1246
+ await _collect_stream_telemetry(
1247
+ request=request,
1248
+ req=req,
1249
+ current_state=current_state,
1250
+ start_time=start_time,
1251
+ source="api-stream",
1252
+ description=request.description,
1253
+ )
1254
+ except Exception:
1255
+ logging.debug("Telemetry collection failed on rate limit", exc_info=True)
1158
1256
  yield send_event("done", {"message": "Workflow ended with error"})
1159
1257
  except Exception:
1160
1258
  logging.exception("Streaming workflow error")
@@ -1165,6 +1263,18 @@ async def annotate_stream(
1165
1263
  "error_type": "internal",
1166
1264
  },
1167
1265
  )
1266
+ # Collect telemetry on error
1267
+ try:
1268
+ await _collect_stream_telemetry(
1269
+ request=request,
1270
+ req=req,
1271
+ current_state=current_state,
1272
+ start_time=start_time,
1273
+ source="api-stream",
1274
+ description=request.description,
1275
+ )
1276
+ except Exception:
1277
+ logging.debug("Telemetry collection failed on error", exc_info=True)
1168
1278
  yield send_event("done", {"message": "Workflow ended with error"})
1169
1279
 
1170
1280
  return StreamingResponse(
@@ -1309,6 +1419,10 @@ async def annotate_from_image_stream(
1309
1419
  # SSE padding comment to force Safari to open the stream
1310
1420
  yield ": stream opened\n\n"
1311
1421
 
1422
+ start_time = time.time()
1423
+ current_state: dict = {}
1424
+ image_description = ""
1425
+
1312
1426
  try:
1313
1427
  # Send initial start event
1314
1428
  yield send_event(
@@ -1429,6 +1543,22 @@ async def annotate_from_image_stream(
1429
1543
  }
1430
1544
 
1431
1545
  yield send_event("result", result)
1546
+
1547
+ # Collect telemetry after sending result but before done event
1548
+ try:
1549
+ await _collect_stream_telemetry(
1550
+ request=request,
1551
+ req=req,
1552
+ current_state=current_state,
1553
+ start_time=start_time,
1554
+ source="api-image-stream",
1555
+ description=image_description,
1556
+ )
1557
+ except Exception:
1558
+ logging.debug(
1559
+ "Telemetry collection failed for image streaming request", exc_info=True
1560
+ )
1561
+
1432
1562
  yield send_event("done", {"message": "Workflow completed"})
1433
1563
 
1434
1564
  except asyncio.CancelledError:
@@ -1442,6 +1572,18 @@ async def annotate_from_image_stream(
1442
1572
  "error_type": "timeout",
1443
1573
  },
1444
1574
  )
1575
+ # Collect telemetry on error
1576
+ try:
1577
+ await _collect_stream_telemetry(
1578
+ request=request,
1579
+ req=req,
1580
+ current_state=current_state,
1581
+ start_time=start_time,
1582
+ source="api-image-stream",
1583
+ description=image_description or "image-annotation-failed",
1584
+ )
1585
+ except Exception:
1586
+ logging.debug("Telemetry collection failed on image timeout", exc_info=True)
1445
1587
  yield send_event("done", {"message": "Workflow ended with error"})
1446
1588
  except RateLimitError:
1447
1589
  logging.exception("Streaming image workflow rate limit")
@@ -1452,6 +1594,18 @@ async def annotate_from_image_stream(
1452
1594
  "error_type": "rate_limit",
1453
1595
  },
1454
1596
  )
1597
+ # Collect telemetry on error
1598
+ try:
1599
+ await _collect_stream_telemetry(
1600
+ request=request,
1601
+ req=req,
1602
+ current_state=current_state,
1603
+ start_time=start_time,
1604
+ source="api-image-stream",
1605
+ description=image_description or "image-annotation-failed",
1606
+ )
1607
+ except Exception:
1608
+ logging.debug("Telemetry collection failed on image rate limit", exc_info=True)
1455
1609
  yield send_event("done", {"message": "Workflow ended with error"})
1456
1610
  except Exception:
1457
1611
  logging.exception("Streaming image annotation workflow error")
@@ -1462,6 +1616,18 @@ async def annotate_from_image_stream(
1462
1616
  "error_type": "internal",
1463
1617
  },
1464
1618
  )
1619
+ # Collect telemetry on error
1620
+ try:
1621
+ await _collect_stream_telemetry(
1622
+ request=request,
1623
+ req=req,
1624
+ current_state=current_state,
1625
+ start_time=start_time,
1626
+ source="api-image-stream",
1627
+ description=image_description or "image-annotation-failed",
1628
+ )
1629
+ except Exception:
1630
+ logging.debug("Telemetry collection failed on image error", exc_info=True)
1465
1631
  yield send_event("done", {"message": "Workflow ended with error"})
1466
1632
 
1467
1633
  return StreamingResponse(
@@ -1,6 +1,6 @@
1
1
  """Version information for HEDit."""
2
2
 
3
- __version__ = "0.7.5.dev1"
3
+ __version__ = "0.7.5.dev2"
4
4
  __version_info__ = (0, 7, 5, "dev")
5
5
 
6
6
 
@@ -8,6 +8,7 @@ App is imported inside the fixture to avoid polluting global state.
8
8
 
9
9
  import importlib
10
10
  import os
11
+ import time
11
12
  from unittest.mock import MagicMock, patch
12
13
 
13
14
  import pytest
@@ -841,3 +842,433 @@ class TestTelemetryCollectorIntegration:
841
842
 
842
843
  assert event.source == "api-image"
843
844
  assert event.model.provider == "deepinfra/fp8"
845
+
846
+ def test_telemetry_event_stream_source(self):
847
+ """Test creating a telemetry event with api-stream source."""
848
+ from src.telemetry import TelemetryEvent
849
+
850
+ event = TelemetryEvent.create(
851
+ description="Streaming annotation request",
852
+ schema_version="8.4.0",
853
+ hed_string="Sensory-event, Visual-presentation",
854
+ iterations=2,
855
+ validation_errors=[],
856
+ model="anthropic/claude-haiku-4.5",
857
+ provider="anthropic",
858
+ temperature=0.1,
859
+ latency_ms=2000,
860
+ source="api-stream",
861
+ )
862
+
863
+ assert event.source == "api-stream"
864
+ assert event.input.description == "Streaming annotation request"
865
+ assert event.performance.latency_ms == 2000
866
+
867
+ def test_telemetry_event_image_stream_source(self):
868
+ """Test creating a telemetry event with api-image-stream source."""
869
+ from src.telemetry import TelemetryEvent
870
+
871
+ event = TelemetryEvent.create(
872
+ description="Image description from vision model",
873
+ schema_version="8.4.0",
874
+ hed_string="Visual-presentation",
875
+ iterations=1,
876
+ validation_errors=[],
877
+ model="openai/gpt-4o",
878
+ provider="deepinfra/fp8",
879
+ temperature=0.3,
880
+ latency_ms=4000,
881
+ source="api-image-stream",
882
+ )
883
+
884
+ assert event.source == "api-image-stream"
885
+ assert event.input.description == "Image description from vision model"
886
+
887
+
888
+ class TestStreamingTelemetry:
889
+ """Tests for telemetry collection in streaming endpoints."""
890
+
891
+ @pytest.fixture
892
+ def client_with_telemetry(self):
893
+ """Create a test client with mocked workflow and telemetry collector."""
894
+ original_env = {}
895
+ for key in ["REQUIRE_API_AUTH", "API_KEYS", "OPENROUTER_API_KEY"]:
896
+ if key in os.environ:
897
+ original_env[key] = os.environ[key]
898
+
899
+ os.environ["REQUIRE_API_AUTH"] = "true"
900
+ os.environ["API_KEYS"] = "test-api-key-for-unit-tests"
901
+ os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
902
+
903
+ from src.api import security
904
+
905
+ importlib.reload(security)
906
+
907
+ # Create mock workflow
908
+ mock_workflow = MagicMock()
909
+ mock_graph = MagicMock()
910
+
911
+ async def mock_stream_events(*args, **kwargs):
912
+ yield {"event": "on_chain_start", "name": "annotate", "data": {}}
913
+ yield {
914
+ "event": "on_chain_end",
915
+ "name": "annotate",
916
+ "data": {
917
+ "output": {
918
+ "current_annotation": "Sensory-event, Visual-presentation",
919
+ "validation_attempts": 1,
920
+ }
921
+ },
922
+ }
923
+ yield {"event": "on_chain_start", "name": "validate", "data": {}}
924
+ yield {
925
+ "event": "on_chain_end",
926
+ "name": "validate",
927
+ "data": {"output": {"is_valid": True, "validation_errors": []}},
928
+ }
929
+ yield {"event": "on_chain_start", "name": "evaluate", "data": {}}
930
+ yield {
931
+ "event": "on_chain_end",
932
+ "name": "evaluate",
933
+ "data": {"output": {"is_faithful": True, "is_complete": True}},
934
+ }
935
+
936
+ mock_graph.astream_events = mock_stream_events
937
+ mock_workflow.graph = mock_graph
938
+
939
+ # Create mock telemetry collector with a list to track collected events
940
+ mock_collector = MagicMock()
941
+ collected_events = []
942
+
943
+ async def track_collect(event):
944
+ collected_events.append(event)
945
+ return True
946
+
947
+ mock_collector.collect = track_collect
948
+
949
+ with (
950
+ patch("src.api.main.workflow", mock_workflow),
951
+ patch("src.api.main.telemetry_collector", mock_collector),
952
+ ):
953
+ from src.api.main import app
954
+
955
+ client = TestClient(app, raise_server_exceptions=False)
956
+ yield client, collected_events
957
+
958
+ for key in ["REQUIRE_API_AUTH", "API_KEYS", "OPENROUTER_API_KEY"]:
959
+ if key in original_env:
960
+ os.environ[key] = original_env[key]
961
+ elif key in os.environ:
962
+ del os.environ[key]
963
+
964
+ importlib.reload(security)
965
+
966
+ @pytest.fixture
967
+ def client_with_telemetry_disabled(self):
968
+ """Create a test client with mocked workflow but no telemetry collector."""
969
+ original_env = {}
970
+ for key in ["REQUIRE_API_AUTH", "API_KEYS", "OPENROUTER_API_KEY"]:
971
+ if key in os.environ:
972
+ original_env[key] = os.environ[key]
973
+
974
+ os.environ["REQUIRE_API_AUTH"] = "true"
975
+ os.environ["API_KEYS"] = "test-api-key-for-unit-tests"
976
+ os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
977
+
978
+ from src.api import security
979
+
980
+ importlib.reload(security)
981
+
982
+ mock_workflow = MagicMock()
983
+ mock_graph = MagicMock()
984
+
985
+ async def mock_stream_events(*args, **kwargs):
986
+ yield {"event": "on_chain_start", "name": "annotate", "data": {}}
987
+ yield {
988
+ "event": "on_chain_end",
989
+ "name": "annotate",
990
+ "data": {
991
+ "output": {
992
+ "current_annotation": "Sensory-event",
993
+ "validation_attempts": 1,
994
+ }
995
+ },
996
+ }
997
+ yield {"event": "on_chain_start", "name": "validate", "data": {}}
998
+ yield {
999
+ "event": "on_chain_end",
1000
+ "name": "validate",
1001
+ "data": {"output": {"is_valid": True, "validation_errors": []}},
1002
+ }
1003
+
1004
+ mock_graph.astream_events = mock_stream_events
1005
+ mock_workflow.graph = mock_graph
1006
+
1007
+ # Track if collect was called
1008
+ collected_events = []
1009
+
1010
+ async def track_collect(event):
1011
+ collected_events.append(event)
1012
+ return True
1013
+
1014
+ mock_collector = MagicMock()
1015
+ mock_collector.collect = track_collect
1016
+
1017
+ with (
1018
+ patch("src.api.main.workflow", mock_workflow),
1019
+ patch("src.api.main.telemetry_collector", mock_collector),
1020
+ ):
1021
+ from src.api.main import app
1022
+
1023
+ client = TestClient(app, raise_server_exceptions=False)
1024
+ yield client, collected_events
1025
+
1026
+ for key in ["REQUIRE_API_AUTH", "API_KEYS", "OPENROUTER_API_KEY"]:
1027
+ if key in original_env:
1028
+ os.environ[key] = original_env[key]
1029
+ elif key in os.environ:
1030
+ del os.environ[key]
1031
+
1032
+ importlib.reload(security)
1033
+
1034
+ @pytest.fixture
1035
+ def client_with_failing_workflow(self):
1036
+ """Create a test client with a workflow that raises an error."""
1037
+ original_env = {}
1038
+ for key in ["REQUIRE_API_AUTH", "API_KEYS", "OPENROUTER_API_KEY"]:
1039
+ if key in os.environ:
1040
+ original_env[key] = os.environ[key]
1041
+
1042
+ os.environ["REQUIRE_API_AUTH"] = "true"
1043
+ os.environ["API_KEYS"] = "test-api-key-for-unit-tests"
1044
+ os.environ["OPENROUTER_API_KEY"] = "test-openrouter-key"
1045
+
1046
+ from src.api import security
1047
+
1048
+ importlib.reload(security)
1049
+
1050
+ mock_workflow = MagicMock()
1051
+ mock_graph = MagicMock()
1052
+
1053
+ async def mock_stream_events_error(*args, **kwargs):
1054
+ yield {"event": "on_chain_start", "name": "annotate", "data": {}}
1055
+ yield {
1056
+ "event": "on_chain_end",
1057
+ "name": "annotate",
1058
+ "data": {
1059
+ "output": {
1060
+ "current_annotation": "Partial-annotation",
1061
+ "validation_attempts": 0,
1062
+ }
1063
+ },
1064
+ }
1065
+ raise RuntimeError("Simulated workflow failure")
1066
+
1067
+ mock_graph.astream_events = mock_stream_events_error
1068
+ mock_workflow.graph = mock_graph
1069
+
1070
+ collected_events = []
1071
+
1072
+ async def track_collect(event):
1073
+ collected_events.append(event)
1074
+ return True
1075
+
1076
+ mock_collector = MagicMock()
1077
+ mock_collector.collect = track_collect
1078
+
1079
+ with (
1080
+ patch("src.api.main.workflow", mock_workflow),
1081
+ patch("src.api.main.telemetry_collector", mock_collector),
1082
+ ):
1083
+ from src.api.main import app
1084
+
1085
+ client = TestClient(app, raise_server_exceptions=False)
1086
+ yield client, collected_events
1087
+
1088
+ for key in ["REQUIRE_API_AUTH", "API_KEYS", "OPENROUTER_API_KEY"]:
1089
+ if key in original_env:
1090
+ os.environ[key] = original_env[key]
1091
+ elif key in os.environ:
1092
+ del os.environ[key]
1093
+
1094
+ importlib.reload(security)
1095
+
1096
+ def test_stream_telemetry_collected_on_success(self, client_with_telemetry):
1097
+ """Test that telemetry is collected for successful streaming requests."""
1098
+ client, collected_events = client_with_telemetry
1099
+ request_data = {
1100
+ "description": "A red circle appears on screen",
1101
+ "schema_version": "8.3.0",
1102
+ "telemetry_enabled": True,
1103
+ }
1104
+ response = client.post("/annotate/stream", json=request_data, headers=TEST_AUTH_HEADERS)
1105
+ assert response.status_code == 200
1106
+ assert len(collected_events) == 1
1107
+
1108
+ event = collected_events[0]
1109
+ assert event.source == "api-stream"
1110
+ assert event.input.description == "A red circle appears on screen"
1111
+ assert event.input.schema_version == "8.3.0"
1112
+ assert event.performance.latency_ms >= 0
1113
+
1114
+ def test_stream_telemetry_not_collected_when_disabled(self, client_with_telemetry_disabled):
1115
+ """Test that telemetry is not collected when telemetry_enabled=False."""
1116
+ client, collected_events = client_with_telemetry_disabled
1117
+ request_data = {
1118
+ "description": "A red circle appears on screen",
1119
+ "schema_version": "8.3.0",
1120
+ "telemetry_enabled": False,
1121
+ }
1122
+ response = client.post("/annotate/stream", json=request_data, headers=TEST_AUTH_HEADERS)
1123
+ assert response.status_code == 200
1124
+ assert len(collected_events) == 0
1125
+
1126
+ def test_stream_telemetry_collected_on_workflow_error(self, client_with_failing_workflow):
1127
+ """Test that telemetry is collected even when workflow fails."""
1128
+ client, collected_events = client_with_failing_workflow
1129
+ request_data = {
1130
+ "description": "A red circle appears on screen",
1131
+ "schema_version": "8.3.0",
1132
+ "telemetry_enabled": True,
1133
+ }
1134
+ response = client.post("/annotate/stream", json=request_data, headers=TEST_AUTH_HEADERS)
1135
+ assert response.status_code == 200
1136
+ # Telemetry should still be collected on error
1137
+ assert len(collected_events) == 1
1138
+
1139
+ event = collected_events[0]
1140
+ assert event.source == "api-stream"
1141
+ # Partial state should be captured
1142
+ assert event.output.hed_string == "Partial-annotation"
1143
+
1144
+ def test_stream_telemetry_has_correct_model_info(self, client_with_telemetry):
1145
+ """Test that telemetry captures model info from headers."""
1146
+ client, collected_events = client_with_telemetry
1147
+ request_data = {
1148
+ "description": "A blue square flashes",
1149
+ "schema_version": "8.4.0",
1150
+ "telemetry_enabled": True,
1151
+ "model": "anthropic/claude-haiku-4.5",
1152
+ "provider": "anthropic",
1153
+ "temperature": 0.3,
1154
+ }
1155
+ response = client.post("/annotate/stream", json=request_data, headers=TEST_AUTH_HEADERS)
1156
+ assert response.status_code == 200
1157
+ assert len(collected_events) == 1
1158
+
1159
+ event = collected_events[0]
1160
+ assert event.model.model == "anthropic/claude-haiku-4.5"
1161
+ assert event.model.provider == "anthropic"
1162
+ assert event.model.temperature == 0.3
1163
+
1164
+ def test_stream_telemetry_result_still_sent(self, client_with_telemetry):
1165
+ """Test that result and done events are still sent with telemetry."""
1166
+ client, collected_events = client_with_telemetry
1167
+ request_data = {
1168
+ "description": "Test event",
1169
+ "schema_version": "8.3.0",
1170
+ "telemetry_enabled": True,
1171
+ }
1172
+ response = client.post("/annotate/stream", json=request_data, headers=TEST_AUTH_HEADERS)
1173
+ assert response.status_code == 200
1174
+ content = response.text
1175
+ # Result and done events should still be present
1176
+ assert "event: result" in content
1177
+ assert "event: done" in content
1178
+
1179
+
1180
+ class TestCollectStreamTelemetryHelper:
1181
+ """Tests for the _collect_stream_telemetry helper function."""
1182
+
1183
+ def test_helper_function_exists(self):
1184
+ """Test that _collect_stream_telemetry is importable."""
1185
+ from src.api.main import _collect_stream_telemetry
1186
+
1187
+ assert callable(_collect_stream_telemetry)
1188
+
1189
+ @pytest.mark.asyncio
1190
+ async def test_helper_skips_when_telemetry_disabled(self):
1191
+ """Test helper returns without collecting when telemetry is disabled."""
1192
+ from src.api.main import _collect_stream_telemetry
1193
+ from src.api.models import AnnotationRequest
1194
+
1195
+ request = AnnotationRequest(description="Test", telemetry_enabled=False)
1196
+ mock_req = MagicMock()
1197
+
1198
+ # Should not raise even with no collector
1199
+ with patch("src.api.main.telemetry_collector", None):
1200
+ await _collect_stream_telemetry(
1201
+ request=request,
1202
+ req=mock_req,
1203
+ current_state={},
1204
+ start_time=time.time(),
1205
+ source="api-stream",
1206
+ description="Test",
1207
+ )
1208
+
1209
+ @pytest.mark.asyncio
1210
+ async def test_helper_skips_when_collector_is_none(self):
1211
+ """Test helper returns without collecting when collector is None."""
1212
+ from src.api.main import _collect_stream_telemetry
1213
+ from src.api.models import AnnotationRequest
1214
+
1215
+ request = AnnotationRequest(description="Test", telemetry_enabled=True)
1216
+ mock_req = MagicMock()
1217
+
1218
+ with patch("src.api.main.telemetry_collector", None):
1219
+ await _collect_stream_telemetry(
1220
+ request=request,
1221
+ req=mock_req,
1222
+ current_state={},
1223
+ start_time=time.time(),
1224
+ source="api-stream",
1225
+ description="Test",
1226
+ )
1227
+
1228
+ @pytest.mark.asyncio
1229
+ async def test_helper_collects_when_enabled(self):
1230
+ """Test helper collects telemetry when enabled with collector."""
1231
+ from src.api.main import _collect_stream_telemetry
1232
+ from src.api.models import AnnotationRequest
1233
+
1234
+ request = AnnotationRequest(
1235
+ description="Test event description",
1236
+ schema_version="8.4.0",
1237
+ telemetry_enabled=True,
1238
+ model="test/model",
1239
+ temperature=0.2,
1240
+ )
1241
+
1242
+ mock_req = MagicMock()
1243
+ mock_req.headers = {}
1244
+
1245
+ collected = []
1246
+
1247
+ async def mock_collect(event):
1248
+ collected.append(event)
1249
+ return True
1250
+
1251
+ mock_collector = MagicMock()
1252
+ mock_collector.collect = mock_collect
1253
+
1254
+ with patch("src.api.main.telemetry_collector", mock_collector):
1255
+ await _collect_stream_telemetry(
1256
+ request=request,
1257
+ req=mock_req,
1258
+ current_state={
1259
+ "current_annotation": "Sensory-event",
1260
+ "validation_attempts": 2,
1261
+ "validation_errors": [],
1262
+ },
1263
+ start_time=time.time() - 1.5,
1264
+ source="api-stream",
1265
+ description="Test event description",
1266
+ )
1267
+
1268
+ assert len(collected) == 1
1269
+ event = collected[0]
1270
+ assert event.source == "api-stream"
1271
+ assert event.input.description == "Test event description"
1272
+ assert event.output.hed_string == "Sensory-event"
1273
+ assert event.output.iterations == 2
1274
+ assert event.performance.latency_ms >= 1400 # ~1.5 seconds
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