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

judgeval/common/utils.py

@@ -103,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}",

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/trace.py

@@ -54,7 +54,8 @@ class TraceSpan(BaseModel):
             "has_evaluation": self.has_evaluation,
             "agent_name": self.agent_name,
             "state_before": self.state_before,
-            "state_after": self.state_after 
+            "state_after": self.state_after,
+            "additional_metadata": self._serialize_value(self.additional_metadata)
         }
     
     def print_span(self):

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/rules.py

@@ -9,13 +9,13 @@ import asyncio
 from concurrent.futures import ThreadPoolExecutor
 import time
 import uuid
+import os
+import re
+import json
+from datetime import datetime
 
 from judgeval.scorers import APIJudgmentScorer, JudgevalScorer
-
-class AlertStatus(str, Enum):
-    """Status of an alert evaluation."""
-    TRIGGERED = "triggered"
-    NOT_TRIGGERED = "not_triggered"
+from judgeval.utils.alerts import AlertStatus, AlertResult
 
 class Condition(BaseModel):
     """
@@ -68,6 +68,36 @@ class Condition(BaseModel):
             # Fallback to default comparison (greater than or equal)
             return value >= self.threshold if self.threshold is not None else False
 
+class PagerDutyConfig(BaseModel):
+    """
+    Configuration for PagerDuty notifications.
+    
+    Attributes:
+        routing_key: PagerDuty integration routing key
+        severity: Severity level (critical, error, warning, info)
+        source: Source of the alert (defaults to "judgeval")
+        component: Optional component that triggered the alert
+        group: Optional logical grouping for the alert
+        class_type: Optional class/type of alert event
+    """
+    routing_key: str
+    severity: str = "error"  # critical, error, warning, info
+    source: str = "judgeval"
+    component: Optional[str] = None
+    group: Optional[str] = None
+    class_type: Optional[str] = None
+    
+    def model_dump(self, **kwargs):
+        """Convert the PagerDutyConfig to a dictionary for JSON serialization."""
+        return {
+            "routing_key": self.routing_key,
+            "severity": self.severity,
+            "source": self.source,
+            "component": self.component,
+            "group": self.group,
+            "class_type": self.class_type
+        }
+
 class NotificationConfig(BaseModel):
     """
     Configuration for notifications when a rule is triggered.
@@ -75,8 +105,12 @@ class NotificationConfig(BaseModel):
     Example:
         {
             "enabled": true,
-            "communication_methods": ["email", "broadcast_slack", "broadcast_email"],
+            "communication_methods": ["email", "broadcast_slack", "broadcast_email", "pagerduty"],
             "email_addresses": ["user1@example.com", "user2@example.com"],
+            "pagerduty_config": {
+                "routing_key": "R0ABCD1234567890123456789",
+                "severity": "error"
+            },
             "send_at": 1632150000  # Unix timestamp (specific date/time)
         }
         
@@ -84,10 +118,12 @@ class NotificationConfig(BaseModel):
         - "email": Send emails to specified email addresses
         - "broadcast_slack": Send broadcast notifications to all configured Slack channels
         - "broadcast_email": Send broadcast emails to all organization emails
+        - "pagerduty": Send alerts to PagerDuty using the configured routing key
     """
     enabled: bool = True
     communication_methods: List[str] = []
     email_addresses: Optional[List[str]] = None
+    pagerduty_config: Optional[PagerDutyConfig] = None
     send_at: Optional[int] = None  # Unix timestamp for scheduled notifications
     
     def model_dump(self, **kwargs):
@@ -96,6 +132,7 @@ class NotificationConfig(BaseModel):
             "enabled": self.enabled,
             "communication_methods": self.communication_methods,
             "email_addresses": self.email_addresses,
+            "pagerduty_config": self.pagerduty_config.model_dump() if self.pagerduty_config else None,
             "send_at": self.send_at
         }
 
@@ -144,7 +181,8 @@ class Rule(BaseModel):
                     # Create standardized metric representation needed by server API
                     metric_data = {
                         "score_type": "",
-                        "threshold": 0.0
+                        "threshold": 0.0,
+                        "name": ""
                     }
                     
                     # First try to use object's own serialization methods
@@ -182,6 +220,16 @@ class Rule(BaseModel):
                             # Use condition threshold if metric doesn't have one
                             metric_data['threshold'] = self.conditions[i].threshold
                     
+                    # Make sure name is set
+                    if not metric_data.get('name'):
+                        if hasattr(metric_obj, '__name__'):
+                            metric_data['name'] = metric_obj.__name__
+                        elif hasattr(metric_obj, 'name'):
+                            metric_data['name'] = metric_obj.name
+                        else:
+                            # Fallback to score_type if available
+                            metric_data['name'] = metric_data.get('score_type', str(metric_obj))
+                    
                     # Update the condition with our properly serialized metric
                     condition["metric"] = metric_data
         
@@ -199,47 +247,6 @@ class Rule(BaseModel):
             raise ValueError(f"combine_type must be 'all' or 'any', got: {v}")
         return v
 
-class AlertResult(BaseModel):
-    """
-    Result of evaluating a rule.
-    
-    Example:
-        {
-            "status": "triggered",
-            "rule_name": "Quality Check",
-            "conditions_result": [
-                {"metric": "faithfulness", "value": 0.6, "threshold": 0.7, "passed": False},
-                {"metric": "relevancy", "value": 0.9, "threshold": 0.8, "passed": True}
-            ],
-            "rule_id": "123e4567-e89b-12d3-a456-426614174000",
-            "metadata": {
-                "example_id": "example_123",
-                "timestamp": "20240321_123456"
-            },
-            "notification": {
-                "enabled": true,
-                "communication_methods": ["slack", "email"],
-                "email_addresses": ["user1@example.com", "user2@example.com"]
-            }
-        }
-    """
-    status: AlertStatus
-    rule_id: Optional[str] = None  # The unique identifier of the rule
-    rule_name: str
-    conditions_result: List[Dict[str, Any]]
-    metadata: Dict[str, Any] = {}
-    notification: Optional[NotificationConfig] = None  # Configuration for notifications
-    
-    @property
-    def example_id(self) -> Optional[str]:
-        """Get example_id from metadata for backward compatibility"""
-        return self.metadata.get("example_id")
-        
-    @property
-    def timestamp(self) -> Optional[str]:
-        """Get timestamp from metadata for backward compatibility"""
-        return self.metadata.get("timestamp")
-
 class RulesEngine:
     """
     Engine for creating and evaluating rules against metrics.
@@ -406,7 +413,7 @@ class RulesEngine:
                 # If rule has a notification config and the alert is triggered, include it in the result
                 notification_config = rule.notification
             
-            # Set the alert status based on whether the rule was triggered
+            # Set the alert status based on whether the rule was triggered using proper enum values
             status = AlertStatus.TRIGGERED if triggered else AlertStatus.NOT_TRIGGERED
             
             # Create the alert result
@@ -416,7 +423,10 @@ class RulesEngine:
                 rule_name=rule.name,
                 conditions_result=condition_results,
                 notification=notification_config,
-                metadata=example_metadata or {}
+                metadata=example_metadata or {},
+                combine_type=rule.combine_type,
+                project_id=example_metadata.get("project_id") if example_metadata else None,
+                trace_span_id=example_metadata.get("trace_span_id") if example_metadata else None
             )
             
             results[rule_id] = alert_result