judgeval 0.7.1__py3-none-any.whl → 0.9.0__py3-none-any.whl

judgeval/integrations/langgraph/__init__.py

@@ -0,0 +1,788 @@
+from __future__ import annotations
+
+import time
+import uuid
+from typing import Any, Dict, List, Optional, Sequence, Set, Type
+from uuid import UUID
+
+try:
+    from langchain_core.callbacks import BaseCallbackHandler
+    from langchain_core.agents import AgentAction, AgentFinish
+    from langchain_core.outputs import LLMResult, ChatGeneration
+    from langchain_core.messages import (
+        AIMessage,
+        BaseMessage,
+        ChatMessage,
+        FunctionMessage,
+        HumanMessage,
+        SystemMessage,
+        ToolMessage,
+    )
+    from langchain_core.documents import Document
+except ImportError as e:
+    raise ImportError(
+        "Judgeval's langgraph integration requires langchain to be installed. Please install it with `pip install judgeval[langchain]`"
+    ) from e
+
+from judgeval.tracer import Tracer
+from judgeval.tracer.keys import AttributeKeys
+from judgeval.tracer.managers import sync_span_context
+from judgeval.utils.serialize import safe_serialize
+from judgeval.logger import judgeval_logger
+from opentelemetry.trace import Status, StatusCode, Span
+
+# Control flow exception types that should not be treated as errors
+CONTROL_FLOW_EXCEPTION_TYPES: Set[Type[BaseException]] = set()
+
+try:
+    from langgraph.errors import GraphBubbleUp
+
+    CONTROL_FLOW_EXCEPTION_TYPES.add(GraphBubbleUp)
+except ImportError:
+    pass
+
+LANGSMITH_TAG_HIDDEN: str = "langsmith:hidden"
+
+
+class JudgevalCallbackHandler(BaseCallbackHandler):
+    """
+    LangGraph/LangChain Callback Handler that creates OpenTelemetry spans
+    using the Judgeval tracer framework.
+
+    This handler tracks the execution of chains, tools, LLMs, and other components
+    in a LangGraph/LangChain application, creating proper span hierarchies for monitoring.
+    """
+
+    # Prevent LangChain serialization issues
+    lc_serializable = False
+    lc_kwargs: dict = {}
+
+    def __init__(self, tracer: Optional[Tracer] = None):
+        """
+        Initialize the callback handler.
+
+        Args:
+            tracer: Optional Tracer instance. If not provided, will try to use an active tracer.
+        """
+        self.tracer = tracer
+        if self.tracer is None:
+            # Try to get an active tracer
+            if Tracer._active_tracers:
+                self.tracer = next(iter(Tracer._active_tracers))
+            else:
+                judgeval_logger.warning(
+                    "No tracer provided and no active tracers found. "
+                    "Callback handler will not create spans."
+                )
+                return
+
+        # Track spans by run_id for proper hierarchy
+        self.spans: Dict[UUID, Span] = {}
+        self.span_start_times: Dict[UUID, float] = {}
+        self.run_id_to_span_id: Dict[UUID, str] = {}
+        self.span_id_to_depth: Dict[str, int] = {}
+        self.root_run_id: Optional[UUID] = None
+
+        # Track execution for debugging
+        self.executed_nodes: List[str] = []
+        self.executed_tools: List[str] = []
+        self.executed_node_tools: List[Dict[str, Any]] = []
+
+    def reset(self):
+        """Reset handler state for reuse across multiple executions."""
+        self.spans.clear()
+        self.span_start_times.clear()
+        self.executed_nodes.clear()
+        self.executed_tools.clear()
+        self.executed_node_tools.clear()
+
+    def _get_run_name(self, serialized: Optional[Dict[str, Any]], **kwargs: Any) -> str:
+        """Extract the name of the operation from serialized data or kwargs."""
+        if "name" in kwargs and kwargs["name"] is not None:
+            return str(kwargs["name"])
+
+        if serialized is None:
+            return "<unknown>"
+
+        try:
+            return str(serialized["name"])
+        except (KeyError, TypeError):
+            pass
+
+        try:
+            return str(serialized["id"][-1])
+        except (KeyError, TypeError):
+            pass
+
+        return "<unknown>"
+
+    def _convert_message_to_dict(self, message: BaseMessage) -> Dict[str, Any]:
+        """Convert a LangChain message to a dictionary for storage."""
+        if isinstance(message, HumanMessage):
+            message_dict = {"role": "user", "content": message.content}
+        elif isinstance(message, AIMessage):
+            message_dict = {"role": "assistant", "content": message.content}
+        elif isinstance(message, SystemMessage):
+            message_dict = {"role": "system", "content": message.content}
+        elif isinstance(message, ToolMessage):
+            message_dict = {
+                "role": "tool",
+                "content": message.content,
+                "tool_call_id": message.tool_call_id,
+            }
+        elif isinstance(message, FunctionMessage):
+            message_dict = {"role": "function", "content": message.content}
+        elif isinstance(message, ChatMessage):
+            message_dict = {"role": message.role, "content": message.content}
+        else:
+            message_dict = {"role": "unknown", "content": str(message.content)}
+
+        if hasattr(message, "additional_kwargs") and message.additional_kwargs:
+            message_dict["additional_kwargs"] = str(message.additional_kwargs)
+
+        return message_dict
+
+    def _create_message_dicts(
+        self, messages: List[BaseMessage]
+    ) -> List[Dict[str, Any]]:
+        """Convert a list of LangChain messages to dictionaries."""
+        return [self._convert_message_to_dict(m) for m in messages]
+
+    def _join_tags_and_metadata(
+        self,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> Optional[Dict[str, Any]]:
+        """Join tags and metadata into a single dictionary."""
+        final_dict = {}
+        if tags is not None and len(tags) > 0:
+            final_dict["tags"] = tags
+        if metadata is not None:
+            final_dict.update(metadata)
+        return final_dict if final_dict else None
+
+    def _start_span(
+        self,
+        run_id: UUID,
+        parent_run_id: Optional[UUID],
+        name: str,
+        span_type: str,
+        inputs: Any = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **extra_attributes: Any,
+    ) -> None:
+        """Start a new span for the given run."""
+        if not self.tracer:
+            return
+
+        # Skip internal spans
+        if name.startswith("__") and name.endswith("__"):
+            return
+
+        try:
+            # Determine if this is a root span
+            is_root = parent_run_id is None
+            if is_root:
+                self.root_run_id = run_id
+
+            # Calculate depth for proper hierarchy
+            current_depth = 0
+            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]
+                current_depth = self.span_id_to_depth.get(parent_span_id, 0) + 1
+
+            # Create span attributes
+            attributes = {
+                AttributeKeys.JUDGMENT_SPAN_KIND.value: span_type,
+            }
+
+            # Add metadata and tags
+            combined_metadata = self._join_tags_and_metadata(tags, metadata)
+            if combined_metadata:
+                metadata_str = safe_serialize(combined_metadata)
+                attributes["metadata"] = metadata_str
+
+            # Add extra attributes
+            for key, value in extra_attributes.items():
+                if value is not None:
+                    attributes[str(key)] = str(value)
+
+            # Create span using the tracer's context manager for proper hierarchy
+            with sync_span_context(self.tracer, name, attributes) as span:
+                # Set input data if provided
+                if inputs is not None:
+                    span.set_attribute(
+                        AttributeKeys.JUDGMENT_INPUT.value, safe_serialize(inputs)
+                    )
+
+                # Store span information for tracking
+                span_id = (
+                    str(span.get_span_context().span_id)
+                    if span.get_span_context()
+                    else str(uuid.uuid4())
+                )
+                self.spans[run_id] = span
+                self.span_start_times[run_id] = time.time()
+                self.run_id_to_span_id[run_id] = span_id
+                self.span_id_to_depth[span_id] = current_depth
+
+        except Exception as e:
+            judgeval_logger.exception(f"Error starting span for {name}: {e}")
+
+    def _end_span(
+        self,
+        run_id: UUID,
+        outputs: Any = None,
+        error: Optional[BaseException] = None,
+        **extra_attributes: Any,
+    ) -> None:
+        """End the span for the given run."""
+        if run_id not in self.spans:
+            return
+
+        try:
+            span = self.spans[run_id]
+
+            # Set output data if provided
+            if outputs is not None:
+                span.set_attribute(
+                    AttributeKeys.JUDGMENT_OUTPUT.value, safe_serialize(outputs)
+                )
+
+            # Set additional attributes
+            for key, value in extra_attributes.items():
+                if value is not None:
+                    span.set_attribute(str(key), str(value))
+
+            # Handle errors
+            if error is not None:
+                # Check if this is a control flow exception
+                is_control_flow = any(
+                    isinstance(error, t) for t in CONTROL_FLOW_EXCEPTION_TYPES
+                )
+                if not is_control_flow:
+                    span.record_exception(error)
+                    span.set_status(Status(StatusCode.ERROR, str(error)))
+                # Control flow exceptions don't set error status
+            else:
+                span.set_status(Status(StatusCode.OK))
+
+            # Note: The span will be ended automatically by the context manager
+
+        except Exception as e:
+            judgeval_logger.exception(f"Error ending span for run_id {run_id}: {e}")
+        finally:
+            # Cleanup tracking data
+            if run_id in self.spans:
+                del self.spans[run_id]
+            if run_id in self.span_start_times:
+                del self.span_start_times[run_id]
+            if run_id in self.run_id_to_span_id:
+                span_id = self.run_id_to_span_id[run_id]
+                del self.run_id_to_span_id[run_id]
+                if span_id in self.span_id_to_depth:
+                    del self.span_id_to_depth[span_id]
+
+            # Check if this is the root run ending
+            if run_id == self.root_run_id:
+                self.root_run_id = None
+
+    def _log_debug_event(
+        self,
+        event_name: str,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Log debug information about callback events."""
+        judgeval_logger.debug(
+            f"Event: {event_name}, run_id: {str(run_id)[:8]}, "
+            f"parent_run_id: {str(parent_run_id)[:8] if parent_run_id else None}"
+        )
+
+    # Chain callbacks
+    def on_chain_start(
+        self,
+        serialized: Optional[Dict[str, Any]],
+        inputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chain starts running."""
+        try:
+            self._log_debug_event(
+                "on_chain_start", run_id, parent_run_id, inputs=inputs
+            )
+
+            name = self._get_run_name(serialized, **kwargs)
+
+            # Check for LangGraph node
+            node_name = metadata.get("langgraph_node") if metadata else None
+            if node_name:
+                name = node_name
+                if name not in self.executed_nodes:
+                    self.executed_nodes.append(name)
+
+            # Determine if this is a root LangGraph execution
+            is_langgraph_root = (
+                kwargs.get("name") == "LangGraph" and parent_run_id is None
+            )
+            if is_langgraph_root:
+                name = "LangGraph"
+
+            span_level = "DEBUG" if tags and LANGSMITH_TAG_HIDDEN in tags else None
+
+            self._start_span(
+                run_id=run_id,
+                parent_run_id=parent_run_id,
+                name=name,
+                span_type="chain",
+                inputs=inputs,
+                tags=tags,
+                metadata=metadata,
+                level=span_level,
+                serialized=safe_serialize(serialized) if serialized else None,
+            )
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_chain_start: {e}")
+
+    def on_chain_end(
+        self,
+        outputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chain ends successfully."""
+        try:
+            self._log_debug_event(
+                "on_chain_end", run_id, parent_run_id, outputs=outputs
+            )
+            self._end_span(run_id=run_id, outputs=outputs)
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_chain_end: {e}")
+
+    def on_chain_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Called when a chain encounters an error."""
+        try:
+            self._log_debug_event("on_chain_error", run_id, parent_run_id, error=error)
+            self._end_span(run_id=run_id, error=error)
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_chain_error: {e}")
+
+    # LLM callbacks
+    def on_llm_start(
+        self,
+        serialized: Optional[Dict[str, Any]],
+        prompts: List[str],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an LLM starts generating."""
+        try:
+            self._log_debug_event(
+                "on_llm_start", run_id, parent_run_id, prompts=prompts
+            )
+
+            name = self._get_run_name(serialized, **kwargs)
+            model_name = self._extract_model_name(serialized, kwargs)
+
+            prompt_data = prompts[0] if len(prompts) == 1 else prompts
+
+            self._start_span(
+                run_id=run_id,
+                parent_run_id=parent_run_id,
+                name=name,
+                span_type="llm",
+                inputs=prompt_data,
+                tags=tags,
+                metadata=metadata,
+                model=model_name,
+                serialized=safe_serialize(serialized) if serialized else None,
+            )
+
+            # Set GenAI specific attributes
+            if run_id in self.spans:
+                span = self.spans[run_id]
+                if model_name:
+                    span.set_attribute(AttributeKeys.GEN_AI_REQUEST_MODEL, model_name)
+                span.set_attribute(
+                    AttributeKeys.GEN_AI_PROMPT, safe_serialize(prompt_data)
+                )
+
+                # Set model parameters if available
+                invocation_params = kwargs.get("invocation_params", {})
+                if "temperature" in invocation_params:
+                    span.set_attribute(
+                        AttributeKeys.GEN_AI_REQUEST_TEMPERATURE,
+                        float(invocation_params["temperature"]),
+                    )
+                if "max_tokens" in invocation_params:
+                    span.set_attribute(
+                        AttributeKeys.GEN_AI_REQUEST_MAX_TOKENS,
+                        int(invocation_params["max_tokens"]),
+                    )
+
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_llm_start: {e}")
+
+    def on_chat_model_start(
+        self,
+        serialized: Optional[Dict[str, Any]],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a chat model starts generating."""
+        try:
+            self._log_debug_event(
+                "on_chat_model_start", run_id, parent_run_id, messages=messages
+            )
+
+            name = self._get_run_name(serialized, **kwargs)
+            model_name = self._extract_model_name(serialized, kwargs)
+
+            # Flatten messages
+            flattened_messages = []
+            for message_list in messages:
+                flattened_messages.extend(self._create_message_dicts(message_list))
+
+            self._start_span(
+                run_id=run_id,
+                parent_run_id=parent_run_id,
+                name=name,
+                span_type="llm",
+                inputs=flattened_messages,
+                tags=tags,
+                metadata=metadata,
+                model=model_name,
+                serialized=safe_serialize(serialized) if serialized else None,
+            )
+
+            # Set GenAI specific attributes
+            if run_id in self.spans:
+                span = self.spans[run_id]
+                if model_name:
+                    span.set_attribute(AttributeKeys.GEN_AI_REQUEST_MODEL, model_name)
+                span.set_attribute(
+                    AttributeKeys.GEN_AI_PROMPT, safe_serialize(flattened_messages)
+                )
+
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_chat_model_start: {e}")
+
+    def on_llm_end(
+        self,
+        response: LLMResult,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an LLM finishes generating."""
+        try:
+            self._log_debug_event(
+                "on_llm_end", run_id, parent_run_id, response=response
+            )
+
+            # Extract response content
+            if response.generations:
+                last_generation = response.generations[-1][-1]
+                if (
+                    isinstance(last_generation, ChatGeneration)
+                    and last_generation.message
+                ):
+                    output = self._convert_message_to_dict(last_generation.message)
+                else:
+                    output = (
+                        last_generation.text
+                        if hasattr(last_generation, "text")
+                        else str(last_generation)
+                    )
+            else:
+                output = ""
+
+            # Extract usage information
+            usage_attrs = {}
+            if response.llm_output and "token_usage" in response.llm_output:
+                token_usage = response.llm_output["token_usage"]
+                if hasattr(token_usage, "prompt_tokens"):
+                    usage_attrs[AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS] = (
+                        token_usage.prompt_tokens
+                    )
+                if hasattr(token_usage, "completion_tokens"):
+                    usage_attrs[AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS] = (
+                        token_usage.completion_tokens
+                    )
+
+            # Set completion attribute
+            if run_id in self.spans:
+                span = self.spans[run_id]
+                span.set_attribute(
+                    AttributeKeys.GEN_AI_COMPLETION, safe_serialize(output)
+                )
+
+                # Set usage attributes
+                for key, value in usage_attrs.items():
+                    span.set_attribute(key, value)
+
+            self._end_span(run_id=run_id, outputs=output, **usage_attrs)
+
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_llm_end: {e}")
+
+    def on_llm_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an LLM encounters an error."""
+        try:
+            self._log_debug_event("on_llm_error", run_id, parent_run_id, error=error)
+            self._end_span(run_id=run_id, error=error)
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_llm_error: {e}")
+
+    # Tool callbacks
+    def on_tool_start(
+        self,
+        serialized: Optional[Dict[str, Any]],
+        input_str: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a tool starts executing."""
+        try:
+            self._log_debug_event(
+                "on_tool_start", run_id, parent_run_id, input_str=input_str
+            )
+
+            name = self._get_run_name(serialized, **kwargs)
+            if name not in self.executed_tools:
+                self.executed_tools.append(name)
+
+            self._start_span(
+                run_id=run_id,
+                parent_run_id=parent_run_id,
+                name=name,
+                span_type="tool",
+                inputs=input_str,
+                tags=tags,
+                metadata=metadata,
+                serialized=safe_serialize(serialized) if serialized else None,
+            )
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_tool_start: {e}")
+
+    def on_tool_end(
+        self,
+        output: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a tool finishes executing."""
+        try:
+            self._log_debug_event("on_tool_end", run_id, parent_run_id, output=output)
+            self._end_span(run_id=run_id, outputs=output)
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_tool_end: {e}")
+
+    def on_tool_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a tool encounters an error."""
+        try:
+            self._log_debug_event("on_tool_error", run_id, parent_run_id, error=error)
+            self._end_span(run_id=run_id, error=error)
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_tool_error: {e}")
+
+    # Agent callbacks
+    def on_agent_action(
+        self,
+        action: AgentAction,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an agent takes an action."""
+        try:
+            self._log_debug_event(
+                "on_agent_action", run_id, parent_run_id, action=action
+            )
+
+            if run_id in self.spans:
+                span = self.spans[run_id]
+                span.set_attribute("agent.action.tool", action.tool)
+                span.set_attribute(
+                    "agent.action.tool_input", safe_serialize(action.tool_input)
+                )
+                span.set_attribute("agent.action.log", action.log)
+
+            self._end_span(
+                run_id=run_id,
+                outputs={"action": action.tool, "input": action.tool_input},
+            )
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_agent_action: {e}")
+
+    def on_agent_finish(
+        self,
+        finish: AgentFinish,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when an agent finishes."""
+        try:
+            self._log_debug_event(
+                "on_agent_finish", run_id, parent_run_id, finish=finish
+            )
+
+            if run_id in self.spans:
+                span = self.spans[run_id]
+                span.set_attribute("agent.finish.log", finish.log)
+
+            self._end_span(run_id=run_id, outputs=finish.return_values)
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_agent_finish: {e}")
+
+    # Retriever callbacks
+    def on_retriever_start(
+        self,
+        serialized: Optional[Dict[str, Any]],
+        query: str,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        tags: Optional[List[str]] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a retriever starts."""
+        try:
+            self._log_debug_event(
+                "on_retriever_start", run_id, parent_run_id, query=query
+            )
+
+            name = self._get_run_name(serialized, **kwargs)
+
+            self._start_span(
+                run_id=run_id,
+                parent_run_id=parent_run_id,
+                name=name,
+                span_type="retriever",
+                inputs=query,
+                tags=tags,
+                metadata=metadata,
+                serialized=safe_serialize(serialized) if serialized else None,
+            )
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_retriever_start: {e}")
+
+    def on_retriever_end(
+        self,
+        documents: Sequence[Document],
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a retriever finishes."""
+        try:
+            self._log_debug_event(
+                "on_retriever_end", run_id, parent_run_id, documents=documents
+            )
+
+            # Convert documents to serializable format
+            doc_data = [
+                {"page_content": doc.page_content, "metadata": doc.metadata}
+                for doc in documents
+            ]
+
+            if run_id in self.spans:
+                span = self.spans[run_id]
+                span.set_attribute("retriever.document_count", len(documents))
+
+            self._end_span(
+                run_id=run_id, outputs=doc_data, document_count=len(documents)
+            )
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_retriever_end: {e}")
+
+    def on_retriever_error(
+        self,
+        error: BaseException,
+        *,
+        run_id: UUID,
+        parent_run_id: Optional[UUID] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Called when a retriever encounters an error."""
+        try:
+            self._log_debug_event(
+                "on_retriever_error", run_id, parent_run_id, error=error
+            )
+            self._end_span(run_id=run_id, error=error)
+        except Exception as e:
+            judgeval_logger.exception(f"Error in on_retriever_error: {e}")
+
+    def _extract_model_name(
+        self, serialized: Optional[Dict[str, Any]], kwargs: Dict[str, Any]
+    ) -> Optional[str]:
+        """Extract model name from serialized data or kwargs."""
+        # Try to get from invocation params
+        invocation_params = kwargs.get("invocation_params", {})
+        if "model_name" in invocation_params:
+            return invocation_params["model_name"]
+        if "model" in invocation_params:
+            return invocation_params["model"]
+
+        # Try to get from serialized data
+        if serialized:
+            if "model_name" in serialized:
+                return serialized["model_name"]
+            if "model" in serialized:
+                return serialized["model"]
+
+        return None
+
+
+__all__ = ["JudgevalCallbackHandler"]