simpleaudit 0.3.1__tar.gz → 0.3.3__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 (127) hide show
  1. {simpleaudit-0.3.1/simpleaudit.egg-info → simpleaudit-0.3.3}/PKG-INFO +1 -1
  2. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/pyproject.toml +1 -1
  3. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/__init__.py +2 -1
  4. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/model_auditor.py +55 -1
  5. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/model.py +14 -0
  6. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/__init__.py +6 -0
  7. simpleaudit-0.3.3/simpleaudit/tracing/auth.py +332 -0
  8. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/otlp.py +43 -13
  9. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/provider.py +11 -2
  10. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/selection.py +29 -3
  11. {simpleaudit-0.3.1 → simpleaudit-0.3.3/simpleaudit.egg-info}/PKG-INFO +1 -1
  12. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/SOURCES.txt +1 -0
  13. simpleaudit-0.3.3/tests/test_otlp_auth.py +401 -0
  14. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_targets.py +69 -0
  15. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_tracing.py +22 -1
  16. simpleaudit-0.3.1/simpleaudit/tracing/auth.py +0 -175
  17. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/LICENSE +0 -0
  18. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/README.md +0 -0
  19. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/setup.cfg +0 -0
  20. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/auditor.py +0 -0
  21. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/checklist.py +0 -0
  22. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/cli.py +0 -0
  23. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/context_marks.py +0 -0
  24. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/cross_judge.py +0 -0
  25. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/experiment.py +0 -0
  26. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/__init__.py +0 -0
  27. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/abstention.py +0 -0
  28. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/binary_abstention.py +0 -0
  29. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/checklist.py +0 -0
  30. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/compose.py +0 -0
  31. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/default.py +0 -0
  32. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/factuality.py +0 -0
  33. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/harm.py +0 -0
  34. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/helpfulness.py +0 -0
  35. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/helsedir_sexhealth_no.py +0 -0
  36. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/helsedir_sexhealth_no_rag.py +0 -0
  37. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/judges/safety.py +0 -0
  38. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/perturbations.py +0 -0
  39. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/reframing.py +0 -0
  40. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/repeated_results.py +0 -0
  41. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/results.py +0 -0
  42. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/__init__.py +0 -0
  43. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/arbeidstilsynet_arbeidstid.py +0 -0
  44. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/bullshitbench_health.py +0 -0
  45. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/bullshitbench_v1_v2.py +0 -0
  46. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/health.py +0 -0
  47. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/hei_refusal.py +0 -0
  48. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/helfo.py +0 -0
  49. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/helpmed.py +0 -0
  50. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/human_rights_education.py +0 -0
  51. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/human_rights_fair_trial.py +0 -0
  52. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/human_rights_water.py +0 -0
  53. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/blank_chart.png +0 -0
  54. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/illegible_label.png +0 -0
  55. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/injection.png +0 -0
  56. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/make_images.py +0 -0
  57. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/product_a.png +0 -0
  58. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/product_b.png +0 -0
  59. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/quarterly_revenue.png +0 -0
  60. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/images/truncated_axis.png +0 -0
  61. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/lanekassen.py +0 -0
  62. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/nav_aap.py +0 -0
  63. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/nb_kryss_ordning.py +0 -0
  64. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/rag.py +0 -0
  65. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/safety.py +0 -0
  66. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/skatteetaten.py +0 -0
  67. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/skatteetaten_legitimasjon.py +0 -0
  68. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/system_prompt.py +0 -0
  69. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/toll_reisegodskvote.py +0 -0
  70. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/ung.py +0 -0
  71. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/scenarios/vision_integrity.py +0 -0
  72. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/spans.py +0 -0
  73. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/stats.py +0 -0
  74. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/__init__.py +0 -0
  75. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/base.py +0 -0
  76. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/callable.py +0 -0
  77. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/targets/http.py +0 -0
  78. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/context.py +0 -0
  79. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/shared.py +0 -0
  80. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/tracing/store.py +0 -0
  81. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/utils.py +0 -0
  82. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/visualization/scenario_viewer.html +0 -0
  83. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/visualization/server.py +0 -0
  84. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/visualization/thumbnail.png +0 -0
  85. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit/visualization/visualizer.html +0 -0
  86. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/dependency_links.txt +0 -0
  87. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/entry_points.txt +0 -0
  88. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/requires.txt +0 -0
  89. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/simpleaudit.egg-info/top_level.txt +0 -0
  90. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_anthropic_judge_default_path.py +0 -0
  91. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_audit_flow.py +0 -0
  92. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_basic.py +0 -0
  93. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_checklist.py +0 -0
  94. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_checklist_judge.py +0 -0
  95. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_checklist_pipeline.py +0 -0
  96. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_checklist_reframing.py +0 -0
  97. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_cross_judge.py +0 -0
  98. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_custom_prompts.py +0 -0
  99. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_document_expansion.py +0 -0
  100. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_experiment_streaming.py +0 -0
  101. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_fakes.py +0 -0
  102. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_file_uri.py +0 -0
  103. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_fragility.py +0 -0
  104. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_generation_params.py +0 -0
  105. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_json_format.py +0 -0
  106. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_compose.py +0 -0
  107. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_notes.py +0 -0
  108. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_output_contract.py +0 -0
  109. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_postprocess.py +0 -0
  110. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_registry.py +0 -0
  111. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_response_schema.py +0 -0
  112. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_judge_robustness_pipeline.py +0 -0
  113. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_model_auditor.py +0 -0
  114. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_perturbations.py +0 -0
  115. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_reframing.py +0 -0
  116. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_rejudge.py +0 -0
  117. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_repeated_experiments.py +0 -0
  118. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_resumable_experiments.py +0 -0
  119. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_scenario_data.py +0 -0
  120. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_scenario_pack_conventions.py +0 -0
  121. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_spans.py +0 -0
  122. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_stats.py +0 -0
  123. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_strip_thinking.py +0 -0
  124. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_target_api_key.py +0 -0
  125. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_token_counting.py +0 -0
  126. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_tracing_auth.py +0 -0
  127. {simpleaudit-0.3.1 → simpleaudit-0.3.3}/tests/test_vision_integrity_pack.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: simpleaudit
3
- Version: 0.3.1
3
+ Version: 0.3.3
4
4
  Summary: Lightweight AI Safety Auditing Framework
5
5
  Author: SimpleAudit Contributors
6
6
  License: MIT
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "simpleaudit"
7
- version = "0.3.1"
7
+ version = "0.3.3"
8
8
  description = "Lightweight AI Safety Auditing Framework"
9
9
  readme = "README.md"
10
10
  license = {text = "MIT"}
@@ -34,7 +34,7 @@ except PackageNotFoundError:
34
34
  __version__ = "0.0.0+unknown"
35
35
  __author__ = "SimpleAudit Contributors"
36
36
 
37
- from .model_auditor import ModelAuditor
37
+ from .model_auditor import ModelAuditor, ScenarioExecution
38
38
  from .auditor import Auditor
39
39
  from .targets import (
40
40
  CallableTarget,
@@ -86,6 +86,7 @@ from .perturbations import (
86
86
 
87
87
  __all__ = [
88
88
  "ModelAuditor",
89
+ "ScenarioExecution",
89
90
  "Auditor",
90
91
  "Target",
91
92
  "TargetContext",
@@ -43,6 +43,31 @@ from .utils import (
43
43
  from .utils import parse_json_response as _parse_json_response
44
44
 
45
45
 
46
+ from dataclasses import dataclass
47
+
48
+
49
+ @dataclass
50
+ class ScenarioExecution:
51
+ """Captures the state of a scenario after target execution but before judgment.
52
+
53
+ This dataclass represents one scenario execution through the target conversation phase.
54
+ It holds all state necessary for the evidence resolver to select trace evidence, which
55
+ can then be passed to the judge.
56
+ """
57
+ scenario_name: str
58
+ scenario_description: str
59
+ conversation: List[Dict]
60
+ trace_correlation: Optional[Any] = None
61
+ audit_run_id: Optional[str] = None
62
+ error: Optional[str] = None
63
+ scenario_meta: Optional[Dict[str, Any]] = None
64
+ expected_behavior: Optional[List[str]] = None
65
+ target_input_tokens: int = 0
66
+ target_output_tokens: int = 0
67
+ auditor_input_tokens: int = 0
68
+ auditor_output_tokens: int = 0
69
+
70
+
46
71
  def _user_agent() -> str:
47
72
  try:
48
73
  return f"simpleaudit/{_pkg_version('simpleaudit')}"
@@ -867,6 +892,7 @@ Evaluate this conversation and respond with this exact JSON structure:
867
892
  evidence_spans: Optional[List[Dict[str, Any]]] = None,
868
893
  audit_run_id: Optional[str] = None,
869
894
  trace_correlation: Optional[Any] = None,
895
+ evidence_resolver: Optional[Callable[[ScenarioExecution], Union[List[Dict[str, Any]], None]]] = None,
870
896
  ) -> AuditResult:
871
897
  turns = max_turns or self.max_turns
872
898
  # Per-scenario correlation ids. A fresh trace id per scenario keeps each
@@ -976,6 +1002,30 @@ Evaluate this conversation and respond with this exact JSON structure:
976
1002
  self._log(f"--- Scenario FAILED mid-conversation: {name} [{error}] ---")
977
1003
 
978
1004
  if error is None:
1005
+ # Call evidence resolver if provided, after target conversation, before judge.
1006
+ resolved_evidence = evidence_spans
1007
+ if evidence_resolver is not None:
1008
+ execution = ScenarioExecution(
1009
+ scenario_name=name,
1010
+ scenario_description=description,
1011
+ conversation=conversation,
1012
+ scenario_meta=scenario_meta,
1013
+ expected_behavior=expected_behavior,
1014
+ trace_correlation=trace_correlation,
1015
+ audit_run_id=audit_run_id,
1016
+ error=None,
1017
+ target_input_tokens=target_input_tokens,
1018
+ target_output_tokens=target_output_tokens,
1019
+ auditor_input_tokens=auditor_input_tokens,
1020
+ auditor_output_tokens=auditor_output_tokens,
1021
+ )
1022
+ try:
1023
+ resolved_evidence = await evidence_resolver(execution) if asyncio.iscoroutinefunction(evidence_resolver) else evidence_resolver(execution)
1024
+ except Exception as exc:
1025
+ self._log(f"--- Evidence resolution failed: {type(exc).__name__}: {exc} ---")
1026
+ # If evidence resolution fails, continue without evidence
1027
+ resolved_evidence = evidence_spans
1028
+
979
1029
  self._log("Judging conversation...", name=name)
980
1030
  judge_prompt, judge_schema, judge_postprocess, fell_back = self._resolve_judge_spec(
981
1031
  self.judge_prompt,
@@ -1003,7 +1053,7 @@ Evaluate this conversation and respond with this exact JSON structure:
1003
1053
  postprocess=judge_postprocess,
1004
1054
  scenario_meta=scenario_meta,
1005
1055
  params=effective_judge or None,
1006
- evidence_spans=evidence_spans,
1056
+ evidence_spans=resolved_evidence,
1007
1057
  )
1008
1058
  judge_input_tokens += j_in
1009
1059
  judge_output_tokens += j_out
@@ -1076,6 +1126,7 @@ Evaluate this conversation and respond with this exact JSON structure:
1076
1126
  evidence_spans: Optional[List[Dict[str, Any]]] = None,
1077
1127
  audit_run_id: Optional[str] = None,
1078
1128
  trace_correlation: Optional[Any] = None,
1129
+ evidence_resolver: Optional[Callable[[ScenarioExecution], Union[List[Dict[str, Any]], None]]] = None,
1079
1130
  ) -> AuditResults:
1080
1131
  if max_workers < 1:
1081
1132
  raise ValueError(
@@ -1150,6 +1201,7 @@ Evaluate this conversation and respond with this exact JSON structure:
1150
1201
  evidence_spans=evidence_spans,
1151
1202
  audit_run_id=audit_run_id,
1152
1203
  trace_correlation=trace_correlation,
1204
+ evidence_resolver=evidence_resolver,
1153
1205
  )
1154
1206
  except Exception as exc:
1155
1207
  # Don't let one failing scenario abort the whole batch and
@@ -1210,6 +1262,7 @@ Evaluate this conversation and respond with this exact JSON structure:
1210
1262
  evidence_spans: Optional[List[Dict[str, Any]]] = None,
1211
1263
  audit_run_id: Optional[str] = None,
1212
1264
  trace_correlation: Optional[Any] = None,
1265
+ evidence_resolver: Optional[Callable[[ScenarioExecution], Union[List[Dict[str, Any]], None]]] = None,
1213
1266
  ) -> AuditResults:
1214
1267
  try:
1215
1268
  asyncio.get_running_loop()
@@ -1228,6 +1281,7 @@ Evaluate this conversation and respond with this exact JSON structure:
1228
1281
  evidence_spans=evidence_spans,
1229
1282
  audit_run_id=audit_run_id,
1230
1283
  trace_correlation=trace_correlation,
1284
+ evidence_resolver=evidence_resolver,
1231
1285
  )
1232
1286
  )
1233
1287
  msg = "ModelAuditor.run() cannot be called from an active event loop. Use await <object>.run_async()."
@@ -88,6 +88,7 @@ class ModelTarget:
88
88
  documents=documents,
89
89
  response_format=response_format,
90
90
  params=params,
91
+ context=context,
91
92
  max_retries=self.max_retries,
92
93
  retry_backoff=self.retry_backoff,
93
94
  )
@@ -104,6 +105,7 @@ async def _client_send(
104
105
  documents: Optional[List[Union[str, Dict[str, Any]]]],
105
106
  response_format: Optional[Dict[str, Any]],
106
107
  params: Optional[Dict[str, Any]],
108
+ context: Optional[TargetContext] = None,
107
109
  max_retries: int,
108
110
  retry_backoff: float,
109
111
  ) -> TargetResponse:
@@ -111,9 +113,21 @@ async def _client_send(
111
113
 
112
114
  Imported lazily so that ``simpleaudit.targets`` does not create an import
113
115
  cycle with ``model_auditor`` at package import time.
116
+
117
+ When the per-turn context carries W3C trace headers (e.g. the engine's
118
+ scenario traceparent), they are merged into ``extra_headers`` so the
119
+ target's outgoing call propagates the trace context. This is what lets an
120
+ instrumented target export spans under the same trace id the engine's
121
+ TraceCorrelation records. User-supplied extra_headers are preserved.
114
122
  """
115
123
  from ..model_auditor import ModelAuditor
116
124
 
125
+ if context is not None and context.trace_headers:
126
+ headers = dict(params.get("extra_headers", {}) if params else {})
127
+ headers.update(context.trace_headers)
128
+ params = dict(params or {})
129
+ params["extra_headers"] = headers
130
+
117
131
  content, input_tokens, output_tokens = await ModelAuditor._call_async(
118
132
  client,
119
133
  model,
@@ -15,11 +15,14 @@ when the target is instrumented.
15
15
  from .auth import (
16
16
  AuthResult,
17
17
  Authenticator,
18
+ OTLPAuth,
19
+ TARGET_ID_ATTR,
18
20
  hash_secret,
19
21
  make_basic_bearer_authenticator,
20
22
  new_salt,
21
23
  parse_basic_header,
22
24
  parse_bearer_header,
25
+ tag_spans,
23
26
  token_lookup_prefix,
24
27
  verify_secret,
25
28
  )
@@ -71,6 +74,9 @@ __all__ = [
71
74
  "DEFAULT_NOISE_KINDS",
72
75
  "AuthResult",
73
76
  "Authenticator",
77
+ "OTLPAuth",
78
+ "TARGET_ID_ATTR",
79
+ "tag_spans",
74
80
  "hash_secret",
75
81
  "new_salt",
76
82
  "verify_secret",
@@ -0,0 +1,332 @@
1
+ """
2
+ OTLP receiver authentication — the "WHO is allowed to submit" layer.
3
+
4
+ This module is deliberately **framework-agnostic**: it has no Django (or any
5
+ web-framework) dependency so both the standalone ``simpleaudit`` library and
6
+ the SimpleAuditStudio app can use the same auth. It answers one question —
7
+ which target is allowed to push spans — and maps the credential to a stable
8
+ ``target_id`` that the receiver tags onto every ingested span.
9
+
10
+ Two auth modes (plus ``none``):
11
+
12
+ - ``none`` — no check. Local / ephemeral receiver on a trusted network.
13
+ - ``basic`` — ``Authorization: Basic base64(username:password)``.
14
+ Best for OpenWebUI, which natively supports
15
+ ``OTEL_BASIC_AUTH_USERNAME`` / ``OTEL_BASIC_AUTH_PASSWORD``.
16
+ - ``bearer`` — ``Authorization: Bearer <token>``. Best for generic OTel
17
+ exporters that can set arbitrary headers
18
+ (``OTEL_EXPORTER_OTLP_HEADERS=Authorization=Bearer <t>``).
19
+
20
+ Credentials are revocable and rotatable. Secrets are stored as salted
21
+ SHA-256 hashes and compared in constant time (``hmac.compare_digest``) to
22
+ avoid timing side-channels.
23
+ """
24
+
25
+ from __future__ import annotations
26
+
27
+ import base64
28
+ import binascii
29
+ import hashlib
30
+ import hmac
31
+ import os
32
+ import secrets
33
+ from dataclasses import dataclass
34
+ from typing import Callable, Dict, Optional, Protocol, Tuple
35
+
36
+ # Salt length (bytes) prepended to each secret before hashing.
37
+ _SALT_LEN = 16
38
+ # Attribute key the receiver stamps onto every authenticated span.
39
+ TARGET_ID_ATTR = "simpleaudit.target_id"
40
+
41
+ SALT_LEN = _SALT_LEN
42
+ BEARER_TOKEN_PREFIX = "sa_otlp_"
43
+ TOKEN_LOOKUP_PREFIX_LEN = 16
44
+
45
+
46
+ def new_salt() -> bytes:
47
+ return _new_salt()
48
+
49
+
50
+ def hash_secret(secret: str, salt: bytes) -> bytes:
51
+ return _hash_secret(secret, salt)
52
+
53
+
54
+ def verify_secret(secret: str, salt: bytes, expected_hash: bytes) -> bool:
55
+ if not salt or not expected_hash:
56
+ return False
57
+ return hmac.compare_digest(expected_hash, _hash_secret(secret, salt))
58
+
59
+
60
+ def generate_password() -> str:
61
+ return secrets.token_urlsafe(32)
62
+
63
+
64
+ def generate_token() -> str:
65
+ return BEARER_TOKEN_PREFIX + secrets.token_urlsafe(32)
66
+
67
+
68
+ def token_lookup_prefix(token: str) -> str:
69
+ return (token or "")[:TOKEN_LOOKUP_PREFIX_LEN]
70
+
71
+
72
+ def parse_basic_header(authorization: Optional[str]) -> Optional[Tuple[str, str]]:
73
+ if not authorization:
74
+ return None
75
+ parts = authorization.strip().split(" ", 1)
76
+ if len(parts) != 2 or parts[0].lower() != "basic":
77
+ return None
78
+ try:
79
+ decoded = base64.b64decode(parts[1].strip(), validate=True).decode("utf-8")
80
+ except (binascii.Error, UnicodeDecodeError, ValueError):
81
+ return None
82
+ if ":" not in decoded:
83
+ return None
84
+ return tuple(decoded.split(":", 1))
85
+
86
+
87
+ def parse_bearer_header(authorization: Optional[str]) -> Optional[str]:
88
+ if not authorization:
89
+ return None
90
+ parts = authorization.strip().split(" ", 1)
91
+ if len(parts) != 2 or parts[0].lower() != "bearer":
92
+ return None
93
+ return parts[1].strip() or None
94
+
95
+
96
+ class Authenticator(Protocol):
97
+ def __call__(self, authorization: Optional[str]) -> "AuthResult": ...
98
+
99
+
100
+ def make_basic_bearer_authenticator(
101
+ verify: Callable[[Optional[str]], Optional[str]],
102
+ ) -> Authenticator:
103
+ def _authenticate(authorization: Optional[str]) -> AuthResult:
104
+ identity = verify(authorization)
105
+ return AuthResult(ok=identity is not None, target_id=identity)
106
+
107
+ return _authenticate
108
+
109
+
110
+ def _hash_secret(secret: str, salt: bytes) -> bytes:
111
+ """Salted SHA-256 of a secret. Deterministic for a given salt."""
112
+ return hashlib.sha256(salt + secret.encode("utf-8")).digest()
113
+
114
+
115
+ def _new_salt() -> bytes:
116
+ return os.urandom(_SALT_LEN)
117
+
118
+
119
+ @dataclass
120
+ class _Credential:
121
+ """A stored credential: salt + hash of the secret, bound to a target_id."""
122
+
123
+ target_id: str
124
+ salt: bytes
125
+ secret_hash: bytes
126
+ revoked: bool = False
127
+
128
+ def matches(self, secret: str) -> bool:
129
+ """Constant-time check of a presented secret against the stored hash."""
130
+ if self.revoked:
131
+ return False
132
+ return hmac.compare_digest(self.secret_hash, _hash_secret(secret, self.salt))
133
+
134
+
135
+ @dataclass(init=False)
136
+ class AuthResult:
137
+ """Outcome of an :meth:`OTLPAuth.verify` call.
138
+
139
+ ``ok`` is True when the request may proceed. ``target_id`` is set (and
140
+ non-empty) when a credential was matched; it is ``None`` in ``none`` mode
141
+ and on failure. ``reason`` is a short, log-safe explanation on failure.
142
+ """
143
+
144
+ ok: bool
145
+ target_id: Optional[str] = None
146
+ reason: Optional[str] = None
147
+
148
+ def __init__(
149
+ self,
150
+ ok: Optional[bool] = None,
151
+ target_id: Optional[str] = None,
152
+ reason: Optional[str] = None,
153
+ *,
154
+ authenticated: Optional[bool] = None,
155
+ identity: Optional[str] = None,
156
+ ) -> None:
157
+ self.ok = bool(ok if ok is not None else authenticated)
158
+ self.target_id = target_id if target_id is not None else identity
159
+ self.reason = reason
160
+
161
+ @property
162
+ def authenticated(self) -> bool:
163
+ return self.ok
164
+
165
+ @property
166
+ def identity(self) -> Optional[str]:
167
+ return self.target_id
168
+
169
+
170
+ class OTLPAuth:
171
+ """Credential store + verifier for an OTLP receiver.
172
+
173
+ Construct with a mode, add credentials, then hand the instance to an
174
+ :class:`~simpleaudit.tracing.otlp.OTLPTraceReceiver` or
175
+ :class:`~simpleaudit.tracing.otlp.EphemeralOTLPReceiver`. The receiver
176
+ calls :meth:`verify` on each request's ``Authorization`` header.
177
+
178
+ Example::
179
+
180
+ auth = OTLPAuth(mode="basic")
181
+ auth.add_basic("sa_owui_17", "p4ss", target_id="owui_17")
182
+ result = auth.verify("Basic " + b64("sa_owui_17:p4ss"))
183
+ assert result.ok and result.target_id == "owui_17"
184
+ """
185
+
186
+ def __init__(self, mode: str = "none") -> None:
187
+ mode = (mode or "none").strip().lower()
188
+ if mode not in ("none", "basic", "bearer"):
189
+ raise ValueError(f"Unknown OTLPAuth mode: {mode!r} (expected 'none', 'basic', or 'bearer')")
190
+ self.mode = mode
191
+ # username -> credential (basic)
192
+ self._basic: Dict[str, _Credential] = {}
193
+ # token-hash-keyed credential (bearer). We index by a *lookup* hash of
194
+ # the token (salted per-credential) so we can find the record to
195
+ # compare against without storing the raw token.
196
+ self._bearer: Dict[str, _Credential] = {}
197
+
198
+ # ── Credential management ────────────────────────────────────────────
199
+
200
+ def add_basic(self, username: str, password: str, target_id: str) -> None:
201
+ """Register (or replace) a Basic-Auth credential for ``target_id``."""
202
+ if not username:
203
+ raise ValueError("username is required")
204
+ salt = _new_salt()
205
+ self._basic[username] = _Credential(
206
+ target_id=target_id, salt=salt, secret_hash=_hash_secret(password, salt)
207
+ )
208
+
209
+ def add_bearer(self, token: str, target_id: str) -> None:
210
+ """Register (or replace) a Bearer credential for ``target_id``."""
211
+ if not token:
212
+ raise ValueError("token is required")
213
+ salt = _new_salt()
214
+ cred = _Credential(target_id=target_id, salt=salt, secret_hash=_hash_secret(token, salt))
215
+ # Index by the salted hash so lookup needs only the presented token.
216
+ self._bearer[_hash_secret(token, salt).hex()] = cred
217
+
218
+ def revoke(self, username_or_token: str) -> bool:
219
+ """Revoke a credential by username (basic) or token (bearer).
220
+
221
+ Returns True if a matching credential was found and revoked.
222
+ """
223
+ if username_or_token in self._basic:
224
+ self._basic[username_or_token].revoked = True
225
+ return True
226
+ # For bearer we must find the record by hashing the token with each
227
+ # stored salt; the index key is that hash, so we can't reverse it.
228
+ # Instead we scan (credential sets are small — one per target).
229
+ for cred in self._bearer.values():
230
+ if cred.matches(username_or_token):
231
+ cred.revoked = True
232
+ return True
233
+ return False
234
+
235
+ def rotate_basic(self, username: str, new_password: str) -> None:
236
+ """Re-issue the password for an existing Basic credential (same target)."""
237
+ cred = self._basic.get(username)
238
+ if cred is None:
239
+ raise KeyError(f"No Basic credential for username {username!r}")
240
+ salt = _new_salt()
241
+ cred.salt = salt
242
+ cred.secret_hash = _hash_secret(new_password, salt)
243
+ cred.revoked = False
244
+
245
+ def rotate_bearer(self, old_token: str, new_token: str) -> None:
246
+ """Replace a Bearer token, keeping the same target_id."""
247
+ old_cred = self._find_bearer(old_token)
248
+ if old_cred is None:
249
+ raise KeyError("No Bearer credential for the given token")
250
+ # Remove the old index entry and register the new token under the
251
+ # same target_id.
252
+ for key, cred in list(self._bearer.items()):
253
+ if cred is old_cred:
254
+ del self._bearer[key]
255
+ break
256
+ self.add_bearer(new_token, old_cred.target_id)
257
+
258
+ def _find_bearer(self, token: str) -> Optional[_Credential]:
259
+ for cred in self._bearer.values():
260
+ if cred.matches(token):
261
+ return cred
262
+ return None
263
+
264
+ # ── Verification ─────────────────────────────────────────────────────
265
+
266
+ def verify(self, authorization_header: Optional[str]) -> AuthResult:
267
+ """Verify an ``Authorization`` header and return the target id.
268
+
269
+ In ``none`` mode this always succeeds (``target_id`` is ``None``).
270
+ """
271
+ if self.mode == "none":
272
+ return AuthResult(ok=True)
273
+
274
+ header = (authorization_header or "").strip()
275
+ if not header:
276
+ return AuthResult(ok=False, reason="missing_authorization_header")
277
+
278
+ parts = header.split(" ", 1)
279
+ if len(parts) != 2 or not parts[1].strip():
280
+ return AuthResult(ok=False, reason="malformed_authorization_header")
281
+ scheme, value = parts[0].strip().lower(), parts[1].strip()
282
+
283
+ # The scheme must match the configured mode.
284
+ if self.mode == "basic":
285
+ if scheme != "basic":
286
+ return AuthResult(ok=False, reason="expected_basic_scheme")
287
+ return self._verify_basic(value)
288
+ # bearer
289
+ if scheme != "bearer":
290
+ return AuthResult(ok=False, reason="expected_bearer_scheme")
291
+ return self._verify_bearer(value)
292
+
293
+ def _verify_basic(self, value: str) -> AuthResult:
294
+ # value should be base64("username:password")
295
+ try:
296
+ decoded = base64.b64decode(value, validate=True).decode("utf-8")
297
+ except (binascii.Error, UnicodeDecodeError, ValueError):
298
+ return AuthResult(ok=False, reason="invalid_basic_encoding")
299
+ if ":" not in decoded:
300
+ return AuthResult(ok=False, reason="invalid_basic_credentials")
301
+ username, password = decoded.split(":", 1)
302
+ cred = self._basic.get(username)
303
+ if cred is None:
304
+ return AuthResult(ok=False, reason="unknown_username")
305
+ if not cred.matches(password):
306
+ return AuthResult(ok=False, reason="invalid_credentials")
307
+ return AuthResult(ok=True, target_id=cred.target_id)
308
+
309
+ def _verify_bearer(self, token: str) -> AuthResult:
310
+ if not token:
311
+ return AuthResult(ok=False, reason="empty_bearer_token")
312
+ cred = self._find_bearer(token)
313
+ if cred is None:
314
+ return AuthResult(ok=False, reason="unknown_token")
315
+ return AuthResult(ok=True, target_id=cred.target_id)
316
+
317
+
318
+ def tag_spans(spans, target_id: Optional[str]) -> None:
319
+ """Stamp ``simpleaudit.target_id`` onto each span's attributes in place.
320
+
321
+ No-op when ``target_id`` is falsy (e.g. ``none`` mode). Mutates the raw
322
+ span dicts before they are normalized/stored, so the tag flows through
323
+ :func:`simpleaudit.tracing.otlp.parse_otlp_json` and is queryable via
324
+ ``SpanStore.by_attribute``.
325
+ """
326
+ if not target_id:
327
+ return
328
+ for span in spans:
329
+ attrs = span.get("attributes")
330
+ if attrs is None:
331
+ span["attributes"] = attrs = {}
332
+ attrs.setdefault(TARGET_ID_ATTR, target_id)
@@ -21,6 +21,7 @@ from __future__ import annotations
21
21
  import json
22
22
  from typing import TYPE_CHECKING, Any, Dict, List, Optional
23
23
 
24
+ from .auth import OTLPAuth, tag_spans
24
25
  from .store import SpanStore, normalize_span
25
26
 
26
27
  if TYPE_CHECKING: # pragma: no cover - annotation only
@@ -132,13 +133,16 @@ class OTLPTraceReceiver:
132
133
  self,
133
134
  store: Optional[SpanStore] = None,
134
135
  authenticator: Optional["Authenticator"] = None,
136
+ auth: Optional[OTLPAuth] = None,
135
137
  ) -> None:
138
+ self._default_store = store is None
136
139
  self.store = store or SpanStore()
137
140
  # Optional auth gate. When set, handle() calls it with the request's
138
141
  # Authorization header and rejects the export (401) on failure. When
139
142
  # None the receiver is open — the historical, backward-compatible
140
143
  # behaviour for the local, single-audit ephemeral case.
141
144
  self.authenticator = authenticator
145
+ self.auth = auth
142
146
 
143
147
  async def handle(self, body: Any, authorization: Optional[str] = None) -> Dict[str, Any]:
144
148
  """Ingest an OTLP JSON export body; return the OTLP ack payload.
@@ -147,14 +151,30 @@ class OTLPTraceReceiver:
147
151
  ``Authorization`` header) is checked first; a failed check returns a
148
152
  401 ``{"error": ...}`` payload and no spans are stored.
149
153
  """
150
- if self.authenticator is not None:
151
- result = self.authenticator(authorization)
152
- if not result.authenticated:
153
- return {"status": 401, "error": {"code": "unauthorized", "message": "Invalid or missing OTLP credentials."}}
154
+ result = self._authenticate(authorization)
155
+ if result is not None and not result.authenticated:
156
+ return {
157
+ "status": 401,
158
+ "authenticated": False,
159
+ "reason": getattr(result, "reason", None),
160
+ "error": {"code": "unauthorized", "message": "Invalid or missing OTLP credentials."},
161
+ }
154
162
  raw_spans = parse_otlp_json(body)
163
+ if result is not None:
164
+ tag_spans(raw_spans, result.identity)
155
165
  self.store.add_many(raw_spans)
156
166
  # OTLP ack: partialSuccess with the number of rejected spans (0 here).
157
- return {"partialSuccess": {"rejectedSpans": 0}}
167
+ ack = {"partialSuccess": {"rejectedSpans": 0}}
168
+ if self.auth is not None or self._default_store:
169
+ ack["authenticated"] = True
170
+ return ack
171
+
172
+ def _authenticate(self, authorization: Optional[str]) -> Any:
173
+ if self.auth is not None:
174
+ return self.auth.verify(authorization)
175
+ if self.authenticator is not None:
176
+ return self.authenticator(authorization)
177
+ return None
158
178
 
159
179
  def trace_ids(self) -> List[str]:
160
180
  seen: List[str] = []
@@ -196,6 +216,7 @@ class EphemeralOTLPReceiver:
196
216
  port: int = 0,
197
217
  store: Optional[SpanStore] = None,
198
218
  authenticator: Optional["Authenticator"] = None,
219
+ auth: Optional[OTLPAuth] = None,
199
220
  ) -> None:
200
221
  self.host = host
201
222
  self.port = port
@@ -204,6 +225,7 @@ class EphemeralOTLPReceiver:
204
225
  # checked against the request's Authorization header and rejected with
205
226
  # 401 on failure. None = open receiver (default).
206
227
  self.authenticator = authenticator
228
+ self.auth = auth
207
229
  self._runner: Optional[Any] = None
208
230
  self._site: Optional[Any] = None
209
231
  self._thread: Optional[Any] = None
@@ -225,13 +247,12 @@ class EphemeralOTLPReceiver:
225
247
  async def _handle_traces(self, request: Any) -> Any:
226
248
  from aiohttp import web
227
249
 
228
- if self.authenticator is not None:
229
- result = self.authenticator(request.headers.get("Authorization"))
230
- if not result.authenticated:
231
- return web.json_response(
232
- {"error": {"code": "unauthorized", "message": "Invalid or missing OTLP credentials."}},
233
- status=401,
234
- )
250
+ result = self._authenticate(request.headers.get("Authorization"))
251
+ if result is not None and not result.authenticated:
252
+ return web.json_response(
253
+ {"authenticated": False, "error": {"code": "unauthorized", "message": "Invalid or missing OTLP credentials."}},
254
+ status=401,
255
+ )
235
256
 
236
257
  content_type = request.headers.get("Content-Type", "")
237
258
  body_bytes = await request.read()
@@ -240,12 +261,21 @@ class EphemeralOTLPReceiver:
240
261
  raw_spans = _parse_otlp_http_protobuf(body_bytes)
241
262
  else:
242
263
  raw_spans = parse_otlp_json(body_bytes.decode("utf-8"))
264
+ if result is not None:
265
+ tag_spans(raw_spans, result.identity)
243
266
  self.store.add_many(raw_spans)
244
267
  except Exception:
245
268
  # Never fail the export; ack with a rejection count so the target
246
269
  # doesn't retry-loop. The audit continues regardless.
247
270
  return web.json_response({"partialSuccess": {"rejectedSpans": 1}}, status=200)
248
- return web.json_response({"partialSuccess": {"rejectedSpans": 0}}, status=200)
271
+ return web.json_response({"partialSuccess": {"rejectedSpans": 0}, "authenticated": True}, status=200)
272
+
273
+ def _authenticate(self, authorization: Optional[str]) -> Any:
274
+ if self.auth is not None:
275
+ return self.auth.verify(authorization)
276
+ if self.authenticator is not None:
277
+ return self.authenticator(authorization)
278
+ return None
249
279
 
250
280
  def _serve(self, loop: Any) -> None:
251
281
  import asyncio