diff --git a/src/splunk_ao/handlers/openai_agents/handler.py b/src/splunk_ao/handlers/openai_agents/handler.py index 9ce91da..4da18cb 100644 --- a/src/splunk_ao/handlers/openai_agents/handler.py +++ b/src/splunk_ao/handlers/openai_agents/handler.py @@ -7,7 +7,7 @@ from agents.tracing import ResponseSpanData, get_current_span, get_trace_provider from galileo_core.schemas.logging.span import LlmMetrics, LlmSpan -from galileo_core.schemas.logging.span import Span as GalileoSpan +from galileo_core.schemas.logging.span import Span as SplunkAOSpan from splunk_ao import SplunkAOLogger, splunk_ao_context from splunk_ao.schema.handlers import Node from splunk_ao.utils import _get_timestamp @@ -492,7 +492,7 @@ def _extract_tool_output(self, item_dict: dict[str, Any], item_type: str) -> str return serialize_to_str(item_dict.get("output") or item_dict.get("results")) @staticmethod - def add_splunk_ao_custom_span(span: GalileoSpan) -> Span[SplunkAOCustomSpan]: + def add_splunk_ao_custom_span(span: SplunkAOSpan) -> Span[SplunkAOCustomSpan]: """Add a Splunk AO custom span to the trace.""" trace_provider = get_trace_provider() current_span = get_current_span() diff --git a/src/splunk_ao/otel.py b/src/splunk_ao/otel.py index 2386c16..6758b23 100644 --- a/src/splunk_ao/otel.py +++ b/src/splunk_ao/otel.py @@ -14,7 +14,7 @@ from opentelemetry.trace import Tracer from galileo_core.schemas.logging.span import AgentSpan, WorkflowSpan -from galileo_core.schemas.logging.span import Span as GalileoSpan +from galileo_core.schemas.logging.span import Span as SplunkAOSpan from splunk_ao.config import SplunkAOConfig from splunk_ao.converter import build_span_attributes from splunk_ao.decorator import ( @@ -298,7 +298,7 @@ def _apply_dataset_attributes( @contextmanager -def start_splunk_ao_span(splunk_ao_span: GalileoSpan) -> Generator[trace.Span, Any, None]: +def start_splunk_ao_span(splunk_ao_span: SplunkAOSpan) -> Generator[trace.Span, Any, None]: tracer_provider = _TRACE_PROVIDER_CONTEXT_VAR.get() if tracer_provider is None: tracer_provider = trace.get_tracer_provider() diff --git a/src/splunk_ao/utils/openai_agents.py b/src/splunk_ao/utils/openai_agents.py index a58c74c..3ae1369 100644 --- a/src/splunk_ao/utils/openai_agents.py +++ b/src/splunk_ao/utils/openai_agents.py @@ -13,7 +13,7 @@ ) from agents.tracing import ResponseSpanData -from galileo_core.schemas.logging.span import Span as GalileoSpan +from galileo_core.schemas.logging.span import Span as SplunkAOSpan from splunk_ao.schema.handlers import SPAN_TYPE from splunk_ao.utils.serialization import serialize_to_str @@ -21,7 +21,7 @@ class SplunkAOCustomSpan(CustomSpanData): - def __init__(self, span: GalileoSpan, data: dict[str, Any]): + def __init__(self, span: SplunkAOSpan, data: dict[str, Any]): self.span = span super().__init__(span.name, data) diff --git a/tests/test_openai_agents_utils.py b/tests/test_openai_agents_utils.py index 69b1bdc..773135b 100644 --- a/tests/test_openai_agents_utils.py +++ b/tests/test_openai_agents_utils.py @@ -296,10 +296,10 @@ def test_span_data_types(self, span_data: Any, expected_type: str) -> None: """Test mapping various span data types.""" assert _map_span_type(span_data) == expected_type - def test_galileo_custom_span(self) -> None: + def test_splunk_ao_custom_span(self) -> None: """Test mapping SplunkAOCustomSpan.""" - galileo_span = WorkflowSpan(name="Test", input="input", output="output", status_code=200) - assert _map_span_type(SplunkAOCustomSpan(galileo_span, {})) == "splunk_ao_custom" + splunk_ao_span = WorkflowSpan(name="Test", input="input", output="output", status_code=200) + assert _map_span_type(SplunkAOCustomSpan(splunk_ao_span, {})) == "splunk_ao_custom" @pytest.mark.parametrize( "type_attr,expected_type", [("function", "tool"), ("generation", "llm"), ("agent", "workflow")] diff --git a/tests/test_otel.py b/tests/test_otel.py index 8924482..654ce4f 100644 --- a/tests/test_otel.py +++ b/tests/test_otel.py @@ -314,7 +314,7 @@ def test_tool_span_with_output_no_tool_call_id(self): assert "gen_ai.tool.call.id" not in attrs -class TestStartGalileoSpan: +class TestStartSplunkAOSpan: """Test suite for start_splunk_ao_span context manager.""" @pytest.fixture(autouse=True) @@ -381,7 +381,7 @@ def test_start_splunk_ao_span_tool_span_with_none_output(self): assert "gen_ai.tool.call.id" not in calls @pytest.mark.parametrize( - "galileo_span", + "splunk_ao_span", [ LlmSpan(input="prompt", output="answer", model="gpt-4o"), ToolSpan(name="search", input="query", output="result"), @@ -390,8 +390,8 @@ def test_start_splunk_ao_span_tool_span_with_none_output(self): AgentSpan(name="agent", input="question", output="answer"), ], ) - def test_start_span_applies_canonical_builder_for_every_supported_type(self, galileo_span): - expected = build_span_attributes(galileo_span, session_id="session-id") + def test_start_span_applies_canonical_builder_for_every_supported_type(self, splunk_ao_span): + expected = build_span_attributes(splunk_ao_span, session_id="session-id") mock_otel_span = Mock() mock_tracer = Mock() mock_tracer.start_as_current_span.return_value.__enter__ = Mock(return_value=mock_otel_span) @@ -402,7 +402,7 @@ def test_start_span_applies_canonical_builder_for_every_supported_type(self, gal token = _session_id_context.set("session-id") try: - with start_splunk_ao_span(galileo_span): + with start_splunk_ao_span(splunk_ao_span): pass finally: _session_id_context.reset(token) @@ -411,7 +411,7 @@ def test_start_span_applies_canonical_builder_for_every_supported_type(self, gal assert {key: calls[key] for key in expected} == expected def test_start_span_applies_canonical_attributes_when_body_raises(self): - galileo_span = ToolSpan(name="search", input="query", output="result") + splunk_ao_span = ToolSpan(name="search", input="query", output="result") mock_otel_span = Mock() mock_tracer = Mock() mock_tracer.start_as_current_span.return_value.__enter__ = Mock(return_value=mock_otel_span) @@ -420,7 +420,7 @@ def test_start_span_applies_canonical_attributes_when_body_raises(self): mock_provider.get_tracer.return_value = mock_tracer _TRACE_PROVIDER_CONTEXT_VAR.set(mock_provider) - with pytest.raises(RuntimeError, match="failure"), start_splunk_ao_span(galileo_span): + with pytest.raises(RuntimeError, match="failure"), start_splunk_ao_span(splunk_ao_span): raise RuntimeError("failure") calls = {args[0]: args[1] for args, _ in mock_otel_span.set_attribute.call_args_list} @@ -430,7 +430,7 @@ def test_start_span_applies_canonical_attributes_when_body_raises(self): @patch("splunk_ao.otel.logger.warning") @patch("splunk_ao.otel.build_span_attributes", side_effect=ValueError("invalid partial span")) def test_start_span_finalization_does_not_mask_body_exception(self, _mock_build, mock_warning): - galileo_span = ToolSpan(name="search", input="query", output="result") + splunk_ao_span = ToolSpan(name="search", input="query", output="result") mock_tracer = Mock() mock_tracer.start_as_current_span.return_value.__enter__ = Mock(return_value=Mock()) mock_tracer.start_as_current_span.return_value.__exit__ = Mock(return_value=False) @@ -438,7 +438,7 @@ def test_start_span_finalization_does_not_mask_body_exception(self, _mock_build, mock_provider.get_tracer.return_value = mock_tracer _TRACE_PROVIDER_CONTEXT_VAR.set(mock_provider) - with pytest.raises(RuntimeError, match="user failure"), start_splunk_ao_span(galileo_span): + with pytest.raises(RuntimeError, match="user failure"), start_splunk_ao_span(splunk_ao_span): raise RuntimeError("user failure") mock_warning.assert_called_once_with("Failed to finalize Splunk AO span attributes", exc_info=True) @@ -446,7 +446,7 @@ def test_start_span_finalization_does_not_mask_body_exception(self, _mock_build, @patch("splunk_ao.otel.logger.warning") @patch("splunk_ao.otel.build_span_attributes", side_effect=ValueError("invalid span")) def test_start_span_finalization_failure_does_not_fail_successful_body(self, _mock_build, mock_warning): - galileo_span = ToolSpan(name="search", input="query", output="result") + splunk_ao_span = ToolSpan(name="search", input="query", output="result") mock_tracer = Mock() mock_tracer.start_as_current_span.return_value.__enter__ = Mock(return_value=Mock()) mock_tracer.start_as_current_span.return_value.__exit__ = Mock(return_value=False) @@ -454,7 +454,7 @@ def test_start_span_finalization_failure_does_not_fail_successful_body(self, _mo mock_provider.get_tracer.return_value = mock_tracer _TRACE_PROVIDER_CONTEXT_VAR.set(mock_provider) - with start_splunk_ao_span(galileo_span): + with start_splunk_ao_span(splunk_ao_span): pass mock_warning.assert_called_once_with("Failed to finalize Splunk AO span attributes", exc_info=True) @@ -462,7 +462,7 @@ def test_start_span_finalization_failure_does_not_fail_successful_body(self, _mo @patch("splunk_ao.otel.logger.warning") @patch("splunk_ao.otel.build_span_attributes", return_value={"gen_ai.operation.name": "execute_tool"}) def test_start_span_contains_attribute_write_failure(self, _mock_build, mock_warning): - galileo_span = ToolSpan(name="search", input="query", output="result") + splunk_ao_span = ToolSpan(name="search", input="query", output="result") mock_otel_span = Mock() mock_otel_span.set_attribute.side_effect = ValueError("attribute rejected") mock_tracer = Mock() @@ -472,19 +472,19 @@ def test_start_span_contains_attribute_write_failure(self, _mock_build, mock_war mock_provider.get_tracer.return_value = mock_tracer _TRACE_PROVIDER_CONTEXT_VAR.set(mock_provider) - with start_splunk_ao_span(galileo_span): + with start_splunk_ao_span(splunk_ao_span): pass mock_warning.assert_called_once_with("Failed to finalize Splunk AO span attributes", exc_info=True) @pytest.mark.parametrize( - "galileo_span", + "splunk_ao_span", [ WorkflowSpan(name="workflow", input="input", output="output"), AgentSpan(name="agent", input="input", output="output"), ], ) - def test_start_splunk_ao_span_marks_eligible_root_without_parent(self, galileo_span): + def test_start_splunk_ao_span_marks_eligible_root_without_parent(self, splunk_ao_span): """Eligible spans with no caller parent receive the standard root marker.""" mock_otel_span = Mock() mock_tracer = Mock() @@ -496,7 +496,7 @@ def test_start_splunk_ao_span_marks_eligible_root_without_parent(self, galileo_s with patch("splunk_ao.otel.trace") as mock_trace: mock_trace.get_current_span.return_value.get_span_context.return_value.is_valid = False - with start_splunk_ao_span(galileo_span): + with start_splunk_ao_span(splunk_ao_span): pass calls = { @@ -525,7 +525,7 @@ def test_start_splunk_ao_span_does_not_mark_span_with_parent(self): assert "gen_ai.conversation_root" not in calls @pytest.mark.parametrize( - "galileo_span", + "splunk_ao_span", [ ToolSpan(name="tool", input="input", output="output"), LlmSpan( @@ -537,7 +537,7 @@ def test_start_splunk_ao_span_does_not_mark_span_with_parent(self): RetrieverSpan(name="retriever", input="input", output=[]), ], ) - def test_start_splunk_ao_span_does_not_mark_ineligible_root(self, galileo_span): + def test_start_splunk_ao_span_does_not_mark_ineligible_root(self, splunk_ao_span): """LLM, tool, and retriever spans are never conversation roots.""" mock_otel_span = Mock() mock_tracer = Mock() @@ -549,7 +549,7 @@ def test_start_splunk_ao_span_does_not_mark_ineligible_root(self, galileo_span): with patch("splunk_ao.otel.trace") as mock_trace: mock_trace.get_current_span.return_value.get_span_context.return_value.is_valid = False - with start_splunk_ao_span(galileo_span): + with start_splunk_ao_span(splunk_ao_span): pass calls = {args[0]: args[1] for args, _ in mock_otel_span.set_attribute.call_args_list}