feat: add support for serializing WorkflowSpan input and output

src/galileo/otel.py

@@ -15,7 +15,7 @@
 from galileo.config import GalileoPythonConfig
 from galileo.decorator import _experiment_id_context, _log_stream_context, _project_context, _session_id_context
 from galileo.utils.retrievers import document_adapter
-from galileo_core.schemas.logging.span import RetrieverSpan
+from galileo_core.schemas.logging.span import RetrieverSpan, WorkflowSpan
 from galileo_core.schemas.logging.span import Span as GalileoSpan
 
 logger = logging.getLogger(__name__)
@@ -325,6 +325,48 @@ def _set_retriever_span_attributes(span: trace.Span, galileo_span: RetrieverSpan
     )
 
 
+def _set_workflow_span_attributes(span: trace.Span, galileo_span: WorkflowSpan) -> None:
+    """Set OpenTelemetry attributes for WorkflowSpan."""
+    # Handle input - Union[str, Sequence[Message]]
+    if isinstance(galileo_span.input, str):
+        input_messages = [{"role": "user", "content": galileo_span.input}]
+    else:
+        # Sequence[Message] - serialize each message
+        input_messages = []
+        for msg in list(galileo_span.input):
+            if hasattr(msg, "model_dump"):
+                input_messages.append(msg.model_dump(exclude_none=True))
+            else:
+                input_messages.append(msg)
+    span.set_attribute("gen_ai.input.messages", json.dumps(input_messages))
+
+    # Handle output - Union[str, Message, Sequence[Document], None]
+    if galileo_span.output is None:
+        return
+
+    output_value = galileo_span.output
+    # Type annotation to handle flexible content types (string or dict)
+    # Content can be: str (simple output), dict (documents), or dict (Message model_dump)
+    output_messages: list[dict[str, Any]] = []
+
+    if isinstance(output_value, str):
+        output_messages = [{"role": "assistant", "content": output_value}]
+    elif hasattr(output_value, "model_dump"):
+        # Single Message
+        output_messages = [output_value.model_dump(exclude_none=True)]
+    else:
+        # Sequence[Document] - wrap in assistant message
+        # Use document_adapter for consistency with _set_retriever_span_attributes
+        output_messages = [
+            {
+                "role": "assistant",
+                "content": {"documents": document_adapter.dump_python(list(output_value), mode="json")},
+            }
+        ]
+
+    span.set_attribute("gen_ai.output.messages", json.dumps(output_messages))
+
+
 @contextmanager
 def start_galileo_span(galileo_span: GalileoSpan) -> Generator[trace.Span, Any, None]:
     tracer_provider = _TRACE_PROVIDER_CONTEXT_VAR.get()
@@ -337,3 +379,5 @@ def start_galileo_span(galileo_span: GalileoSpan) -> Generator[trace.Span, Any,
         span.set_attribute("gen_ai.system", "galileo-otel")
         if isinstance(galileo_span, RetrieverSpan):
             _set_retriever_span_attributes(span, galileo_span)
+        elif isinstance(galileo_span, WorkflowSpan):
+            _set_workflow_span_attributes(span, galileo_span)

tests/test_otel.py

@@ -9,6 +9,12 @@
 from galileo.decorator import _experiment_id_context, _log_stream_context, _project_context, _session_id_context
 from galileo.otel import INSTALL_ERR_MSG, OTEL_AVAILABLE, GalileoOTLPExporter, GalileoSpanProcessor
 
+if OTEL_AVAILABLE:
+    from galileo.otel import _set_workflow_span_attributes, start_galileo_span
+    from galileo_core.schemas.logging.llm import Message, MessageRole
+    from galileo_core.schemas.logging.span import WorkflowSpan
+    from galileo_core.schemas.shared.document import Document
+
 
 class TestGalileoOTLPExporter:
     """Test suite for GalileoOTLPExporter class."""
@@ -453,3 +459,125 @@ def test_exporter_export_merges_resource_attributes(self, mock_resource_class, m
         mock_resource_class.assert_not_called()
         mock_span2.resource.merge.assert_not_called()
         mock_parent_export.assert_called_once_with([mock_span2])
+
+
+class TestWorkflowSpanAttributes:
+    """Test suite for WorkflowSpan OpenTelemetry attribute mapping."""
+
+    @pytest.fixture
+    def mock_dependencies(self):
+        """Set up mocks for testing workflow span attributes."""
+        with patch("galileo.otel.trace") as mock_trace_module, patch("galileo.otel.json") as mock_json_module:
+            mock_span = Mock()
+            mock_json_module.dumps.return_value = '"test"'
+            yield {"span": mock_span, "trace": mock_trace_module, "json": mock_json_module}
+
+    @pytest.mark.skipif(not OTEL_AVAILABLE, reason="OpenTelemetry not available")
+    def test_workflow_span_with_string_input_output(self, mock_dependencies):
+        """Test WorkflowSpan with string input and output."""
+        # Given: a WorkflowSpan with string input and output
+        workflow_span = WorkflowSpan(name="test-workflow", input="input text", output="output text", status_code=200)
+        mock_span = mock_dependencies["span"]
+        mock_json = mock_dependencies["json"]
+
+        # When: setting workflow span attributes
+        _set_workflow_span_attributes(mock_span, workflow_span)
+
+        # Then: input and output should be wrapped in message format
+        assert mock_span.set_attribute.call_count == 2
+
+        # Check first call (input)
+        input_call = mock_span.set_attribute.call_args_list[0]
+        assert input_call[0][0] == "gen_ai.input.messages"
+        mock_json.dumps.assert_any_call([{"role": "user", "content": "input text"}])
+
+        # Check second call (output)
+        output_call = mock_span.set_attribute.call_args_list[1]
+        assert output_call[0][0] == "gen_ai.output.messages"
+        mock_json.dumps.assert_any_call([{"role": "assistant", "content": "output text"}])
+
+    @pytest.mark.skipif(not OTEL_AVAILABLE, reason="OpenTelemetry not available")
+    def test_workflow_span_with_message_input_output(self, mock_dependencies):
+        """Test WorkflowSpan with Message input and output."""
+        # Given: a WorkflowSpan with Message input and output
+        input_msg = Message(role=MessageRole.user, content="user question")
+        workflow_span = WorkflowSpan(name="test-workflow", input=[input_msg], output=input_msg, status_code=200)
+        mock_span = mock_dependencies["span"]
+        mock_dependencies["json"]
+
+        # When: setting workflow span attributes
+        _set_workflow_span_attributes(mock_span, workflow_span)
+
+        # Then: input and output should serialize Message objects
+        assert mock_span.set_attribute.call_count == 2
+        input_call = mock_span.set_attribute.call_args_list[0]
+        assert input_call[0][0] == "gen_ai.input.messages"
+        output_call = mock_span.set_attribute.call_args_list[1]
+        assert output_call[0][0] == "gen_ai.output.messages"
+
+    @pytest.mark.skipif(not OTEL_AVAILABLE, reason="OpenTelemetry not available")
+    def test_workflow_span_with_document_sequence_output(self, mock_dependencies):
+        """Test WorkflowSpan with Document sequence output."""
+        # Given: a WorkflowSpan with string input and Document sequence output
+        documents = [
+            Document(content="doc1 content", metadata={"source": "db"}),
+            Document(content="doc2 content", metadata={"source": "api"}),
+        ]
+        # Note: Using model_construct to bypass a galileo_core validator bug where
+        # Message._allow_null_content_with_tool_calling assumes dict input during
+        # Union type validation, but receives a list when validating Sequence[Document]
+        workflow_span = WorkflowSpan.model_construct(
+            name="test-workflow", input="query", output=documents, status_code=200
+        )
+        mock_span = mock_dependencies["span"]
+        mock_dependencies["json"]
+
+        # When: setting workflow span attributes
+        _set_workflow_span_attributes(mock_span, workflow_span)
+
+        # Then: output should be wrapped in assistant message with documents
+        assert mock_span.set_attribute.call_count == 2
+        output_call = mock_span.set_attribute.call_args_list[1]
+        assert output_call[0][0] == "gen_ai.output.messages"
+
+    @pytest.mark.skipif(not OTEL_AVAILABLE, reason="OpenTelemetry not available")
+    def test_workflow_span_with_none_output(self, mock_dependencies):
+        """Test WorkflowSpan with None output (should not set output attribute)."""
+        # Given: a WorkflowSpan with None output
+        workflow_span = WorkflowSpan(name="test-workflow", input="input text", output=None, status_code=200)
+        mock_span = mock_dependencies["span"]
+
+        # When: setting workflow span attributes
+        _set_workflow_span_attributes(mock_span, workflow_span)
+
+        # Then: only input attribute should be set, not output
+        assert mock_span.set_attribute.call_count == 1
+        input_call = mock_span.set_attribute.call_args_list[0]
+        assert input_call[0][0] == "gen_ai.input.messages"
+
+    @pytest.mark.skipif(not OTEL_AVAILABLE, reason="OpenTelemetry not available")
+    def test_workflow_span_in_start_galileo_span(self, mock_dependencies):
+        """Test that WorkflowSpan is handled in start_galileo_span context manager."""
+        # Given: a WorkflowSpan
+        workflow_span = WorkflowSpan(name="test-workflow", input="input", output="output", status_code=200)
+        mock_span = mock_dependencies["span"]
+        mock_dependencies["json"]
+
+        # Setup the mock tracer
+        mock_tracer = Mock()
+        mock_tracer.start_as_current_span.return_value.__enter__ = Mock(return_value=mock_span)
+        mock_tracer.start_as_current_span.return_value.__exit__ = Mock(return_value=None)
+
+        mock_trace_provider = Mock()
+        mock_trace_provider.get_tracer.return_value = mock_tracer
+
+        # Patch get_tracer_provider to return our mock
+        with patch("galileo.otel.trace.get_tracer_provider", return_value=mock_trace_provider):
+            # When: using start_galileo_span with WorkflowSpan
+            with start_galileo_span(workflow_span):
+                pass
+
+        # Then: WorkflowSpan attributes should be set
+        assert mock_span.set_attribute.call_count >= 2  # gen_ai.system + at least input
+        system_call = [call for call in mock_span.set_attribute.call_args_list if call[0][0] == "gen_ai.system"]
+        assert len(system_call) == 1