From 5bede49e723db1ee19aa2affe5586ec120ebf3dc Mon Sep 17 00:00:00 2001 From: Namrata Ghadi Date: Thu, 8 Oct 2026 12:06:43 -0700 Subject: [PATCH] add retriever, trace and session --- CHANGELOG.md | 2 ++ openapi.yaml | 3 +++ .../handlers/agent_control/bridge.py | 2 +- src/splunk_ao/logger/control.py | 23 ++++++++++++++----- src/splunk_ao/logger/logger.py | 1 + tests/test_agent_control.py | 9 ++++++++ tests/test_agent_control_bridge.py | 20 ++++++++++++---- tests/test_attribute_mapping.py | 7 +++--- tests/test_span_converter.py | 6 +++-- 9 files changed, 57 insertions(+), 16 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index dd51d1ce..d3c9cf47 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- Agent Control bridges now preserve all supported `applies_to` values: + `llm_call`, `tool_call`, `retriever_call`, `trace_call`, and `session_call`. - `SplunkAOLogger.add_llm_span()` and `add_single_llm_span_trace()` accept `image_input_tokens`, `audio_input_tokens`, `audio_output_tokens`, and `image_output_tokens`: the image/audio share of the input and output token diff --git a/openapi.yaml b/openapi.yaml index 14fd3ddb..3fcc4797 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -15136,6 +15136,9 @@ components: enum: - llm_call - tool_call + - retriever_call + - trace_call + - session_call title: ControlAppliesTo ControlCheckStage: type: string diff --git a/src/splunk_ao/handlers/agent_control/bridge.py b/src/splunk_ao/handlers/agent_control/bridge.py index 74b2ed12..b77a6e26 100644 --- a/src/splunk_ao/handlers/agent_control/bridge.py +++ b/src/splunk_ao/handlers/agent_control/bridge.py @@ -294,7 +294,7 @@ def _control_span_kwargs(event: Any) -> dict[str, Any]: "control_id": getattr(event, "control_id", None), "agent_name": getattr(event, "agent_name", None), "check_stage": ControlCheckStage(check_stage) if check_stage else None, - "applies_to": ControlAppliesTo(applies_to) if applies_to else None, + "applies_to": ControlAppliesTo(getattr(applies_to, "value", applies_to)) if applies_to else None, "evaluator_name": getattr(event, "evaluator_name", None), "selector_path": getattr(event, "selector_path", None), } diff --git a/src/splunk_ao/logger/control.py b/src/splunk_ao/logger/control.py index 09af1b69..8a7875d8 100644 --- a/src/splunk_ao/logger/control.py +++ b/src/splunk_ao/logger/control.py @@ -9,18 +9,29 @@ from galileo_core.schemas.logging.step import Metrics + +class ControlAppliesTo(StrEnum): + llm_call = "llm_call" + tool_call = "tool_call" + retriever_call = "retriever_call" + trace_call = "trace_call" + session_call = "session_call" + + try: - from galileo_core.schemas.logging.control import ControlAppliesTo, ControlCheckStage, ControlResult - from galileo_core.schemas.logging.span import ControlSpan + from galileo_core.schemas.logging.control import ControlCheckStage, ControlResult + from galileo_core.schemas.logging.span import ControlSpan as _CoreControlSpan HAS_NATIVE_CONTROL_SPAN = True + + class ControlSpan(_CoreControlSpan): + """ControlSpan whose applies_to field supports Agent Control event values.""" + + applies_to: ControlAppliesTo | None = Field(default=None, description="Parent execution type.") + except ImportError: HAS_NATIVE_CONTROL_SPAN = False - class ControlAppliesTo(StrEnum): - llm_call = "llm_call" - tool_call = "tool_call" - class ControlCheckStage(StrEnum): pre = "pre" post = "post" diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index 2089ac9e..033b454a 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -1841,6 +1841,7 @@ def add_control_span( "status_code": status_code, "metrics": Metrics(duration_ns=duration_ns), "id": span_id or uuid.uuid4(), + "session_id": self.session_id, "trace_id": trace_id, "parent_id": parent_id, "step_number": step_number, diff --git a/tests/test_agent_control.py b/tests/test_agent_control.py index 6bb6e09c..cd5ac4f2 100644 --- a/tests/test_agent_control.py +++ b/tests/test_agent_control.py @@ -7,6 +7,7 @@ from splunk_ao import AgentControlTarget, AgentControlTargetUnresolvedError, get_agent_control_target from splunk_ao.constants import DEFAULT_AGENT_STREAM_NAME, DEFAULT_PROJECT_NAME from splunk_ao.decorator import splunk_ao_context +from splunk_ao.logger.control import ControlAppliesTo from splunk_ao.utils.singleton import SplunkAOLoggerSingleton @@ -42,6 +43,14 @@ def _stub_cached_loggers(monkeypatch, loggers: dict[tuple[str, ...], SimpleNames monkeypatch.setattr(SplunkAOLoggerSingleton, "get_all_loggers", lambda self: loggers) +def test_control_applies_to_supports_agent_control_values() -> None: + # Given: the supported Agent Control parent execution types + applies_to = {value.value for value in ControlAppliesTo} + + # Then: the SDK preserves all values without adding unsupported types + assert applies_to == {"llm_call", "tool_call", "retriever_call", "trace_call", "session_call"} + + def test_get_agent_control_target_uses_explicit_log_stream_id() -> None: # Given: an explicit Galileo log stream ID agent_stream_id = str(uuid4()) diff --git a/tests/test_agent_control_bridge.py b/tests/test_agent_control_bridge.py index fa916339..bed25b11 100644 --- a/tests/test_agent_control_bridge.py +++ b/tests/test_agent_control_bridge.py @@ -13,7 +13,7 @@ from splunk_ao.exporter.span_transform import copy_span_for_export from splunk_ao.handlers.agent_control import setup_agent_control_bridge -from splunk_ao.logger.control import ControlResult, ControlSpan +from splunk_ao.logger.control import ControlAppliesTo, ControlResult, ControlSpan from splunk_ao.logger.logger import SplunkAOLogger from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client @@ -320,8 +320,13 @@ def test_idle_new_logger_does_not_mask_active_logger_context( @patch("splunk_ao.logger.logger.AgentStreams") @patch("splunk_ao.logger.logger.Projects") @patch("splunk_ao.logger.logger.Traces") +@pytest.mark.parametrize("applies_to", [value.value for value in ControlAppliesTo]) def test_agent_control_event_converts_to_control_span_in_batch_mode( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, fake_agent_control_modules + mock_traces_client: Mock, + mock_projects_client: Mock, + mock_logstreams_client: Mock, + fake_agent_control_modules, + applies_to: str, ) -> None: # Given: a batch logger with an active parent and a matching Agent Control event mock_traces_client_instance = setup_mock_traces_client(mock_traces_client) @@ -330,8 +335,10 @@ def test_agent_control_event_converts_to_control_span_in_batch_mode( logger = SplunkAOLogger(project="my_project", agent_stream="my_log_stream") logger.start_trace(input="trace input") workflow = logger.add_workflow_span(input="workflow input", name="workflow") + session_id = str(uuid.uuid4()) + logger.set_session(session_id) bridge = setup_agent_control_bridge(logger) - event = _make_event(logger) + event = _make_event(logger, applies_to=applies_to) # When: the bridge receives the event through the public sink contract result = bridge.write_events([event]) @@ -345,7 +352,10 @@ def test_agent_control_event_converts_to_control_span_in_batch_mode( assert isinstance(control_span, ControlSpan) assert control_span.id == uuid.UUID(event.control_execution_id) assert control_span.trace_id == logger.traces[0].id + assert control_span.session_id == uuid.UUID(session_id) assert control_span.parent_id == workflow.id + assert control_span.applies_to == ControlAppliesTo(applies_to) + assert control_span.model_dump(mode="json")["applies_to"] == applies_to assert control_span.name == "toxicity-guardrail" assert control_span.input == "selected text" assert control_span.output == ControlResult(action="observe", matched=True, confidence=0.91) @@ -380,7 +390,7 @@ def test_agent_control_event_converts_to_control_span_in_batch_mode( "agent_control.control_name": "toxicity-guardrail", "agent_control.agent_name": "assistant", "agent_control.check_stage": "pre", - "agent_control.applies_to": "llm_call", + "agent_control.applies_to": applies_to, "agent_control.evaluator_name": "regex", "agent_control.selector_path": "input", "agent_control.action": "observe", @@ -388,12 +398,14 @@ def test_agent_control_event_converts_to_control_span_in_batch_mode( "agent_control.confidence": 0.91, } assert exported_attrs["gen_ai.operation.name"] == "control" + assert exported_attrs["gen_ai.conversation.id"] == session_id assert "splunk_ao.operation.name" not in exported_attrs assert json.loads(exported_attrs["gen_ai.input.messages"]) == [ {"parts": [{"content": "selected text", "type": "text"}], "role": "user"} ] assert exported.context == emitted.context assert exported.parent == logger._otel_ids[workflow.id].span_context + assert exported.context.trace_id == logger._otel_ids[workflow.id].span_context.trace_id epoch = datetime.datetime(1970, 1, 1, tzinfo=datetime.UTC) elapsed = event.timestamp - epoch expected_start_ns = ((elapsed.days * 86_400 + elapsed.seconds) * 1_000_000_000) + elapsed.microseconds * 1_000 diff --git a/tests/test_attribute_mapping.py b/tests/test_attribute_mapping.py index 8d88d7a7..c2b1b668 100644 --- a/tests/test_attribute_mapping.py +++ b/tests/test_attribute_mapping.py @@ -852,7 +852,8 @@ def test_orchestration_keeps_non_json_strings_as_text_messages() -> None: ] -def test_control_mapping_exports_fully_populated_backend_contract() -> None: +@pytest.mark.parametrize("applies_to", [value.value for value in ControlAppliesTo]) +def test_control_mapping_exports_fully_populated_backend_contract(applies_to: str) -> None: span = LoggedControlSpan( name="PII Guard", input="question", @@ -860,7 +861,7 @@ def test_control_mapping_exports_fully_populated_backend_contract() -> None: control_id=42, agent_name="planner", check_stage=ControlCheckStage.pre, - applies_to=ControlAppliesTo.llm_call, + applies_to=ControlAppliesTo(applies_to), evaluator_name="pii-check", selector_path="$.input", tags=["agent_control", "control"], @@ -877,7 +878,7 @@ def test_control_mapping_exports_fully_populated_backend_contract() -> None: "agent_control.control_name": "PII Guard", "agent_control.agent_name": "planner", "agent_control.check_stage": "pre", - "agent_control.applies_to": "llm_call", + "agent_control.applies_to": applies_to, "agent_control.evaluator_name": "pii-check", "agent_control.selector_path": "$.input", "agent_control.action": "deny", diff --git a/tests/test_span_converter.py b/tests/test_span_converter.py index 61d93b06..dc75ef93 100644 --- a/tests/test_span_converter.py +++ b/tests/test_span_converter.py @@ -165,7 +165,8 @@ def test_converter_leaves_final_export_normalization_to_the_sink() -> None: assert "splunk_ao.system" not in attributes -def test_fully_populated_logged_control_span_preserves_otlp_envelope_and_attributes() -> None: +@pytest.mark.parametrize("applies_to", [value.value for value in ControlAppliesTo]) +def test_fully_populated_logged_control_span_preserves_otlp_envelope_and_attributes(applies_to: str) -> None: parent_context = make_context(span_id=PARENT_SPAN_ID) span = LoggedControlSpan( name="PII Guard", @@ -181,7 +182,7 @@ def test_fully_populated_logged_control_span_preserves_otlp_envelope_and_attribu control_id=42, agent_name="planner", check_stage=ControlCheckStage.pre, - applies_to=ControlAppliesTo.llm_call, + applies_to=ControlAppliesTo(applies_to), evaluator_name="pii-check", selector_path="$.input", ) @@ -197,6 +198,7 @@ def test_fully_populated_logged_control_span_preserves_otlp_envelope_and_attribu assert result.status.status_code is StatusCode.ERROR assert attrs["galileo.span.kind"] == "control" assert attrs["agent_control.control_id"] == 42 + assert attrs["agent_control.applies_to"] == applies_to assert attrs["agent_control.action"] == "deny" assert attrs["agent_control.matched"] is False assert attrs["agent_control.error_message"] == "blocked"