diff --git a/src/splunk_ao/converter/attribute_mapping.py b/src/splunk_ao/converter/attribute_mapping.py index 3b84f80..3fe5c45 100644 --- a/src/splunk_ao/converter/attribute_mapping.py +++ b/src/splunk_ao/converter/attribute_mapping.py @@ -127,11 +127,12 @@ def _text_part(value: Any) -> dict[str, Any]: def _content_part(value: Any) -> dict[str, Any] | None: part = _mapping_value(value) - if part is None or "type" not in part: + if part is None: return None - part_type = _json_compatible(part["type"]) - part["type"] = str(part_type) + part_type = part.get("type") + if not isinstance(part_type, str) or not part_type.strip(): + return None if part_type == "text" and "content" not in part and "text" in part: part["content"] = part.pop("text") return part @@ -142,7 +143,7 @@ def _content_parts(value: Any) -> list[dict[str, Any]]: if part is not None: return [part] - if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray): + if value and isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray): parts = [_content_part(item) for item in value] if all(item is not None for item in parts): return [item for item in parts if item is not None] @@ -204,7 +205,7 @@ def _message_sequence(value: Any, default_role: str) -> list[dict[str, Any]]: return [_message(value, default_role)] -def _message_container(value: Any) -> tuple[Any | None, bool]: +def _message_container(value: Any) -> tuple[Any | None, bool, bool]: parsed = _parse_json_string(value) source = _mapping_view(parsed) @@ -214,35 +215,44 @@ def _message_container(value: Any) -> tuple[Any | None, bool]: if container is None or "messages" not in container: continue messages = _parse_json_string(container["messages"]) - if _is_message_sequence(messages): - return messages, container is source + if _is_message_sequence(messages, allow_empty=True): + return messages, container is source, True if "role" in source or "type" in source: - return parsed, False - return None, False + return parsed, False, True + return parsed, False, False if isinstance(parsed, Sequence) and not isinstance(parsed, str | bytes | bytearray): if _is_message_sequence(parsed) or _is_content_part_sequence(parsed): - return parsed, False - return (parsed, False) if not isinstance(value, str | bytes | bytearray) else (None, False) + return parsed, False, True + return parsed, False, False - return parsed, False + return parsed, False, False -def _is_message_sequence(value: Any) -> bool: +def _is_message_sequence(value: Any, *, allow_empty: bool = False) -> bool: if not isinstance(value, Sequence) or isinstance(value, str | bytes | bytearray): message = _mapping_view(value) return message is not None and "role" in message - return not value or all((message := _mapping_view(item)) is not None and "role" in message for item in value) + return (allow_empty and not value) or ( + bool(value) and all((message := _mapping_view(item)) is not None and "role" in message for item in value) + ) def _is_content_part_sequence(value: Any) -> bool: - return bool(value) and all((part := _mapping_view(item)) is not None and "type" in part for item in value) + return bool(value) and all(_content_part(item) is not None for item in value) def _orchestration_messages(value: Any, default_role: str) -> tuple[list[dict[str, Any]] | None, bool]: - container, full_history = _message_container(value) - messages = None if container is None else _message_sequence(container, default_role) + container, full_history, recognized_messages = _message_container(value) + if container is None: + messages = None + elif recognized_messages and isinstance(container, Sequence) and not container: + messages = [] + elif isinstance(container, Sequence) and not isinstance(container, str | bytes | bytearray) and not container: + messages = [_message(container, default_role)] + else: + messages = _message_sequence(container, default_role) return messages, full_history @@ -428,8 +438,7 @@ def _set_orchestration_content(attrs: MutableMapping[str, AttributeValue], span: return if full_history and input_messages is not None and output_messages[: len(input_messages)] == input_messages: output_messages = output_messages[len(input_messages) :] - if output_messages: - attrs["gen_ai.output.messages"] = _json_string(_with_finish_reasons(output_messages)) + attrs["gen_ai.output.messages"] = _json_string(_with_finish_reasons(output_messages)) def set_workflow_attributes(attrs: MutableMapping[str, AttributeValue], span: WorkflowSpan) -> None: @@ -495,6 +504,7 @@ def normalize_attributes_for_export( ) -> dict[str, AttributeValue]: """Return final wire attributes without mutating the source mapping.""" result = dict(attrs) + result[SPLUNK_AO_SYSTEM] = SPLUNK_AO_SYSTEM_VALUE if not enabled: return result @@ -507,5 +517,4 @@ def normalize_attributes_for_export( for source_key in CONTENT_ALIAS_BY_GEN_AI: result.pop(source_key, None) - result[SPLUNK_AO_SYSTEM] = SPLUNK_AO_SYSTEM_VALUE return result diff --git a/src/splunk_ao/exporter/span_transform.py b/src/splunk_ao/exporter/span_transform.py index 60daab2..c8bac6c 100644 --- a/src/splunk_ao/exporter/span_transform.py +++ b/src/splunk_ao/exporter/span_transform.py @@ -22,16 +22,16 @@ ) _NORMALIZATION_ENV = "SPLUNK_AO_DEV_ENABLE_ATTRIBUTE_NORMALIZATION" -_FALSE_VALUES = frozenset({"0", "false", "no", "off"}) +_TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) def _normalization_enabled() -> bool: value = os.environ.get(_NORMALIZATION_ENV) - return value is None or value.strip().lower() not in _FALSE_VALUES + return value is not None and value.strip().lower() in _TRUE_VALUES def copy_span_for_export( - span: ReadableSpan, routing_resource: Resource | None = None, *, normalize_attributes: bool = True + span: ReadableSpan, routing_resource: Resource | None = None, *, normalize_attributes: bool = False ) -> ReadableSpan: """Return an immutable span copy with final attributes and routing.""" source_attributes = { diff --git a/tests/test_attribute_mapping.py b/tests/test_attribute_mapping.py index 3e4cc85..9d983fd 100644 --- a/tests/test_attribute_mapping.py +++ b/tests/test_attribute_mapping.py @@ -397,17 +397,88 @@ def test_orchestration_preserves_schema_valid_parts_and_tool_calls() -> None: ] -def test_orchestration_does_not_label_arbitrary_state_as_messages() -> None: - span = WorkflowSpan( - name="state-machine", - input=json.dumps({"current_agent": "coordinator", "travellers": 2}), - output=json.dumps({"next_agent": "flight_specialist"}), - ) - +@pytest.mark.parametrize( + "span", + [ + WorkflowSpan( + name="state-machine", + input=json.dumps({"current_agent": "coordinator", "travellers": 2}), + output=json.dumps({"next_agent": "flight_specialist"}), + ), + AgentSpan( + name="state-agent", + agent_type=AgentType.planner, + input=json.dumps({"current_agent": "coordinator", "travellers": 2}), + output=json.dumps({"next_agent": "flight_specialist"}), + ), + ], +) +def test_orchestration_preserves_arbitrary_state_as_messages(span: WorkflowSpan | AgentSpan) -> None: attrs = build_span_attributes(span) - assert "gen_ai.input.messages" not in attrs - assert "gen_ai.output.messages" not in attrs + assert json.loads(attrs["gen_ai.input.messages"]) == [ + _text_message("user", '{"current_agent":"coordinator","travellers":2}') + ] + assert json.loads(attrs["gen_ai.output.messages"]) == [ + _text_message("assistant", '{"next_agent":"flight_specialist"}', finish_reason="unknown") + ] + + +@pytest.mark.parametrize( + ("serialized", "expected_content"), [("false", "false"), ("0", "0"), ("{}", "{}"), ("[]", "[]")] +) +def test_orchestration_preserves_false_zero_and_empty_values(serialized: str, expected_content: str) -> None: + attrs = build_span_attributes(WorkflowSpan(name="workflow", input=serialized, output=serialized)) + + assert json.loads(attrs["gen_ai.input.messages"]) == [_text_message("user", expected_content)] + assert json.loads(attrs["gen_ai.output.messages"]) == [ + _text_message("assistant", expected_content, finish_reason="unknown") + ] + + +def test_orchestration_preserves_explicit_empty_message_history() -> None: + attrs = build_span_attributes(WorkflowSpan(name="workflow", input='{"messages":[]}', output='{"messages":[]}')) + + assert json.loads(attrs["gen_ai.input.messages"]) == [] + assert json.loads(attrs["gen_ai.output.messages"]) == [] + + +def test_orchestration_preserves_empty_full_history_suffix() -> None: + history = '{"messages":[{"role":"user","content":"question"}]}' + attrs = build_span_attributes(AgentSpan(name="agent", input=history, output=history)) + + assert json.loads(attrs["gen_ai.output.messages"]) == [] + + +def test_orchestration_falls_back_for_malformed_message_container() -> None: + value = '{"messages":[{"content":"missing role"}],"state":"kept"}' + attrs = build_span_attributes(WorkflowSpan(name="workflow", input=value, output=value)) + + assert json.loads(attrs["gen_ai.input.messages"]) == [_text_message("user", value)] + assert json.loads(attrs["gen_ai.output.messages"]) == [_text_message("assistant", value, finish_reason="unknown")] + + +def test_orchestration_preserves_typed_extension_part_without_inventing_fields() -> None: + parts = [ + {"type": "audio", "url": "https://example.com/answer.wav", "format": "wav"}, + {"type": "file", "file_id": "file-1"}, + ] + serialized_parts = json.dumps(parts) + attrs = build_span_attributes(AgentSpan(name="agent", input=serialized_parts, output=serialized_parts)) + + assert json.loads(attrs["gen_ai.input.messages"]) == [{"role": "user", "parts": parts}] + assert json.loads(attrs["gen_ai.output.messages"]) == [ + {"role": "assistant", "parts": parts, "finish_reason": "unknown"} + ] + + +def test_orchestration_falls_back_when_part_type_is_empty() -> None: + value = [{"type": "", "content": "not a valid typed part"}] + attrs = build_span_attributes(WorkflowSpan(name="workflow", input=json.dumps(value))) + + assert json.loads(attrs["gen_ai.input.messages"]) == [ + _text_message("user", json.dumps(value, separators=(",", ":"), sort_keys=True)) + ] def test_orchestration_keeps_non_json_strings_as_text_messages() -> None: @@ -510,7 +581,7 @@ def test_normalizer_can_be_disabled_for_developer_comparison() -> None: "splunk_ao.system": "source-value", } - assert normalize_attributes_for_export(source, enabled=False) == source + assert normalize_attributes_for_export(source, enabled=False) == {**source, "splunk_ao.system": "splunk_ao_python"} def test_every_alias_uses_an_explicit_destination_namespace() -> None: diff --git a/tests/test_decorator_distributed.py b/tests/test_decorator_distributed.py index f872c57..a5334b7 100644 --- a/tests/test_decorator_distributed.py +++ b/tests/test_decorator_distributed.py @@ -109,6 +109,27 @@ def workflow(input_value: str) -> str: distributed_clients.update_span.assert_not_called() +def test_decorator_workflow_arguments_and_result_reach_otel_span( + reset_context: None, distributed_clients: Mock +) -> None: + logger = init_logger() + + @log(span_type="workflow") + def add(arg1: int, arg2: int) -> dict[str, int]: + return {"sum": arg1 + arg2} + + assert add(1, 2) == {"sum": 3} + + [span] = logger._sink.spans + attrs = span.attributes or {} + assert json.loads(attrs["gen_ai.input.messages"]) == [ + {"role": "user", "parts": [{"type": "text", "content": '{"arg1":1,"arg2":2}'}]} + ] + assert json.loads(attrs["gen_ai.output.messages"]) == [ + {"role": "assistant", "parts": [{"type": "text", "content": '{"sum":3}'}], "finish_reason": "unknown"} + ] + + def test_workflow_empty_output_is_preserved(reset_context: None, distributed_clients: Mock) -> None: logger = init_logger() diff --git a/tests/test_otel_native_paths.py b/tests/test_otel_native_paths.py index f655718..394f8d4 100644 --- a/tests/test_otel_native_paths.py +++ b/tests/test_otel_native_paths.py @@ -12,11 +12,11 @@ from opentelemetry.trace.status import Status, StatusCode from splunk_ao.decorator import ( + _agent_stream_context, _dataset_input_context, _dataset_metadata_context, _dataset_output_context, _experiment_id_context, - _agent_stream_context, _project_context, _session_id_context, ) @@ -295,9 +295,9 @@ def test_exporter_preserves_every_unaffected_span_field() -> None: assert getattr(exported, field) == getattr(source, field) assert exported.resource.schema_url == source.resource.schema_url assert exported.attributes["gen_ai.request.model"] == "gpt-4o" - assert exported.attributes["splunk_ao.request.model"] == "gpt-4o" + assert "splunk_ao.request.model" not in exported.attributes assert exported.attributes["gen_ai.provider.name"] == "openai" - assert exported.attributes["splunk_ao.provider.name"] == "openai" + assert "splunk_ao.provider.name" not in exported.attributes assert exported.attributes["gen_ai.system"] == "legacy-upstream-provider" assert exported.attributes["splunk_ao.system"] == "splunk_ao_python" assert exported.attributes["custom.attribute"] == "preserved" diff --git a/tests/test_span_transform.py b/tests/test_span_transform.py index 19852db..8f61c88 100644 --- a/tests/test_span_transform.py +++ b/tests/test_span_transform.py @@ -63,7 +63,9 @@ def test_copy_span_for_export_is_immutable_and_combines_normalization_with_routi source_resource = source.resource exported = copy_span_for_export( - source, Resource({"splunk_ao.project.id": "project-id", "splunk_ao.logstream.id": "log-stream-id"}) + source, + Resource({"splunk_ao.project.id": "project-id", "splunk_ao.logstream.id": "log-stream-id"}), + normalize_attributes=True, ) assert exported is not source @@ -112,7 +114,9 @@ def test_copy_span_preserves_unaffected_readable_span_fields() -> None: def test_exporter_normalizes_once_and_delegates_lifecycle_once() -> None: delegate = RecordingExporter() - exporter = NormalizingSpanExporter(delegate, Resource({"splunk_ao.project.name": "project"})) + exporter = NormalizingSpanExporter( + delegate, Resource({"splunk_ao.project.name": "project"}), normalize_attributes=True + ) assert exporter.export((make_span(),)) == SpanExportResult.SUCCESS assert exporter.force_flush(1234) is True @@ -137,8 +141,7 @@ def test_private_environment_switch_disables_only_attribute_normalization( exporter.export((source,)) exported = delegate.batches[0][0] - assert exported.attributes == source.attributes - assert "splunk_ao.system" not in exported.attributes + assert exported.attributes == {**source.attributes, "splunk_ao.system": "splunk_ao_python"} assert exported.resource.attributes["splunk_ao.project.name"] == "project" @@ -153,11 +156,30 @@ def test_private_environment_switch_is_read_at_exporter_construction(monkeypatch assert "splunk_ao.request.model" not in delegate.batches[0][0].attributes -def test_default_normalization_is_enabled(monkeypatch: pytest.MonkeyPatch) -> None: +def test_default_normalization_is_disabled_but_sdk_marker_is_present(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("SPLUNK_AO_DEV_ENABLE_ATTRIBUTE_NORMALIZATION", raising=False) delegate = RecordingExporter() exporter = NormalizingSpanExporter(delegate) exporter.export((make_span(),)) - assert delegate.batches[0][0].attributes["splunk_ao.system"] == "splunk_ao_python" + attributes = delegate.batches[0][0].attributes + assert attributes["gen_ai.request.model"] == "gpt-4o" + assert "splunk_ao.request.model" not in attributes + assert attributes["splunk_ao.system"] == "splunk_ao_python" + + +@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "on"]) +def test_private_environment_switch_explicitly_enables_normalization( + monkeypatch: pytest.MonkeyPatch, value: str +) -> None: + monkeypatch.setenv("SPLUNK_AO_DEV_ENABLE_ATTRIBUTE_NORMALIZATION", value) + delegate = RecordingExporter() + exporter = NormalizingSpanExporter(delegate) + + exporter.export((make_span({"gen_ai.input.messages": "input-json"}),)) + + attributes = delegate.batches[0][0].attributes + assert "gen_ai.input.messages" not in attributes + assert attributes["splunk_ao.input.messages"] == "input-json" + assert attributes["splunk_ao.system"] == "splunk_ao_python"