judgeval 0.0.54__py3-none-any.whl → 0.1.0__py3-none-any.whl

judgeval/integrations/langgraph.py

@@ -19,15 +19,9 @@ from langchain_core.outputs import LLMResult
 from langchain_core.messages.base import BaseMessage
 from langchain_core.documents import Document
 
-# --- Get context vars from tracer module ---
-# Assuming tracer.py defines these and they are accessible
-# If not, redefine them here or adjust import
-
-# from judgeval.common.tracer import current_span_var
 # TODO: Figure out how to handle context variables. Current solution is to keep track of current span id in Tracer class
 
 
-# --- NEW __init__ ---
 class JudgevalCallbackHandler(BaseCallbackHandler):
     """
     LangChain Callback Handler using run_id/parent_run_id for hierarchy.
@@ -40,19 +34,11 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
     lc_serializable = False
     lc_kwargs: dict = {}
 
-    # --- NEW __init__ ---
     def __init__(self, tracer: Tracer):
         self.tracer = tracer
-        # Initialize tracking/logging variables (preserved across resets)
         self.executed_nodes: List[str] = []
-        self.executed_tools: List[str] = []
-        self.executed_node_tools: List[str] = []
-        self.traces: List[Dict[str, Any]] = []
-        # Initialize execution state (reset between runs)
         self._reset_state()
 
-    # --- END NEW __init__ ---
-
     def _reset_state(self):
         """Reset only the critical execution state for reuse across multiple executions"""
         # Reset core execution state that must be cleared between runs
@@ -61,28 +47,16 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         self._span_id_to_start_time: Dict[str, float] = {}
         self._span_id_to_depth: Dict[str, int] = {}
         self._root_run_id: Optional[UUID] = None
-        self._trace_saved: bool = False  # Flag to prevent actions after trace is saved
+        self._trace_saved: bool = False
         self.span_id_to_token: Dict[str, Any] = {}
         self.trace_id_to_token: Dict[str, Any] = {}
 
         # Add timestamp to track when we last reset
         self._last_reset_time: float = time.time()
 
-        # Preserve tracking/logging variables across executions:
-        # - self.executed_nodes: List[str] = [] # Keep as running log
-        # - self.executed_tools: List[str] = [] # Keep as running log
-        # - self.executed_node_tools: List[str] = [] # Keep as running log
-        # - self.traces: List[Dict[str, Any]] = [] # Keep for collecting multiple traces
-
         # Also reset tracking/logging variables
-        self.executed_nodes: List[
-            str
-        ] = []  # These last four members are only appended to and never accessed; can probably be removed but still might be useful for future reference?
-        self.executed_tools: List[str] = []
-        self.executed_node_tools: List[str] = []
-        self.traces: List[Dict[str, Any]] = []
-
-    # --- END NEW __init__ ---
+        self.executed_nodes: List[str] = []
+
     def reset(self):
         """Public method to manually reset handler execution state for reuse"""
         self._reset_state()
@@ -91,7 +65,6 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         """Public method to reset ALL handler state including tracking/logging data"""
         self._reset_state()
 
-    # --- MODIFIED _ensure_trace_client ---
     def _ensure_trace_client(
         self, run_id: UUID, parent_run_id: Optional[UUID], event_name: str
     ) -> Optional[TraceClient]:
@@ -127,19 +100,14 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             token = self.tracer.set_current_trace(self._trace_client)
             if token:
                 self.trace_id_to_token[trace_id] = token
+
             if self._trace_client:
-                self._root_run_id = (
-                    run_id  # Assign the first run_id encountered as the tentative root
-                )
-                self._trace_saved = False  # Ensure flag is reset
-                # Set active client on Tracer (important for potential fallbacks)
+                self._root_run_id = run_id
+                self._trace_saved = False
                 self.tracer._active_trace_client = self._trace_client
 
-                # NEW: Initial save for live tracking (follows the new practice)
                 try:
-                    trace_id_saved, server_response = self._trace_client.save(
-                        final_save=False,  # Initial save for live tracking
-                    )
+                    self._trace_client.save(final_save=False)
                 except Exception as e:
                     import warnings
 
@@ -207,18 +175,13 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             # Set both fields on the span
             new_span.inputs = clean_inputs
             new_span.additional_metadata = metadata
-            new_span.increment_update_id()  # Thread-safe increment for span modification
         else:
             new_span.inputs = {}
             new_span.additional_metadata = {}
 
         trace_client.add_span(new_span)
 
-        # Queue span with initial state (input phase) through background service
-        if trace_client.background_span_service:
-            trace_client.background_span_service.queue_span(
-                new_span, span_state="input"
-            )
+        trace_client.otel_span_processor.queue_span_update(new_span, span_state="input")
 
         token = self.tracer.set_current_span(span_id)
         if token:
@@ -247,12 +210,10 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             trace_span = trace_client.span_id_to_span.get(span_id)
             if trace_span:
                 trace_span.duration = duration
-                trace_span.increment_update_id()  # Thread-safe increment for span modification
 
                 # Handle outputs and error
                 if error:
                     trace_span.output = error
-                    trace_span.increment_update_id()  # Thread-safe increment for span modification
                 elif outputs:
                     # Separate metadata from outputs
                     metadata = {}
@@ -272,7 +233,6 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
 
                     # Set both fields on the span
                     trace_span.output = clean_outputs
-                    trace_span.increment_update_id()  # Thread-safe increment for span modification
                     if metadata:
                         # Merge with existing metadata
                         existing_metadata = trace_span.additional_metadata or {}
@@ -280,14 +240,11 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                             **existing_metadata,
                             **metadata,
                         }
-                        trace_span.increment_update_id()  # Thread-safe increment for span modification
 
-                # Queue span with completed state through background service
-                if trace_client.background_span_service:
-                    span_state = "error" if error else "completed"
-                    trace_client.background_span_service.queue_span(
-                        trace_span, span_state=span_state
-                    )
+                span_state = "error" if error else "completed"
+                trace_client.otel_span_processor.queue_span_update(
+                    trace_span, span_state=span_state
+                )
 
             # Clean up dictionaries for this specific span
             if span_id in self._span_id_to_start_time:
@@ -298,15 +255,10 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         # Check if this is the root run ending
         if run_id == self._root_run_id:
             try:
-                # Reset root run id after attempt
                 self._root_run_id = None
-                # Reset input storage for this handler instance
-
                 if (
                     self._trace_client and not self._trace_saved
                 ):  # Check if not already saved
-                    # Flush background spans before saving the final trace
-
                     complete_trace_data = {
                         "trace_id": self._trace_client.trace_id,
                         "name": self._trace_client.name,
@@ -321,6 +273,9 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                         "parent_trace_id": self._trace_client.parent_trace_id,
                         "parent_name": self._trace_client.parent_name,
                     }
+
+                    self.tracer.flush_background_spans()
+
                     trace_id, trace_data = self._trace_client.save(
                         final_save=True,  # Final save with usage counter updates
                     )
@@ -331,14 +286,12 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                     self.tracer.traces.append(complete_trace_data)
                     self._trace_saved = True  # Set flag only after successful save
             finally:
-                # --- NEW: Consolidated Cleanup Logic ---
                 # This block executes regardless of save success/failure
                 # Reset root run id
                 self._root_run_id = None
                 # Reset input storage for this handler instance
                 if self.tracer._active_trace_client == self._trace_client:
                     self.tracer._active_trace_client = None
-                # --- End Cleanup Logic ---
 
     # --- Callback Methods ---
     # Each method now ensures the trace client exists before proceeding
@@ -361,10 +314,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         )
 
         name = f"RETRIEVER_{(serialized_name).upper()}"
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, name
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, name)
         if not trace_client:
             return
 
@@ -392,17 +342,17 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> Any:
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, "RetrieverEnd"
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, "RetrieverEnd")
         if not trace_client:
             return
         doc_summary = [
             {
                 "index": i,
-                "page_content": doc.page_content[:100] + "..."
-                if len(doc.page_content) > 100
-                else doc.page_content,
+                "page_content": (
+                    doc.page_content[:100] + "..."
+                    if len(doc.page_content) > 100
+                    else doc.page_content
+                ),
                 "metadata": doc.metadata,
             }
             for i, doc in enumerate(documents)
@@ -431,7 +381,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
 
         # --- Determine Name and Span Type ---
         span_type: SpanType = "chain"
-        name = serialized_name if serialized_name else "Unknown Chain"  # Default name
+        name = serialized_name if serialized_name else "Unknown Chain"
         node_name = metadata.get("langgraph_node") if metadata else None
         is_langgraph_root_kwarg = (
             kwargs.get("name") == "LangGraph"
@@ -449,25 +399,17 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             name = "LangGraph"  # Explicit root detected
         # Add handling for other potential LangChain internal chains if needed, e.g., "RunnableSequence"
 
-        # --- Ensure Trace Client ---
-        # Pass parent_run_id to _ensure_trace_client
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, name
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, name)
         if not trace_client:
             return
 
-        # --- Update Trace Name if Root ---
-        # If this is the root event (parent_run_id is None) and the trace client was just created,
-        # ensure the trace name reflects the graph's name ('LangGraph' usually).
         if (
             is_potential_root_event
             and run_id == self._root_run_id
             and trace_client.name != name
         ):
-            trace_client.name = name  # Update trace name to the determined root name
+            trace_client.name = name
 
-        # --- Start Span Tracking ---
         combined_inputs = {
             "inputs": inputs,
             "tags": tags,
@@ -493,28 +435,20 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         tags: Optional[List[str]] = None,
         **kwargs: Any,
     ) -> Any:
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, "ChainEnd"
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, "ChainEnd")
         if not trace_client:
             return
 
         span_id = self._run_id_to_span_id.get(run_id)
-        # If it's the root run ending, _end_span_tracking will handle cleanup/save
         if not span_id and run_id != self._root_run_id:
-            return  # Don't call end tracking if it's not the root and span wasn't tracked
+            return
 
-        # Prepare outputs for end tracking (moved down)
         combined_outputs = {"outputs": outputs, "tags": tags, "kwargs": kwargs}
 
-        # Call end_span_tracking with potentially determined span_type
         self._end_span_tracking(trace_client, run_id, outputs=combined_outputs)
 
-        # --- Root node cleanup (Existing logic - slightly modified save call) ---
         if run_id == self._root_run_id:
             if trace_client and not self._trace_saved:
-                # Store complete trace data instead of server response
                 complete_trace_data = {
                     "trace_id": trace_client.trace_id,
                     "name": trace_client.name,
@@ -529,19 +463,19 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                     "parent_trace_id": trace_client.parent_trace_id,
                     "parent_name": trace_client.parent_name,
                 }
-                trace_id_saved, trace_data = trace_client.save(
+
+                self.tracer.flush_background_spans()
+
+                trace_client.save(
                     final_save=True,
                 )
 
                 self.tracer.traces.append(complete_trace_data)
                 self._trace_saved = True
-                # Reset tracer's active client *after* successful save
                 if self.tracer._active_trace_client == trace_client:
                     self.tracer._active_trace_client = None
 
-            # Reset root run id after attempt
             self._root_run_id = None
-            # Reset input storage for this handler instance
 
     def on_chain_error(
         self,
@@ -551,16 +485,12 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> Any:
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, "ChainError"
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, "ChainError")
         if not trace_client:
             return
 
         span_id = self._run_id_to_span_id.get(run_id)
 
-        # Let _end_span_tracking handle potential root run cleanup
         if not span_id and run_id != self._root_run_id:
             return
 
@@ -584,10 +514,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             else "Unknown Tool (Serialized=None)"
         )
 
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, name
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, name)
         if not trace_client:
             return
 
@@ -608,23 +535,6 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             inputs=combined_inputs,
         )
 
-        # --- Track executed tools (remains the same) ---
-        if name not in self.executed_tools:
-            self.executed_tools.append(
-                name
-            )  # Leaving this in for now but can probably be removed
-        parent_node_name = None
-        if parent_run_id and parent_run_id in self._run_id_to_span_id:
-            parent_span_id = self._run_id_to_span_id[parent_run_id]
-            parent_node_name = trace_client.span_id_to_span[parent_span_id].function
-
-        node_tool = f"{parent_node_name}:{name}" if parent_node_name else name
-        if node_tool not in self.executed_node_tools:
-            self.executed_node_tools.append(
-                node_tool
-            )  # Leaving this in for now but can probably be removed
-        # --- End Track executed tools ---
-
     def on_tool_end(
         self,
         output: Any,
@@ -633,10 +543,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> Any:
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, "ToolEnd"
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, "ToolEnd")
         if not trace_client:
             return
         outputs = {"output": output, "kwargs": kwargs}
@@ -650,10 +557,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> Any:
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, "ToolError"
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, "ToolError")
         if not trace_client:
             return
         self._end_span_tracking(trace_client, run_id, error=error)
@@ -674,9 +578,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
     ) -> Any:
         llm_name = name or serialized.get("name", "LLM Call")
 
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, llm_name
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, llm_name)
         if not trace_client:
             return
         inputs = {
@@ -704,15 +606,11 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> Any:
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, "LLMEnd"
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, "LLMEnd")
         if not trace_client:
             return
         outputs = {"response": response, "kwargs": kwargs}
 
-        # --- Token Usage Extraction and Cost Calculation ---
         prompt_tokens = None
         completion_tokens = None
         total_tokens = None
@@ -761,9 +659,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                     if prompt_tokens is not None and completion_tokens is not None:
                         total_tokens = prompt_tokens + completion_tokens
 
-            # --- Create TraceUsage object and set on span ---
             if prompt_tokens is not None or completion_tokens is not None:
-                # Calculate costs if model name is available
                 prompt_cost = None
                 completion_cost = None
                 total_cost_usd = None
@@ -792,7 +688,6 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                             f"Failed to calculate token costs for model {model_name}: {e}"
                         )
 
-                # Create TraceUsage object
                 usage = TraceUsage(
                     prompt_tokens=prompt_tokens,
                     completion_tokens=completion_tokens,
@@ -808,15 +703,12 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                     model_name=model_name,
                 )
 
-                # Set usage on the actual span (not in outputs)
                 span_id = self._run_id_to_span_id.get(run_id)
                 if span_id and span_id in trace_client.span_id_to_span:
                     trace_span = trace_client.span_id_to_span[span_id]
                     trace_span.usage = usage
-                    trace_span.increment_update_id()  # Thread-safe increment for span modification
 
         self._end_span_tracking(trace_client, run_id, outputs=outputs)
-        # --- End Token Usage ---
 
     def on_llm_error(
         self,
@@ -826,10 +718,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> Any:
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, "LLMError"
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, "LLMError")
         if not trace_client:
             return
         self._end_span_tracking(trace_client, run_id, error=error)
@@ -848,9 +737,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         name: Optional[str] = None,
         **kwargs: Any,
     ) -> Any:
-        # Reuse on_llm_start logic, adding message formatting if needed
         chat_model_name = name or serialized.get("name", "ChatModel Call")
-        # Add OPENAI_API_CALL suffix if model is OpenAI and not present
         is_openai = (
             any(
                 key.startswith("openai") for key in serialized.get("secrets", {}).keys()
@@ -872,7 +759,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             )
             or "together" in chat_model_name.lower()
         )
-        # Add more checks for other providers like Google if needed
+
         is_google = (
             any(
                 key.startswith("google") for key in serialized.get("secrets", {}).keys()
@@ -891,9 +778,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         elif is_google and "GOOGLE_API_CALL" not in chat_model_name:
             chat_model_name = f"{chat_model_name} GOOGLE_API_CALL"
 
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, chat_model_name
-        )  # Corrected call with parent_run_id
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, chat_model_name)
         if not trace_client:
             return
         inputs = {
@@ -911,7 +796,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             chat_model_name,
             span_type="llm",
             inputs=inputs,
-        )  # Use 'llm' span_type for consistency
+        )
 
     def on_agent_action(
         self,
@@ -923,10 +808,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
     ) -> Any:
         action_tool = action.tool
         name = f"AGENT_ACTION_{(action_tool).upper()}"
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, name
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, name)
         if not trace_client:
             return
 
@@ -948,10 +830,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         parent_run_id: Optional[UUID] = None,
         **kwargs: Any,
     ) -> Any:
-        # Pass parent_run_id
-        trace_client = self._ensure_trace_client(
-            run_id, parent_run_id, "AgentFinish"
-        )  # Corrected call
+        trace_client = self._ensure_trace_client(run_id, parent_run_id, "AgentFinish")
         if not trace_client:
             return