feat(langchain): implement duration_ns metric for langchain callback async handler (#150)

Author: andriisoldatenko

src/galileo/handlers/langchain/async_handler.py

@@ -116,7 +116,14 @@ async def _log_node_tree(self, node: Node) -> None:
 
         # Log the current node based on its type
         if node.node_type in ("agent", "chain"):
-            self._galileo_logger.add_workflow_span(input=input_, output=output, name=name, metadata=metadata, tags=tags)
+            self._galileo_logger.add_workflow_span(
+                input=input_,
+                output=output,
+                name=name,
+                duration_ns=node.span_params.get("duration_ns"),
+                metadata=metadata,
+                tags=tags,
+            )
             is_workflow_span = True
         elif node.node_type in ("llm", "chat"):
             self._galileo_logger.add_llm_span(
@@ -126,6 +133,7 @@ async def _log_node_tree(self, node: Node) -> None:
                 temperature=node.span_params.get("temperature"),
                 tools=node.span_params.get("tools"),
                 name=name,
+                duration_ns=node.span_params.get("duration_ns"),
                 metadata=metadata,
                 tags=tags,
                 num_input_tokens=node.span_params.get("num_input_tokens"),
@@ -135,10 +143,22 @@ async def _log_node_tree(self, node: Node) -> None:
             )
         elif node.node_type == "retriever":
             self._galileo_logger.add_retriever_span(
-                input=input_, output=output, name=name, metadata=metadata, tags=tags
+                input=input_,
+                output=output,
+                name=name,
+                duration_ns=node.span_params.get("duration_ns"),
+                metadata=metadata,
+                tags=tags,
             )
         elif node.node_type == "tool":
-            self._galileo_logger.add_tool_span(input=input_, output=output, name=name, metadata=metadata, tags=tags)
+            self._galileo_logger.add_tool_span(
+                input=input_,
+                output=output,
+                name=name,
+                duration_ns=node.span_params.get("duration_ns"),
+                metadata=metadata,
+                tags=tags,
+            )
         else:
             _logger.warning(f"Unknown node type: {node.node_type}")
 
@@ -187,6 +207,10 @@ async def _start_node(
 
         # Create new node
         node = Node(node_type=node_type, span_params=kwargs, run_id=run_id, parent_run_id=parent_run_id)
+
+        if "start_time" not in node.span_params:
+            node.span_params["start_time"] = time.perf_counter_ns()
+
         self._nodes[node_id] = node
 
         # Set as root node if needed
@@ -222,6 +246,8 @@ async def _end_node(self, run_id: UUID, **kwargs: Any) -> None:
             _logger.debug(f"No node exists for run_id {node_id}")
             return
 
+        node.span_params["duration_ns"] = time.perf_counter_ns() - node.span_params["start_time"]
+
         # Update node parameters
         node.span_params.update(**kwargs)
 
@@ -302,7 +328,7 @@ async def on_llm_start(
             model=model,
             temperature=temperature,
             metadata={k: str(v) for k, v in metadata.items()} if metadata else None,
-            start_time=time.perf_counter(),
+            start_time=time.perf_counter_ns(),
             time_to_first_token_ns=None,
         )
 
@@ -315,7 +341,7 @@ async def on_llm_new_token(self, token: str, *, run_id: UUID, **kwargs: Any) ->
         if "time_to_first_token_ns" not in node.span_params or node.span_params["time_to_first_token_ns"] is None:
             start_time = node.span_params.get("start_time")
             if start_time is not None:
-                node.span_params["time_to_first_token_ns"] = (time.perf_counter() - start_time) * 1e9
+                node.span_params["time_to_first_token_ns"] = time.perf_counter_ns() - start_time
 
     async def on_chat_model_start(
         self,

tests/test_langchain_async.py

@@ -66,6 +66,7 @@ async def test_start_node(self, callback: GalileoAsyncCallback):
         assert node.parent_run_id is None
         assert "name" in node.span_params
         assert node.span_params["name"] == "Parent Chain"
+        assert node.span_params["start_time"] > 0
         assert str(parent_id) in callback._nodes
 
         # Create a child node
@@ -117,6 +118,7 @@ async def test_on_chain_start_end(self, callback: GalileoAsyncCallback, galileo_
         assert str(run_id) in callback._nodes
         assert callback._nodes[str(run_id)].node_type == "chain"
         assert callback._nodes[str(run_id)].span_params["input"] == '{"query": "test question"}'
+        assert callback._nodes[str(run_id)].span_params["start_time"] > 0
 
         # End chain
         await callback.on_chain_end(outputs='{"result": "test answer"}', run_id=run_id)
@@ -223,6 +225,7 @@ async def test_on_llm_start_end(self, callback: GalileoAsyncCallback):
         assert callback._nodes[str(run_id)].span_params["num_input_tokens"] == 10
         assert callback._nodes[str(run_id)].span_params["num_output_tokens"] == 20
         assert callback._nodes[str(run_id)].span_params["total_tokens"] == 30
+        assert callback._nodes[str(run_id)].span_params["duration_ns"] > 0
 
     @mark.asyncio
     async def test_on_chat_model_start(self, callback: GalileoAsyncCallback):