diff --git a/pyproject.toml b/pyproject.toml index a43c67d..0c44a64 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "simpleaudit" -version = "0.3.1" +version = "0.3.0" description = "Lightweight AI Safety Auditing Framework" readme = "README.md" license = {text = "MIT"} diff --git a/simpleaudit/targets/model.py b/simpleaudit/targets/model.py index 03ad27e..8d2a121 100644 --- a/simpleaudit/targets/model.py +++ b/simpleaudit/targets/model.py @@ -88,6 +88,7 @@ async def send( documents=documents, response_format=response_format, params=params, + context=context, max_retries=self.max_retries, retry_backoff=self.retry_backoff, ) @@ -104,6 +105,7 @@ async def _client_send( documents: Optional[List[Union[str, Dict[str, Any]]]], response_format: Optional[Dict[str, Any]], params: Optional[Dict[str, Any]], + context: Optional[TargetContext] = None, max_retries: int, retry_backoff: float, ) -> TargetResponse: @@ -111,9 +113,21 @@ async def _client_send( Imported lazily so that ``simpleaudit.targets`` does not create an import cycle with ``model_auditor`` at package import time. + + When the per-turn context carries W3C trace headers (e.g. the engine's + scenario traceparent), they are merged into ``extra_headers`` so the + target's outgoing call propagates the trace context. This is what lets an + instrumented target export spans under the same trace id the engine's + TraceCorrelation records. User-supplied extra_headers are preserved. """ from ..model_auditor import ModelAuditor + if context is not None and context.trace_headers: + headers = dict(params.get("extra_headers", {}) if params else {}) + headers.update(context.trace_headers) + params = dict(params or {}) + params["extra_headers"] = headers + content, input_tokens, output_tokens = await ModelAuditor._call_async( client, model, diff --git a/tests/test_targets.py b/tests/test_targets.py index e67add5..d0c57cc 100644 --- a/tests/test_targets.py +++ b/tests/test_targets.py @@ -75,6 +75,40 @@ async def test_model_target_client_mode(): assert client.calls[0]["model"] == "m" +@pytest.mark.asyncio +async def test_model_target_client_mode_forwards_trace_headers(): + """trace_headers on the per-turn context reach the wire as extra_headers.""" + client = _FakeClient() + t = ModelTarget(client=client, model="m") + ctx = TargetContext(trace_headers={"traceparent": "00-11223344556677889900112233445566-abcdef0123456789-01"}) + await t.send(user="hi", context=ctx) + assert client.calls[0]["extra_headers"] == { + "traceparent": "00-11223344556677889900112233445566-abcdef0123456789-01" + } + + +@pytest.mark.asyncio +async def test_model_target_trace_headers_merge_with_user_params(): + """User-supplied extra_headers and the framework's trace headers coexist.""" + client = _FakeClient() + t = ModelTarget(client=client, model="m") + ctx = TargetContext(trace_headers={"traceparent": "00-11223344556677889900112233445566-abcdef0123456789-01"}) + await t.send(user="hi", params={"extra_headers": {"X-Custom": "1"}}, context=ctx) + assert client.calls[0]["extra_headers"] == { + "X-Custom": "1", + "traceparent": "00-11223344556677889900112233445566-abcdef0123456789-01", + } + + +@pytest.mark.asyncio +async def test_model_target_no_context_no_extra_headers(): + """Without a context the call is byte-identical to the legacy path.""" + client = _FakeClient() + t = ModelTarget(client=client, model="m") + await t.send(user="hi") + assert "extra_headers" not in client.calls[0] + + @pytest.mark.asyncio async def test_model_target_transport_mode(): async def transport(**kw): @@ -324,3 +358,38 @@ async def send(self, *, user, history=None, context=None, **kw): assert parts[3] == "01" # The correlation recorded at least one turn -> trace link. assert correlation.all_trace_ids() + + +def test_engine_traceparent_reaches_wire_on_default_target(): + """With the default ModelTarget (no override), the scenario's traceparent + is forwarded on the target's outgoing acompletion as extra_headers.""" + import asyncio + + from simpleaudit.tracing.context import TraceCorrelation + from tests.fakes import fixed_probe_auditor, fixed_severity_judge, make_auditor + + client = _FakeClient() + auditor = make_auditor( + target=client, + judge=fixed_severity_judge("pass"), + auditor=fixed_probe_auditor("probe"), + max_turns=2, + show_progress=False, + ) + + correlation = TraceCorrelation(audit_run_id="audit_test") + scenarios = [{"name": "Wire", "description": "wire traceparent test"}] + asyncio.run( + auditor.run_async( + scenarios=scenarios, + max_turns=2, + audit_run_id="audit_test", + trace_correlation=correlation, + ) + ) + + assert client.calls + last = client.calls[-1] + tp = last["extra_headers"]["traceparent"] + # Same trace id the correlation recorded, not a fresh unrelated one. + assert tp.split("-")[1] in correlation.all_trace_ids() diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 552ec76..084b09a 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -422,6 +422,26 @@ def test_parse_otlp_json_string_body(): assert len(spans) == 1 +def test_parse_otlp_json_preserves_genai_content_attributes(): + payload = _otlp_http_payload() + payload["resourceSpans"][0]["scopeSpans"][0]["spans"][0]["attributes"].extend( + [ + {"key": "gen_ai.input.messages", "value": {"stringValue": "user prompt"}}, + {"key": "gen_ai.output.messages", "value": {"stringValue": "assistant output"}}, + {"key": "gen_ai.tool.call.arguments", "value": {"stringValue": '{"x":1}'}}, + { + "key": "gen_ai.retrieval.documents", + "value": {"arrayValue": {"values": [{"stringValue": "document text"}]}}, + }, + ] + ) + span = parse_otlp_json(payload)[0] + assert span["attributes"]["gen_ai.input.messages"] == "user prompt" + assert span["attributes"]["gen_ai.output.messages"] == "assistant output" + assert span["attributes"]["gen_ai.tool.call.arguments"] == '{"x":1}' + assert span["attributes"]["gen_ai.retrieval.documents"] == ["document text"] + + @pytest.mark.asyncio async def test_otlp_receiver_ingests(): receiver = OTLPTraceReceiver()