judgeval 0.0.55__py3-none-any.whl → 0.2.0__py3-none-any.whl

judgeval/common/tracer/otel_exporter.py

@@ -0,0 +1,108 @@
+"""
+Custom OpenTelemetry exporter for Judgment API.
+
+This exporter sends spans to the Judgment API using the existing format.
+The BatchSpanProcessor handles all batching, threading, and retry logic.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Sequence
+
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from opentelemetry.sdk.trace import ReadableSpan
+
+from judgeval.common.tracer.span_transformer import SpanTransformer
+from judgeval.common.logger import judgeval_logger
+from judgeval.common.api.api import JudgmentApiClient
+
+
+class JudgmentAPISpanExporter(SpanExporter):
+    """
+    Custom OpenTelemetry exporter that sends spans to Judgment API.
+
+    This exporter is used by BatchSpanProcessor which handles all the
+    batching, threading, and retry logic for us.
+    """
+
+    def __init__(
+        self,
+        judgment_api_key: str,
+        organization_id: str,
+    ):
+        self.api_client = JudgmentApiClient(judgment_api_key, organization_id)
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        """
+        Export spans to Judgment API.
+
+        This method is called by BatchSpanProcessor with a batch of spans.
+        We send them synchronously since BatchSpanProcessor handles threading.
+        """
+        if not spans:
+            return SpanExportResult.SUCCESS
+
+        try:
+            spans_data = []
+            eval_runs_data = []
+
+            for span in spans:
+                span_data = self._convert_span_to_judgment_format(span)
+
+                if span.attributes and span.attributes.get("judgment.evaluation_run"):
+                    eval_runs_data.append(span_data)
+                else:
+                    spans_data.append(span_data)
+
+            if spans_data:
+                self._send_spans_batch(spans_data)
+
+            if eval_runs_data:
+                self._send_evaluation_runs_batch(eval_runs_data)
+
+            return SpanExportResult.SUCCESS
+
+        except Exception as e:
+            judgeval_logger.error(f"Error in JudgmentAPISpanExporter.export: {e}")
+            return SpanExportResult.FAILURE
+
+    def _convert_span_to_judgment_format(self, span: ReadableSpan) -> Dict[str, Any]:
+        """Convert OpenTelemetry span to existing Judgment API format."""
+        if span.attributes and span.attributes.get("judgment.evaluation_run"):
+            return SpanTransformer.otel_span_to_evaluation_run_format(span)
+        else:
+            return SpanTransformer.otel_span_to_judgment_format(span)
+
+    def _send_spans_batch(self, spans: List[Dict[str, Any]]):
+        """Send a batch of spans to the spans endpoint."""
+        spans_data = [span["data"] for span in spans]
+        self.api_client.send_spans_batch(spans_data)
+
+    def _send_evaluation_runs_batch(self, eval_runs: List[Dict[str, Any]]):
+        """Send a batch of evaluation runs to the evaluation runs endpoint."""
+        evaluation_entries = []
+        for eval_run in eval_runs:
+            eval_data = eval_run["data"]
+            entry = {
+                "evaluation_run": {
+                    key: value
+                    for key, value in eval_data.items()
+                    if key not in ["associated_span_id", "span_data", "queued_at"]
+                },
+                "associated_span": {
+                    "span_id": eval_data.get("associated_span_id"),
+                    "span_data": eval_data.get("span_data"),
+                },
+                "queued_at": eval_data.get("queued_at"),
+            }
+            evaluation_entries.append(entry)
+
+        self.api_client.send_evaluation_runs_batch(evaluation_entries)
+
+    def shutdown(self, timeout_millis: int = 30000) -> None:
+        """Shutdown the exporter."""
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush any pending requests."""
+        return True

judgeval/common/tracer/otel_span_processor.py

@@ -0,0 +1,234 @@
+"""
+Custom OpenTelemetry span processor for Judgment API.
+
+This processor uses BatchSpanProcessor to handle batching and export
+of TraceSpan objects converted to OpenTelemetry format.
+"""
+
+from __future__ import annotations
+
+import threading
+from typing import Any, Dict, Optional
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanProcessor
+from opentelemetry.trace import Span, Status, StatusCode, SpanContext, TraceFlags
+from opentelemetry.trace.span import TraceState, INVALID_SPAN_CONTEXT
+from opentelemetry.util.types import Attributes
+
+from judgeval.common.logger import judgeval_logger
+from judgeval.common.tracer.otel_exporter import JudgmentAPISpanExporter
+from judgeval.common.tracer.span_processor import SpanProcessorBase
+from judgeval.common.tracer.span_transformer import SpanTransformer
+from judgeval.data import TraceSpan
+from judgeval.evaluation_run import EvaluationRun
+
+
+class SimpleReadableSpan(ReadableSpan):
+    """Simple ReadableSpan implementation that wraps TraceSpan data."""
+
+    def __init__(self, trace_span: TraceSpan, span_state: str = "completed"):
+        self._name = trace_span.function
+        self._span_id = trace_span.span_id
+        self._trace_id = trace_span.trace_id
+
+        self._start_time = (
+            int(trace_span.created_at * 1_000_000_000)
+            if trace_span.created_at
+            else None
+        )
+        self._end_time: Optional[int] = None
+
+        if (
+            span_state == "completed"
+            and trace_span.duration is not None
+            and self._start_time is not None
+        ):
+            self._end_time = self._start_time + int(trace_span.duration * 1_000_000_000)
+
+        self._status = (
+            Status(StatusCode.ERROR) if trace_span.error else Status(StatusCode.OK)
+        )
+
+        self._attributes = SpanTransformer.trace_span_to_otel_attributes(
+            trace_span, span_state
+        )
+
+        try:
+            trace_id_int = (
+                int(trace_span.trace_id.replace("-", ""), 16)
+                if trace_span.trace_id
+                else 0
+            )
+            span_id_int = (
+                int(trace_span.span_id.replace("-", ""), 16)
+                if trace_span.span_id
+                else 0
+            )
+
+            self._context = SpanContext(
+                trace_id=trace_id_int,
+                span_id=span_id_int,
+                is_remote=False,
+                trace_flags=TraceFlags(0x01),
+                trace_state=TraceState(),
+            )
+        except (ValueError, TypeError) as e:
+            judgeval_logger.warning(f"Failed to create proper SpanContext: {e}")
+            self._context = INVALID_SPAN_CONTEXT
+
+        self._parent: Optional[SpanContext] = None
+        self._events: list[Any] = []
+        self._links: list[Any] = []
+        self._resource: Optional[Any] = None
+        self._instrumentation_info: Optional[Any] = None
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def context(self) -> SpanContext:
+        return self._context
+
+    @property
+    def parent(self) -> Optional[SpanContext]:
+        return self._parent
+
+    @property
+    def start_time(self) -> Optional[int]:
+        return self._start_time
+
+    @property
+    def end_time(self) -> Optional[int]:
+        return self._end_time
+
+    @property
+    def status(self) -> Status:
+        return self._status
+
+    @property
+    def attributes(self) -> Optional[Attributes]:
+        return self._attributes
+
+    @property
+    def events(self):
+        return self._events
+
+    @property
+    def links(self):
+        return self._links
+
+    @property
+    def resource(self) -> Optional[Any]:
+        return self._resource
+
+    @property
+    def instrumentation_info(self) -> Optional[Any]:
+        return self._instrumentation_info
+
+
+class JudgmentSpanProcessor(SpanProcessor, SpanProcessorBase):
+    """
+    Span processor that converts TraceSpan objects to OpenTelemetry format
+    and uses BatchSpanProcessor for export.
+    """
+
+    def __init__(
+        self,
+        judgment_api_key: str,
+        organization_id: str,
+        batch_size: int = 50,
+        flush_interval: float = 1.0,
+        max_queue_size: int = 2048,
+        export_timeout: int = 30000,
+    ):
+        self.judgment_api_key = judgment_api_key
+        self.organization_id = organization_id
+
+        self._span_cache: Dict[str, TraceSpan] = {}
+        self._span_states: Dict[str, str] = {}
+        self._cache_lock = threading.RLock()
+
+        self.batch_processor = BatchSpanProcessor(
+            JudgmentAPISpanExporter(
+                judgment_api_key=judgment_api_key,
+                organization_id=organization_id,
+            ),
+            max_queue_size=max_queue_size,
+            schedule_delay_millis=int(flush_interval * 1000),
+            max_export_batch_size=batch_size,
+            export_timeout_millis=export_timeout,
+        )
+
+    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+        self.batch_processor.on_start(span, parent_context)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        self.batch_processor.on_end(span)
+
+    def queue_span_update(self, span: TraceSpan, span_state: str = "input") -> None:
+        if span_state == "completed":
+            span.set_update_id_to_ending_number()
+        else:
+            span.increment_update_id()
+
+        with self._cache_lock:
+            span_id = span.span_id
+
+            self._span_cache[span_id] = span
+            self._span_states[span_id] = span_state
+
+            self._send_span_update(span, span_state)
+
+            if span_state == "completed" or span_state == "error":
+                self._span_cache.pop(span_id, None)
+                self._span_states.pop(span_id, None)
+
+    def _send_span_update(self, span: TraceSpan, span_state: str) -> None:
+        readable_span = SimpleReadableSpan(span, span_state)
+        self.batch_processor.on_end(readable_span)
+
+    def flush_pending_spans(self) -> None:
+        with self._cache_lock:
+            if not self._span_cache:
+                return
+
+            for span_id, span in self._span_cache.items():
+                span_state = self._span_states.get(span_id, "input")
+                self._send_span_update(span, span_state)
+
+    def queue_evaluation_run(
+        self, evaluation_run: EvaluationRun, span_id: str, span_data: TraceSpan
+    ) -> None:
+        attributes = SpanTransformer.evaluation_run_to_otel_attributes(
+            evaluation_run, span_id, span_data
+        )
+
+        readable_span = SimpleReadableSpan(span_data, "evaluation_run")
+        readable_span._attributes.update(attributes)
+
+        self.batch_processor.on_end(readable_span)
+
+    def shutdown(self) -> None:
+        try:
+            self.flush_pending_spans()
+        except Exception as e:
+            judgeval_logger.warning(
+                f"Error flushing pending spans during shutdown: {e}"
+            )
+
+        self.batch_processor.shutdown()
+
+        with self._cache_lock:
+            self._span_cache.clear()
+            self._span_states.clear()
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        try:
+            self.flush_pending_spans()
+        except Exception as e:
+            judgeval_logger.warning(f"Error flushing pending spans: {e}")
+
+        return self.batch_processor.force_flush(timeout_millis)

judgeval/common/tracer/span_processor.py

@@ -0,0 +1,37 @@
+"""
+Base class for span processors with default no-op implementations.
+
+This eliminates the need for optional typing and null checks.
+When monitoring is disabled, we use this base class directly.
+When monitoring is enabled, we use JudgmentSpanProcessor which overrides the methods.
+"""
+
+from judgeval.data import TraceSpan
+from judgeval.evaluation_run import EvaluationRun
+
+
+class SpanProcessorBase:
+    """
+    Base class for Judgment span processors with default no-op implementations.
+
+    This eliminates the need for optional typing and null checks.
+    When monitoring is disabled, we use this base class directly.
+    When monitoring is enabled, we use JudgmentSpanProcessor which overrides the methods.
+    """
+
+    def queue_span_update(self, span: TraceSpan, span_state: str = "input") -> None:
+        pass
+
+    def queue_evaluation_run(
+        self, evaluation_run: EvaluationRun, span_id: str, span_data: TraceSpan
+    ) -> None:
+        pass
+
+    def flush_pending_spans(self) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+    def shutdown(self) -> None:
+        pass

judgeval/common/tracer/span_transformer.py

@@ -0,0 +1,211 @@
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from datetime import datetime, timezone
+from typing import Any, Dict, Optional, Union
+
+from opentelemetry.sdk.trace import ReadableSpan
+from pydantic import BaseModel
+
+from judgeval.data import TraceSpan
+from judgeval.evaluation_run import EvaluationRun
+
+
+class SpanTransformer:
+    @staticmethod
+    def _needs_json_serialization(value: Any) -> bool:
+        if value is None:
+            return False
+
+        simple_types = (str, int, float, bool)
+        if isinstance(value, simple_types):
+            return False
+
+        complex_types = (dict, list, tuple, set, BaseModel)
+        if isinstance(value, complex_types):
+            return True
+
+        try:
+            json.dumps(value)
+            return False
+        except (TypeError, ValueError):
+            return True
+
+    @staticmethod
+    def _safe_json_handle(obj: Any, serialize: bool = True) -> Any:
+        if serialize:
+            if obj is None:
+                return None
+            try:
+                return json.dumps(obj, default=str)
+            except Exception:
+                return json.dumps(str(obj))
+        else:
+            if not isinstance(obj, str):
+                return obj
+            try:
+                return json.loads(obj)
+            except (json.JSONDecodeError, TypeError):
+                return obj
+
+    @staticmethod
+    def _format_timestamp(timestamp: Optional[Union[float, int, str]]) -> str:
+        if timestamp is None:
+            return datetime.now(timezone.utc).isoformat()
+
+        if isinstance(timestamp, str):
+            return timestamp
+
+        try:
+            dt = datetime.fromtimestamp(timestamp, tz=timezone.utc)
+            return dt.isoformat()
+        except (ValueError, OSError):
+            return datetime.now(timezone.utc).isoformat()
+
+    @staticmethod
+    def trace_span_to_otel_attributes(
+        trace_span: TraceSpan, span_state: str = "completed"
+    ) -> Dict[str, Any]:
+        serialized_data = trace_span.model_dump()
+        attributes: Dict[str, Any] = {}
+
+        for field_name, value in serialized_data.items():
+            if value is None:
+                continue
+
+            attr_name = f"judgment.{field_name}"
+
+            if field_name == "created_at":
+                attributes[attr_name] = SpanTransformer._format_timestamp(value)
+            elif field_name == "expected_tools" and value:
+                attributes[attr_name] = SpanTransformer._safe_json_handle(
+                    [tool.model_dump() for tool in trace_span.expected_tools]
+                )
+            elif field_name == "usage" and value:
+                attributes[attr_name] = SpanTransformer._safe_json_handle(
+                    trace_span.usage.model_dump()
+                )
+            elif SpanTransformer._needs_json_serialization(value):
+                attributes[attr_name] = SpanTransformer._safe_json_handle(value)
+            else:
+                attributes[attr_name] = value
+
+        attributes["judgment.span_state"] = span_state
+        if not attributes.get("judgment.span_type"):
+            attributes["judgment.span_type"] = "span"
+
+        return attributes
+
+    @staticmethod
+    def otel_attributes_to_judgment_data(attributes: Dict[str, Any]) -> Dict[str, Any]:
+        judgment_data: Dict[str, Any] = {}
+
+        for key, value in attributes.items():
+            if not key.startswith("judgment."):
+                continue
+
+            field_name = key[9:]
+
+            if isinstance(value, str):
+                deserialized = SpanTransformer._safe_json_handle(value, serialize=False)
+                judgment_data[field_name] = deserialized
+            else:
+                judgment_data[field_name] = value
+
+        return judgment_data
+
+    @staticmethod
+    def otel_span_to_judgment_format(span: ReadableSpan) -> Dict[str, Any]:
+        attributes = span.attributes or {}
+        judgment_data = SpanTransformer.otel_attributes_to_judgment_data(attributes)
+
+        duration = judgment_data.get("duration")
+        if duration is None and span.end_time and span.start_time:
+            duration = (span.end_time - span.start_time) / 1_000_000_000
+
+        span_id = judgment_data.get("span_id") or str(uuid.uuid4())
+        trace_id = judgment_data.get("trace_id") or str(uuid.uuid4())
+
+        created_at = judgment_data.get("created_at")
+        if not created_at:
+            created_at = (
+                span.start_time / 1_000_000_000 if span.start_time else time.time()
+            )
+
+        return {
+            "type": "span",
+            "data": {
+                "span_id": span_id,
+                "trace_id": trace_id,
+                "function": span.name,
+                "depth": judgment_data.get("depth", 0),
+                "created_at": SpanTransformer._format_timestamp(created_at),
+                "parent_span_id": judgment_data.get("parent_span_id"),
+                "span_type": judgment_data.get("span_type", "span"),
+                "inputs": judgment_data.get("inputs"),
+                "error": judgment_data.get("error"),
+                "output": judgment_data.get("output"),
+                "usage": judgment_data.get("usage"),
+                "duration": duration,
+                "expected_tools": judgment_data.get("expected_tools"),
+                "additional_metadata": judgment_data.get("additional_metadata"),
+                "has_evaluation": judgment_data.get("has_evaluation", False),
+                "agent_name": judgment_data.get("agent_name"),
+                "state_before": judgment_data.get("state_before"),
+                "state_after": judgment_data.get("state_after"),
+                "update_id": judgment_data.get("update_id", 1),
+                "span_state": judgment_data.get("span_state", "completed"),
+                "queued_at": time.time(),
+            },
+        }
+
+    @staticmethod
+    def evaluation_run_to_otel_attributes(
+        evaluation_run: EvaluationRun, span_id: str, span_data: TraceSpan
+    ) -> Dict[str, Any]:
+        attributes = {
+            "judgment.evaluation_run": True,
+            "judgment.associated_span_id": span_id,
+            "judgment.span_data": SpanTransformer._safe_json_handle(
+                span_data.model_dump()
+            ),
+        }
+
+        eval_data = evaluation_run.model_dump()
+        for key, value in eval_data.items():
+            if value is None:
+                continue
+
+            attr_name = f"judgment.{key}"
+            if SpanTransformer._needs_json_serialization(value):
+                attributes[attr_name] = SpanTransformer._safe_json_handle(value)
+            else:
+                attributes[attr_name] = value
+
+        return attributes
+
+    @staticmethod
+    def otel_span_to_evaluation_run_format(span: ReadableSpan) -> Dict[str, Any]:
+        attributes = span.attributes or {}
+        judgment_data = SpanTransformer.otel_attributes_to_judgment_data(attributes)
+
+        associated_span_id = judgment_data.get("associated_span_id") or str(
+            uuid.uuid4()
+        )
+
+        eval_run_data = {
+            key: value
+            for key, value in judgment_data.items()
+            if key not in ["associated_span_id", "span_data", "evaluation_run"]
+        }
+
+        eval_run_data["associated_span_id"] = associated_span_id
+        eval_run_data["span_data"] = judgment_data.get("span_data")
+        eval_run_data["queued_at"] = time.time()
+
+        return {
+            "type": "evaluation_run",
+            "data": eval_run_data,
+        }

judgeval/common/tracer/trace_manager.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+from typing import List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from judgeval.common.tracer import Tracer
+
+from judgeval.common.logger import judgeval_logger
+from judgeval.common.api import JudgmentApiClient
+from rich import print as rprint
+
+
+class TraceManagerClient:
+    """
+    Client for handling trace endpoints with the Judgment API
+
+
+    Operations include:
+    - Fetching a trace by id
+    - Saving a trace
+    - Deleting a trace
+    """
+
+    def __init__(
+        self,
+        judgment_api_key: str,
+        organization_id: str,
+        tracer: Optional[Tracer] = None,
+    ):
+        self.api_client = JudgmentApiClient(judgment_api_key, organization_id)
+        self.tracer = tracer
+
+    def fetch_trace(self, trace_id: str):
+        """
+        Fetch a trace by its id
+        """
+        return self.api_client.fetch_trace(trace_id)
+
+    def upsert_trace(
+        self,
+        trace_data: dict,
+        offline_mode: bool = False,
+        show_link: bool = True,
+        final_save: bool = True,
+    ):
+        """
+        Upserts a trace to the Judgment API (always overwrites if exists).
+
+        Args:
+            trace_data: The trace data to upsert
+            offline_mode: Whether running in offline mode
+            show_link: Whether to show the UI link (for live tracing)
+            final_save: Whether this is the final save (controls S3 saving)
+
+        Returns:
+            dict: Server response containing UI URL and other metadata
+        """
+        server_response = self.api_client.upsert_trace(trace_data)
+
+        if self.tracer and self.tracer.use_s3 and final_save:
+            try:
+                s3_key = self.tracer.s3_storage.save_trace(
+                    trace_data=trace_data,
+                    trace_id=trace_data["trace_id"],
+                    project_name=trace_data["project_name"],
+                )
+                judgeval_logger.info(f"Trace also saved to S3 at key: {s3_key}")
+            except Exception as e:
+                judgeval_logger.warning(f"Failed to save trace to S3: {str(e)}")
+
+        if not offline_mode and show_link and "ui_results_url" in server_response:
+            pretty_str = f"\n🔍 You can view your trace data here: [rgb(106,0,255)][link={server_response['ui_results_url']}]View Trace[/link]\n"
+            rprint(pretty_str)
+
+        return server_response
+
+    def delete_trace(self, trace_id: str):
+        """
+        Delete a trace from the database.
+        """
+        return self.api_client.delete_trace(trace_id)
+
+    def delete_traces(self, trace_ids: List[str]):
+        """
+        Delete a batch of traces from the database.
+        """
+        return self.api_client.delete_traces(trace_ids)
+
+    def delete_project(self, project_name: str):
+        """
+        Deletes a project from the server. Which also deletes all evaluations and traces associated with the project.
+        """
+        return self.api_client.delete_project(project_name)