finished

Author: xzrderek

eval_protocol/adapters/langfuse.py

@@ -4,6 +4,7 @@
 to EvaluationRow format for use in evaluation pipelines.
 """
 
+from langfuse.api.resources.commons.types.observations_view import ObservationsView
 import logging
 from datetime import datetime, timedelta
 from typing import Any, Dict, Iterator, List, Optional, cast
@@ -15,6 +16,7 @@
 try:
     from langfuse import get_client  # pyright: ignore[reportPrivateImportUsage]
     from langfuse.api.resources.trace.types.traces import Traces
+    from langfuse.api.resources.commons.types.trace import Trace
     from langfuse.api.resources.commons.types.trace_with_full_details import TraceWithFullDetails
 
     LANGFUSE_AVAILABLE = True
@@ -86,6 +88,7 @@ def get_evaluation_rows(
             from_timestamp = None
 
         eval_rows = []
+
         traces: Traces = self.client.api.trace.list(
             limit=limit,
             tags=tags,
@@ -131,7 +134,9 @@ def get_evaluation_rows_by_ids(
                 continue
         return eval_rows
 
-    def _convert_trace_to_evaluation_row(self, trace: Any, include_tool_calls: bool = True) -> Optional[EvaluationRow]:
+    def _convert_trace_to_evaluation_row(
+        self, trace: Trace, include_tool_calls: bool = True
+    ) -> Optional[EvaluationRow]:
         """Convert a Langfuse trace to EvaluationRow format.
 
         Args:
@@ -142,96 +147,78 @@ def _convert_trace_to_evaluation_row(self, trace: Any, include_tool_calls: bool
             EvaluationRow or None if conversion fails
         """
         try:
-            # Get observations (generations, spans) from the trace
-            observations_response = self.client.api.observations.get_many(trace_id=trace.id, limit=100)
-            # print(observations_response)
-            observations = (
-                observations_response.data if hasattr(observations_response, "data") else list(observations_response)
-            )
+            trace = self.client.api.trace.get("2d9f3474-83ab-4431-9788-049ca4219023")
 
-            messages = []
+            # Extract messages from trace input and output
+            messages = self._extract_messages_from_trace(trace, include_tool_calls)
 
-            for obs in observations:
-                if obs.name == "agent run":
-                    messages = self._extract_conversation_from_output(obs.output)
-                    break
+            # Extract tools if available
+            tools = None
+            if include_tool_calls and isinstance(trace.input, dict) and "tools" in trace.input:
+                tools = trace.input["tools"]
 
             if not messages:
                 return None
 
-            # Extract metadata
-            input_metadata = self._create_input_metadata(trace, observations)
-
-            # Extract ground truth if available (from trace metadata or tags)
-            ground_truth = self._extract_ground_truth(trace)
-
-            # Extract tools if available
-            tools = self._extract_tools(observations, trace) if include_tool_calls else None
-
             return EvaluationRow(
                 messages=messages,
                 tools=tools,
-                input_metadata=input_metadata,
-                ground_truth=ground_truth,
             )
 
         except (AttributeError, ValueError, KeyError) as e:
             logger.error("Error converting trace %s: %s", trace.id, e)
             return None
 
-    def _extract_messages_from_observations(
-        self, observations: List[Any], include_tool_calls: bool = True
-    ) -> List[Message]:
-        """Extract messages from Langfuse observations.
+    def _extract_messages_from_trace(self, trace: Any, include_tool_calls: bool = True) -> List[Message]:
+        """Extract messages from Langfuse trace input and output.
 
         Args:
-            observations: List of Langfuse observation objects
+            trace: Langfuse trace object
             include_tool_calls: Whether to include tool calling information
 
         Returns:
             List of Message objects
         """
         messages = []
 
-        # Sort observations by timestamp
-        sorted_observations = sorted(observations, key=lambda x: x.start_time or datetime.min)
-
-        for obs in sorted_observations:
-            try:
-                if hasattr(obs, "input") and obs.input:
-                    # Handle different input formats
-                    if isinstance(obs.input, dict):
-                        if "messages" in obs.input:
-                            # OpenAI-style messages format
-                            for msg in obs.input["messages"]:
-                                messages.append(self._dict_to_message(msg, include_tool_calls))
-                        elif "role" in obs.input:
-                            # Single message format
-                            messages.append(self._dict_to_message(obs.input, include_tool_calls))
-                        elif "prompt" in obs.input:
-                            # Simple prompt format
-                            messages.append(Message(role="user", content=str(obs.input["prompt"])))
-                    elif isinstance(obs.input, str):
-                        # Simple string input
-                        messages.append(Message(role="user", content=obs.input))
-
-                if hasattr(obs, "output") and obs.output:
-                    # Handle output
-                    if isinstance(obs.output, dict):
-                        if "content" in obs.output:
-                            messages.append(Message(role="assistant", content=str(obs.output["content"])))
-                        elif "message" in obs.output:
-                            msg_dict = obs.output["message"]
-                            messages.append(self._dict_to_message(msg_dict, include_tool_calls))
-                        else:
-                            # Fallback: convert entire output to string
-                            messages.append(Message(role="assistant", content=str(obs.output)))
-                    elif isinstance(obs.output, str):
-                        messages.append(Message(role="assistant", content=obs.output))
+        try:
+            # Handle trace input
+            if hasattr(trace, "input") and trace.input:
+                if isinstance(trace.input, dict):
+                    if "messages" in trace.input:
+                        # OpenAI-style messages format
+                        for msg in trace.input["messages"]:
+                            messages.append(self._dict_to_message(msg, include_tool_calls))
+                    elif "role" in trace.input:
+                        # Single message format
+                        messages.append(self._dict_to_message(trace.input, include_tool_calls))
+                    elif "prompt" in trace.input:
+                        # Simple prompt format
+                        messages.append(Message(role="user", content=str(trace.input["prompt"])))
+                elif isinstance(trace.input, list):
+                    # Direct list of message dicts
+                    for msg in trace.input:
+                        messages.append(self._dict_to_message(msg, include_tool_calls))
+                elif isinstance(trace.input, str):
+                    # Simple string input
+                    messages.append(Message(role="user", content=trace.input))
+
+            # Handle trace output
+            if hasattr(trace, "output") and trace.output:
+                if isinstance(trace.output, dict):
+                    if "content" in trace.output:
+                        messages.append(Message(role="assistant", content=str(trace.output["content"])))
+                    elif "message" in trace.output:
+                        msg_dict = trace.output["message"]
+                        messages.append(self._dict_to_message(msg_dict, include_tool_calls))
+                    else:
+                        # Fallback: convert entire output to string
+                        messages.append(Message(role="assistant", content=str(trace.output)))
+                elif isinstance(trace.output, str):
+                    messages.append(Message(role="assistant", content=trace.output))
 
-            except (AttributeError, ValueError, KeyError) as e:
-                logger.warning("Error processing observation %s: %s", obs.id, e)
-                continue
+        except (AttributeError, ValueError, KeyError) as e:
+            logger.warning("Error processing trace %s: %s", trace.id, e)
 
         return messages
 
@@ -272,220 +259,6 @@ def _dict_to_message(self, msg_dict: Dict[str, Any], include_tool_calls: bool =
             function_call=function_call,
         )
 
-    def _extract_conversation_from_output(self, output: Any) -> Optional[List[Message]]:
-        """Extract conversation history from PydanticAI agent run output.
-
-        This looks for the conversation format like:
-        [
-            {"role": "user", "content": "..."},
-            {"role": "assistant", "content": "...", "tool_calls": [...]},
-            {"role": "tool", "content": "...", "name": "execute_sql"},
-            ...
-        ]
-
-        Args:
-            output: The output object to search for conversation history
-
-        Returns:
-            List of Message objects or None if no conversation found
-        """
-        messages = []
-
-        try:
-            # Handle different output formats
-            conversation_data = None
-
-            if isinstance(output, list):
-                # Direct list of messages
-                conversation_data = output
-            elif isinstance(output, dict):
-                # Look for conversation in various nested formats
-                if "messages" in output:
-                    conversation_data = output["messages"]
-                elif "conversation" in output:
-                    conversation_data = output["conversation"]
-                elif "history" in output:
-                    conversation_data = output["history"]
-                elif "agent_run" in output:  # Handle nested conversation data PydanticAI style
-                    agent_run = output["agent_run"]
-                    if isinstance(agent_run, dict) and "messages" in agent_run:
-                        conversation_data = agent_run["messages"]
-                elif len(output.keys()) == 1:
-                    # Single key, check if its value is a list
-                    single_key = list(output.keys())[0]
-                    if isinstance(output[single_key], list):
-                        conversation_data = output[single_key]
-            elif isinstance(output, str):
-                # Try to parse JSON string
-                import json
-
-                try:
-                    parsed = json.loads(output)
-                    return self._extract_conversation_from_output(parsed)
-                except (json.JSONDecodeError, ValueError):
-                    pass
-
-            # Parse conversation data into messages
-            if conversation_data and isinstance(conversation_data, list):
-                for msg_data in conversation_data:
-                    if isinstance(msg_data, dict) and "role" in msg_data:
-                        role = msg_data.get("role")
-                        if role is None:
-                            continue
-                        content = msg_data.get("content", "")
-
-                        # Handle tool calls in assistant messages
-                        tool_calls = None
-                        if role == "assistant" and "tool_calls" in msg_data:
-                            tool_calls = msg_data["tool_calls"]
-
-                        # Handle tool responses
-                        name = None
-                        tool_call_id = None
-                        if role == "tool":
-                            name = msg_data.get("name")
-                            tool_call_id = msg_data.get("id")
-
-                        messages.append(
-                            Message(
-                                role=role, content=content, name=name, tool_calls=tool_calls, tool_call_id=tool_call_id
-                            )
-                        )
-
-            return messages if messages else None
-
-        except Exception as e:
-            logger.debug("Error extracting conversation from output: %s", e)
-            return None
-
-    def _create_input_metadata(self, trace: Any, observations: List[Any]) -> InputMetadata:
-        """Create InputMetadata from trace and observations.
-
-        Args:
-            trace: Langfuse trace object
-            observations: List of observation objects
-
-        Returns:
-            InputMetadata object
-        """
-        # Extract completion parameters from trace input first, then observations
-        completion_params = {}
-
-        # First check trace input for evaluation test completion_params
-        if hasattr(trace, "input") and trace.input:
-            if isinstance(trace.input, dict):
-                kwargs = trace.input.get("kwargs", {})
-                if "completion_params" in kwargs:
-                    trace_completion_params = kwargs["completion_params"]
-                    if trace_completion_params and isinstance(trace_completion_params, dict):
-                        completion_params.update(trace_completion_params)
-
-        # Fallback: Look for model parameters in observations if not found in trace input
-        if not completion_params:
-            for obs in observations:
-                if hasattr(obs, "model") and obs.model:
-                    completion_params["model"] = obs.model
-                if hasattr(obs, "model_parameters") and obs.model_parameters:
-                    params = obs.model_parameters
-                    if "temperature" in params:
-                        completion_params["temperature"] = params["temperature"]
-                    if "max_tokens" in params:
-                        completion_params["max_tokens"] = params["max_tokens"]
-                    if "top_p" in params:
-                        completion_params["top_p"] = params["top_p"]
-                    break
-
-        # Create dataset info from trace metadata
-        dataset_info = {
-            "trace_id": trace.id,
-            "trace_name": getattr(trace, "name", None),
-            "trace_tags": getattr(trace, "tags", []),
-        }
-
-        # Add trace metadata if available
-        if hasattr(trace, "metadata") and trace.metadata:
-            dataset_info["trace_metadata"] = trace.metadata
-
-        # Create session data
-        session_data = {
-            "session_id": getattr(trace, "session_id", None),
-            "user_id": getattr(trace, "user_id", None),
-            "timestamp": getattr(trace, "timestamp", None),
-        }
-
-        return InputMetadata(
-            row_id=trace.id,
-            completion_params=completion_params,
-            dataset_info=dataset_info,
-            session_data=session_data,
-        )
-
-    def _extract_ground_truth(self, trace: Any) -> Optional[str]:
-        """Extract ground truth from trace if available.
-
-        Args:
-            trace: Langfuse trace object
-
-        Returns:
-            Ground truth string or None
-        """
-        # First check trace input for evaluation test data structure
-        if hasattr(trace, "input") and trace.input:
-            if isinstance(trace.input, dict):
-                # Handle EP test format: kwargs.input_rows[0].ground_truth
-                kwargs = trace.input.get("kwargs", {})
-                if "input_rows" in kwargs:
-                    input_rows = kwargs["input_rows"]
-                    if input_rows and len(input_rows) > 0:
-                        first_row = input_rows[0]
-                        if isinstance(first_row, dict) and "ground_truth" in first_row:
-                            ground_truth = first_row["ground_truth"]
-                            if ground_truth:  # Only return if not None/empty
-                                return str(ground_truth)
-
-        # Check trace metadata for ground truth
-        if hasattr(trace, "metadata") and trace.metadata:
-            if isinstance(trace.metadata, dict):
-                return trace.metadata.get("ground_truth") or trace.metadata.get("expected_answer")
-
-        # Check tags for ground truth indicators
-        if hasattr(trace, "tags") and trace.tags:
-            for tag in trace.tags:
-                if tag.startswith("ground_truth:"):
-                    return tag.replace("ground_truth:", "", 1)
-
-        return None
-
-    def _extract_tools(self, observations: List[Any], trace: Any = None) -> Optional[List[Dict[str, Any]]]:
-        """Extract tool definitions from trace metadata or observations.
-
-        Args:
-            observations: List of observation objects
-            trace: Trace object that may contain metadata with tools
-
-        Returns:
-            List of tool definitions or None
-        """
-        # First, try to extract tools from trace metadata (preferred)
-        if trace and hasattr(trace, "metadata") and trace.metadata:
-            if isinstance(trace.metadata, dict) and "tools" in trace.metadata:
-                tools_from_metadata = trace.metadata["tools"]
-                if tools_from_metadata:
-                    return tools_from_metadata
-
-        # Fallback: extract from observations
-        tools = []
-        for obs in observations:
-            if hasattr(obs, "input") and obs.input and isinstance(obs.input, dict):
-                if "tools" in obs.input:
-                    tools.extend(obs.input["tools"])
-                elif "functions" in obs.input:
-                    # Convert functions to tools format
-                    for func in obs.input["functions"]:
-                        tools.append({"type": "function", "function": func})
-
-        return tools if tools else None
-
 
 def create_langfuse_adapter() -> LangfuseAdapter:
     """Factory function to create a Langfuse adapter."""