feat: add add_agent_span method to GalileoLogger

pyproject.toml

@@ -18,7 +18,7 @@ langchain-core = { version = "^0.3.61", optional = true }
 openai = { version = "<1.84.0", optional = true }
 openai-agents = { version = "<0.0.13", optional = true }
 
-galileo-core = "~=3.38.0"
+galileo-core = "~=3.40.0"
 
 [tool.poetry.group.test.dependencies]
 pytest = "^8.3.5"
@@ -28,7 +28,8 @@ pytest-xdist = "^3.3.1"
 pytest-socket = "^0.7"
 pytest-asyncio = "^0.26.0"
 requests-mock = "^1.11.0"
-galileo-core = { extras = ["testing"], version = "~=3.38.0" }
+galileo-core = { extras = ["testing"], version = "~=3.40.0" }
+
 pytest-env = "^1.1.5"
 langchain-core = "^0.3.61"
 pytest-sugar = "^1.0.0"

src/galileo/logger.py

@@ -19,7 +19,9 @@
 from galileo.utils.metrics import populate_local_metrics
 from galileo.utils.nop_logger import nop_async, nop_sync
 from galileo.utils.serialization import serialize_to_str
+from galileo_core.schemas.logging.agent import AgentType
 from galileo_core.schemas.logging.span import (
+    AgentSpan,
     LlmSpan,
     LlmSpanAllowedInputType,
     LlmSpanAllowedOutputType,
@@ -519,6 +521,46 @@ def add_workflow_span(
             tags=tags,
         )
 
+    @nop_sync
+    def add_agent_span(
+        self,
+        input: str,
+        output: Optional[str] = None,
+        name: Optional[str] = None,
+        duration_ns: Optional[int] = None,
+        created_at: Optional[datetime] = None,
+        metadata: Optional[dict[str, str]] = None,
+        tags: Optional[list[str]] = None,
+        agent_type: Optional[AgentType] = None,
+    ) -> AgentSpan:
+        """
+        Add an agent type span to the current parent.
+
+        Parameters:
+        ----------
+            input: str: Input to the node.
+            output: Optional[str]: Output of the node. This can also be set on conclude().
+            name: Optional[str]: Name of the span.
+            duration_ns: Optional[int]: duration_ns of the node in nanoseconds.
+            created_at: Optional[datetime]: Timestamp of the span's creation.
+            metadata: Optional[Dict[str, str]]: Metadata associated with this span.
+            agent_type: Optional[AgentType]: Agent type of the span.
+
+        Returns:
+        -------
+            AgentSpan: The created span.
+        """
+        return super().add_agent_span(
+            input=input,
+            output=output,
+            name=name,
+            duration_ns=duration_ns,
+            created_at=created_at,
+            user_metadata=metadata,
+            tags=tags,
+            agent_type=agent_type,
+        )
+
     @nop_sync
     def conclude(
         self,

tasks.py

@@ -49,3 +49,16 @@ def type_check(ctx: Context) -> None:
 @task
 def docs_build(ctx: Context) -> None:
     ctx.run("poetry run mkdocs build --verbose", **COMMON_PARAMS)
+
+
+@task
+def poetry_lock(ctx: Context) -> None:
+    """
+    Update poetry.lock file.
+
+    Parameters
+    ----------
+    ctx : Context
+        Invoke context.
+    """
+    ctx.run("poetry lock", **COMMON_PARAMS)

tests/test_logger.py

@@ -10,7 +10,8 @@
 from galileo.logger import GalileoLogger
 from galileo.schema.metrics import LocalMetricConfig
 from galileo.schema.trace import TracesIngestRequest
-from galileo_core.schemas.logging.span import LlmSpan, RetrieverSpan, Span, ToolSpan, WorkflowSpan
+from galileo_core.schemas.logging.agent import AgentType
+from galileo_core.schemas.logging.span import AgentSpan, LlmSpan, RetrieverSpan, Span, ToolSpan, WorkflowSpan
 from galileo_core.schemas.logging.step import Metrics
 from galileo_core.schemas.logging.trace import Trace
 from galileo_core.schemas.shared.document import Document
@@ -203,6 +204,35 @@ def test_nested_span_trace_to_galileo(
     assert logger._parent_stack == deque()
 
 
+@patch("galileo.logger.LogStreams")
+@patch("galileo.logger.Projects")
+@patch("galileo.logger.GalileoCoreApiClient")
+def test_add_agent_span(mock_core_api_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock) -> None:
+    mock_core_api_instance = setup_mock_core_api_client(mock_core_api_client)
+    setup_mock_projects_client(mock_projects_client)
+    setup_mock_logstreams_client(mock_logstreams_client)
+
+    created_at = datetime.datetime.now()
+    metadata = {"key": "value"}
+    logger = GalileoLogger(project="my_project", log_stream="my_log_stream")
+    trace = logger.start_trace(
+        input="input", name="test-trace", duration_ns=1_000_000, created_at=created_at, metadata=metadata
+    )
+
+    logger.add_agent_span(input="prompt", name="test-agent-span", created_at=created_at, metadata=metadata)
+
+    logger.conclude(output="response", duration_ns=1_000_000, status_code=200)
+    logger.flush()
+
+    payload = mock_core_api_instance.ingest_traces_sync.call_args[0][0]
+    expected_payload = TracesIngestRequest(log_stream_id=None, experiment_id=None, traces=[trace])
+    assert payload == expected_payload
+    assert isinstance(payload.traces[0].spans[0], AgentSpan)
+    assert payload.traces[0].spans[0].agent_type == AgentType.default
+    assert logger.traces == list()
+    assert logger._parent_stack == deque()
+
+
 @patch("galileo.logger.LogStreams")
 @patch("galileo.logger.Projects")
 @patch("galileo.logger.GalileoCoreApiClient")