judgeval 0.0.40__py3-none-any.whl → 0.0.42__py3-none-any.whl

judgeval/common/utils.py

@@ -12,9 +12,10 @@ NOTE: any function beginning with 'a', e.g. 'afetch_together_api_response', is a
 import asyncio
 import concurrent.futures
 import os
+from types import TracebackType
 import requests
 import pprint
-from typing import Any, Dict, List, Literal, Mapping, Optional, Union
+from typing import Any, Dict, List, Literal, Mapping, Optional, TypeAlias, Union
 
 # Third-party imports
 import litellm
@@ -102,7 +103,7 @@ def validate_api_key(judgment_api_key: str):
     Validates that the user api key is valid
     """
     response = requests.post(
-        f"{ROOT_API}/validate_api_key/",
+        f"{ROOT_API}/auth/validate_api_key/",
         headers={
             "Content-Type": "application/json",
             "Authorization": f"Bearer {judgment_api_key}",
@@ -782,3 +783,6 @@ if __name__ == "__main__":
             ]
         ]
     ))
+    
+ExcInfo: TypeAlias = tuple[type[BaseException], BaseException, TracebackType]
+OptExcInfo: TypeAlias = ExcInfo | tuple[None, None, None]

judgeval/constants.py

@@ -58,8 +58,13 @@ JUDGMENT_PROJECT_DELETE_API_URL = f"{ROOT_API}/projects/delete/"
 JUDGMENT_PROJECT_CREATE_API_URL = f"{ROOT_API}/projects/add/"
 JUDGMENT_TRACES_FETCH_API_URL = f"{ROOT_API}/traces/fetch/"
 JUDGMENT_TRACES_SAVE_API_URL = f"{ROOT_API}/traces/save/"
+JUDGMENT_TRACES_UPSERT_API_URL = f"{ROOT_API}/traces/upsert/"
+JUDGMENT_TRACES_USAGE_CHECK_API_URL = f"{ROOT_API}/traces/usage/check/"
+JUDGMENT_TRACES_USAGE_UPDATE_API_URL = f"{ROOT_API}/traces/usage/update/"
 JUDGMENT_TRACES_DELETE_API_URL = f"{ROOT_API}/traces/delete/"
 JUDGMENT_TRACES_ADD_ANNOTATION_API_URL = f"{ROOT_API}/traces/add_annotation/"
+JUDGMENT_TRACES_SPANS_BATCH_API_URL = f"{ROOT_API}/traces/spans/batch/"
+JUDGMENT_TRACES_EVALUATION_RUNS_BATCH_API_URL = f"{ROOT_API}/traces/evaluation_runs/batch/"
 JUDGMENT_ADD_TO_RUN_EVAL_QUEUE_API_URL = f"{ROOT_API}/add_to_run_eval_queue/"
 JUDGMENT_GET_EVAL_STATUS_API_URL = f"{ROOT_API}/get_evaluation_status/"
 # RabbitMQ

judgeval/data/datasets/dataset.py

@@ -5,14 +5,15 @@ import json
 import os
 import yaml
 from dataclasses import dataclass, field
-from typing import List, Union, Literal
+from typing import List, Union, Literal, Optional
 
-from judgeval.data import Example
+from judgeval.data import Example, Trace
 from judgeval.common.logger import debug, error, warning, info
 
 @dataclass
 class EvalDataset:
     examples: List[Example]
+    traces: List[Trace]
     _alias: Union[str, None] = field(default=None)
     _id: Union[str, None] = field(default=None)
     judgment_api_key: str = field(default="")
@@ -20,12 +21,13 @@ class EvalDataset:
     def __init__(self, 
                  judgment_api_key: str = os.getenv("JUDGMENT_API_KEY"),  
                  organization_id: str = os.getenv("JUDGMENT_ORG_ID"),
-                 examples: List[Example] = [],
+                 examples: Optional[List[Example]] = None,
+                 traces: Optional[List[Trace]] = None
                  ):
-        debug(f"Initializing EvalDataset with {len(examples)} examples")
         if not judgment_api_key:
             warning("No judgment_api_key provided")
-        self.examples = examples
+        self.examples = examples or []
+        self.traces = traces or []
         self._alias = None
         self._id = None
         self.judgment_api_key = judgment_api_key
@@ -218,8 +220,11 @@ class EvalDataset:
             self.add_example(e)
 
     def add_example(self, e: Example) -> None:
-        self.examples = self.examples + [e]
+        self.examples.append(e)
         # TODO if we need to add rank, then we need to do it here
+    
+    def add_trace(self, t: Trace) -> None:
+        self.traces.append(t)
 
     def save_as(self, file_type: Literal["json", "csv", "yaml"], dir_path: str, save_name: str = None) -> None:
         """
@@ -307,6 +312,7 @@ class EvalDataset:
         return (
             f"{self.__class__.__name__}("
             f"examples={self.examples}, "
+            f"traces={self.traces}, "
             f"_alias={self._alias}, "
             f"_id={self._id}"
             f")"

judgeval/data/datasets/eval_dataset_client.py

@@ -13,7 +13,7 @@ from judgeval.constants import (
     JUDGMENT_DATASETS_INSERT_API_URL,
     JUDGMENT_DATASETS_EXPORT_JSONL_API_URL
 )
-from judgeval.data import Example
+from judgeval.data import Example, Trace
 from judgeval.data.datasets import EvalDataset
 
 
@@ -58,6 +58,7 @@ class EvalDatasetClient:
                     "dataset_alias": alias,
                     "project_name": project_name,
                     "examples": [e.to_dict() for e in dataset.examples],
+                    "traces": [t.model_dump() for t in dataset.traces],
                     "overwrite": overwrite,
                 }
             try:
@@ -202,6 +203,7 @@ class EvalDatasetClient:
                 info(f"Successfully pulled dataset with alias '{alias}'")
                 payload = response.json()
                 dataset.examples = [Example(**e) for e in payload.get("examples", [])]
+                dataset.traces = [Trace(**t) for t in payload.get("traces", [])]
                 dataset._alias = payload.get("alias")
                 dataset._id = payload.get("id")
                 progress.update(

judgeval/data/trace.py

@@ -33,6 +33,8 @@ class TraceSpan(BaseModel):
     additional_metadata: Optional[Dict[str, Any]] = None
     has_evaluation: Optional[bool] = False
     agent_name: Optional[str] = None
+    state_before: Optional[Dict[str, Any]] = None
+    state_after: Optional[Dict[str, Any]] = None
 
     def model_dump(self, **kwargs):
         return {
@@ -50,7 +52,10 @@ class TraceSpan(BaseModel):
             "span_type": self.span_type,
             "usage": self.usage.model_dump() if self.usage else None,
             "has_evaluation": self.has_evaluation,
-            "agent_name": self.agent_name
+            "agent_name": self.agent_name,
+            "state_before": self.state_before,
+            "state_after": self.state_after,
+            "additional_metadata": self._serialize_value(self.additional_metadata)
         }
     
     def print_span(self):
@@ -113,7 +118,7 @@ class Trace(BaseModel):
     name: str
     created_at: str
     duration: float
-    entries: List[TraceSpan]
+    trace_spans: List[TraceSpan]
     overwrite: bool = False
     offline_mode: bool = False
     rules: Optional[Dict[str, Any]] = None

judgeval/integrations/langgraph.py

@@ -3,9 +3,11 @@ from uuid import UUID
 import time
 import uuid
 import contextvars # <--- Import contextvars
+from datetime import datetime
 
-from judgeval.common.tracer import TraceClient, TraceSpan, Tracer, SpanType, EvaluationConfig
+from judgeval.common.tracer import TraceClient, TraceSpan, Tracer, SpanType, EvaluationConfig, cost_per_token
 from judgeval.data import Example # Import Example
+from judgeval.data.trace import TraceUsage
 
 from langchain_core.callbacks import BaseCallbackHandler
 from langchain_core.agents import AgentAction, AgentFinish
@@ -36,18 +38,48 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
     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
         self._trace_client: Optional[TraceClient] = None
         self._run_id_to_span_id: Dict[UUID, str] = {}
         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.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._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
+
+    def reset(self):
+        """Public method to manually reset handler execution state for reuse"""
+        self._reset_state()
+
+    def reset_all(self):
+        """Public method to reset ALL handler state including tracking/logging data"""
+        self._reset_state()
+        # Also reset tracking/logging variables
+        self.executed_nodes: List[str] = []
         self.executed_tools: List[str] = []
         self.executed_node_tools: List[str] = []
         self.traces: List[Dict[str, Any]] = []
-    # --- END NEW __init__ ---
 
     # --- MODIFIED _ensure_trace_client ---
     def _ensure_trace_client(self, run_id: UUID, parent_run_id: Optional[UUID], event_name: str) -> Optional[TraceClient]:
@@ -57,6 +89,11 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         Returns the client or None.
         """
 
+        # If this is a potential new root execution (no parent_run_id) and we had a previous trace saved,
+        # reset state to allow reuse of the handler
+        if parent_run_id is None and self._trace_saved:
+            self._reset_state()
+
         # If a client already exists, return it.
         if self._trace_client:
             return self._trace_client
@@ -73,11 +110,25 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                 enable_evaluations=self.tracer.enable_evaluations
             )
             self._trace_client = client_instance
+            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.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_with_rate_limiting(
+                        overwrite=self._trace_client.overwrite, 
+                        final_save=False  # Initial save for live tracking
+                    )
+                except Exception as e:
+                    import warnings
+                    warnings.warn(f"Failed to save initial trace for live tracking: {e}")
+                
                 return self._trace_client
             else:
                 return None
@@ -112,12 +163,7 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         self._span_id_to_start_time[span_id] = start_time
         self._span_id_to_depth[span_id] = current_depth
 
-
-        # --- Set SPAN context variable ONLY for chain (node) spans (Sync version) ---
-        if span_type == "chain":
-            self.tracer.set_current_span(span_id)
-
-        new_trace = TraceSpan(
+        new_span = TraceSpan(
             span_id=span_id,
             trace_id=trace_client.trace_id,
             parent_span_id=parent_span_id,
@@ -127,9 +173,36 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             span_type=span_type
         )
 
-        new_trace.inputs = inputs
-
-        trace_client.add_span(new_trace)
+        # Separate metadata from inputs
+        if inputs:
+            metadata = {}
+            clean_inputs = {}
+            
+            # Extract metadata fields
+            metadata_fields = ['tags', 'metadata', 'kwargs', 'serialized']
+            for field in metadata_fields:
+                if field in inputs:
+                    metadata[field] = inputs.pop(field)
+            
+            # Store the remaining inputs
+            clean_inputs = inputs
+            
+            # Set both fields on the span
+            new_span.inputs = clean_inputs
+            new_span.additional_metadata = metadata
+        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")
+        
+        token = self.tracer.set_current_span(span_id)
+        if token:
+            self.span_id_to_token[span_id] = token
 
     def _end_span_tracking(
         self,
@@ -142,6 +215,8 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
 
         # Get span ID and check if it exists
         span_id = self._run_id_to_span_id.get(run_id)
+        token = self.span_id_to_token.pop(span_id, None)
+        self.tracer.reset_current_span(token, span_id)
 
         start_time = self._span_id_to_start_time.get(span_id) if span_id else None
         duration = time.time() - start_time if start_time is not None else None
@@ -151,7 +226,38 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
             trace_span = trace_client.span_id_to_span.get(span_id)
             if trace_span:
                 trace_span.duration = duration
-                trace_span.output = error if error else outputs
+                
+                # Handle outputs and error
+                if error:
+                    trace_span.output = error
+                elif outputs:
+                    # Separate metadata from outputs
+                    metadata = {}
+                    clean_outputs = {}
+                    
+                    # Extract metadata fields
+                    metadata_fields = ['tags', 'kwargs']
+                    if isinstance(outputs, dict):
+                        for field in metadata_fields:
+                            if field in outputs:
+                                metadata[field] = outputs.pop(field)
+                        
+                        # Store the remaining outputs
+                        clean_outputs = outputs
+                    else:
+                        clean_outputs = outputs
+                    
+                    # Set both fields on the span
+                    trace_span.output = clean_outputs
+                    if metadata:
+                        # Merge with existing metadata
+                        existing_metadata = trace_span.additional_metadata or {}
+                        trace_span.additional_metadata = {**existing_metadata, **metadata}
+                
+                # 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)
 
             # Clean up dictionaries for this specific span
             if span_id in self._span_id_to_start_time: del self._span_id_to_start_time[span_id]
@@ -165,9 +271,30 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                 # Reset input storage for this handler instance
 
                 if self._trace_client and not self._trace_saved: # Check if not already saved
-                    # TODO: Check if trace_client.save needs await if TraceClient becomes async
-                    trace_id, trace_data = self._trace_client.save(overwrite=self._trace_client.overwrite) # Use client's overwrite setting
-                    self.traces.append(trace_data) # Leaving this in for now but can probably be removed
+                    # Flush background spans before saving the final trace
+
+                    complete_trace_data = {
+                        "trace_id": self._trace_client.trace_id,
+                        "name": self._trace_client.name,
+                        "created_at": datetime.utcfromtimestamp(self._trace_client.start_time).isoformat(),
+                        "duration": self._trace_client.get_duration(),
+                        "trace_spans": [span.model_dump() for span in self._trace_client.trace_spans],
+                        "overwrite": self._trace_client.overwrite,
+                        "offline_mode": self.tracer.offline_mode,
+                        "parent_trace_id": self._trace_client.parent_trace_id,
+                        "parent_name": self._trace_client.parent_name
+                    }
+
+                    # NEW: Use save_with_rate_limiting with final_save=True for final save
+                    trace_id, trace_data = self._trace_client.save_with_rate_limiting(
+                        overwrite=self._trace_client.overwrite, 
+                        final_save=True  # Final save with usage counter updates
+                    )
+                    token = self.trace_id_to_token.pop(trace_id, None)
+                    self.tracer.reset_current_trace(token, trace_id)
+                    
+                    # Store complete trace data instead of server response
+                    self.tracer.traces.append(complete_trace_data)
                     self._trace_saved = True # Set flag only after successful save
             finally:
                 # --- NEW: Consolidated Cleanup Logic --- 
@@ -254,10 +381,26 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         # --- Root node cleanup (Existing logic - slightly modified save call) ---
         if run_id == self._root_run_id:
             if trace_client and not self._trace_saved:
-                # Save might need to be async if TraceClient methods become async
-                # Pass overwrite=True based on client's setting
-                trace_id_saved, trace_data = trace_client.save(overwrite=trace_client.overwrite)
-                self.traces.append(trace_data) # Leaving this in for now but can probably be removed
+                # Store complete trace data instead of server response
+                complete_trace_data = {
+                    "trace_id": trace_client.trace_id,
+                    "name": trace_client.name,
+                    "created_at": datetime.utcfromtimestamp(trace_client.start_time).isoformat(),
+                    "duration": trace_client.get_duration(),
+                    "trace_spans": [span.model_dump() for span in trace_client.trace_spans],
+                    "overwrite": trace_client.overwrite,
+                    "offline_mode": self.tracer.offline_mode,
+                    "parent_trace_id": trace_client.parent_trace_id,
+                    "parent_name": trace_client.parent_name
+                }
+                # NEW: Use save_with_rate_limiting with final_save=True for final save
+                trace_id_saved, trace_data = trace_client.save_with_rate_limiting(
+                    overwrite=trace_client.overwrite,
+                    final_save=True  # Final save with usage counter updates
+                )
+                
+                
+                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:
@@ -333,11 +476,23 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         if not trace_client:
             return
         outputs = {"response": response, "kwargs": kwargs}
-        # --- Token Usage Extraction and Accumulation ---
-        token_usage = None
-        prompt_tokens = None  # Use standard name
-        completion_tokens = None # Use standard name
+        
+        # --- Token Usage Extraction and Cost Calculation ---
+        prompt_tokens = None
+        completion_tokens = None
         total_tokens = None
+        model_name = None
+        
+        # Extract model name from response if available
+        if hasattr(response, 'llm_output') and response.llm_output and isinstance(response.llm_output, dict):
+            model_name = response.llm_output.get('model_name') or response.llm_output.get('model')
+        
+        # Try to get model from the first generation if available
+        if not model_name and response.generations and len(response.generations) > 0:
+            if hasattr(response.generations[0][0], 'generation_info') and response.generations[0][0].generation_info:
+                gen_info = response.generations[0][0].generation_info
+                model_name = gen_info.get('model') or gen_info.get('model_name')
+
         if response.llm_output and isinstance(response.llm_output, dict):
             # Check for OpenAI/standard 'token_usage' first
             if 'token_usage' in response.llm_output:
@@ -356,14 +511,43 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
                     if prompt_tokens is not None and completion_tokens is not None:
                         total_tokens = prompt_tokens + completion_tokens
 
-            # --- Store individual usage in span output and Accumulate --- 
+            # --- Create TraceUsage object and set on span ---
             if prompt_tokens is not None or completion_tokens is not None:
-                # Store individual usage for this span
-                outputs['usage'] = { 
-                    'prompt_tokens': prompt_tokens,
-                    'completion_tokens': completion_tokens,
-                    'total_tokens': total_tokens
-                }
+                # Calculate costs if model name is available
+                prompt_cost = None
+                completion_cost = None
+                total_cost_usd = None
+                
+                if model_name and prompt_tokens is not None and completion_tokens is not None:
+                    try:
+                        prompt_cost, completion_cost = cost_per_token(
+                            model=model_name,
+                            prompt_tokens=prompt_tokens,
+                            completion_tokens=completion_tokens
+                        )
+                        total_cost_usd = (prompt_cost + completion_cost) if prompt_cost and completion_cost else None
+                    except Exception as e:
+                        # If cost calculation fails, continue without costs
+                        import warnings
+                        warnings.warn(f"Failed to calculate token costs for model {model_name}: {e}")
+                
+                # Create TraceUsage object
+                usage = TraceUsage(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    total_tokens=total_tokens or (prompt_tokens + completion_tokens if prompt_tokens and completion_tokens else None),
+                    prompt_tokens_cost_usd=prompt_cost,
+                    completion_tokens_cost_usd=completion_cost,
+                    total_cost_usd=total_cost_usd,
+                    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
+                    
             
         self._end_span_tracking(trace_client, run_id, outputs=outputs)
         # --- End Token Usage ---
@@ -416,4 +600,4 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         if not trace_client: return
 
         outputs = {'return_values': finish.return_values, 'log': finish.log, 'messages': finish.messages, 'kwargs': kwargs}
-        self._end_span_tracking(trace_client, run_id, outputs=outputs)
+        self._end_span_tracking(trace_client, run_id, outputs=outputs)

judgeval/judgment_client.py

@@ -63,7 +63,15 @@ class SingletonMeta(type):
         return cls._instances[cls]
 
 class JudgmentClient(metaclass=SingletonMeta):
-    def __init__(self, judgment_api_key: str = os.getenv("JUDGMENT_API_KEY"), organization_id: str = os.getenv("JUDGMENT_ORG_ID")):
+    def __init__(self, judgment_api_key: Optional[str] = os.getenv("JUDGMENT_API_KEY"), organization_id: Optional[str] = os.getenv("JUDGMENT_ORG_ID")):
+        # Check if API key is None
+        if judgment_api_key is None:
+            raise ValueError("JUDGMENT_API_KEY cannot be None. Please provide a valid API key or set the JUDGMENT_API_KEY environment variable.")
+        
+        # Check if organization ID is None
+        if organization_id is None:
+            raise ValueError("JUDGMENT_ORG_ID cannot be None. Please provide a valid organization ID or set the JUDGMENT_ORG_ID environment variable.")
+        
         self.judgment_api_key = judgment_api_key
         self.organization_id = organization_id
         self.eval_dataset_client = EvalDatasetClient(judgment_api_key, organization_id)