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

judgeval/tracer/constants.py

@@ -0,0 +1 @@
+JUDGEVAL_TRACER_INSTRUMENTING_MODULE_NAME = "opentelemetry.instrumentation.judgeval"

judgeval/tracer/exporters/__init__.py

@@ -0,0 +1,37 @@
+from opentelemetry.sdk.trace.export import (
+    SpanExportResult,
+    SpanExporter,
+)
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from typing import Sequence
+
+from judgeval.tracer.exporters.store import ABCSpanStore
+from judgeval.tracer.exporters.s3 import S3Exporter
+from judgeval.tracer.exporters.utils import deduplicate_spans
+
+
+class JudgmentSpanExporter(OTLPSpanExporter):
+    def __init__(self, endpoint: str, api_key: str, organization_id: str):
+        super().__init__(
+            endpoint=endpoint,
+            headers={
+                "Authorization": f"Bearer {api_key}",
+                "X-Organization-Id": organization_id,
+            },
+        )
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        return super().export(deduplicate_spans(spans))
+
+
+class InMemorySpanExporter(SpanExporter):
+    def __init__(self, store: ABCSpanStore):
+        self.store = store
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        self.store.add(*spans)
+        return SpanExportResult.SUCCESS
+
+
+__all__ = ("JudgmentSpanExporter", "InMemorySpanExporter", "S3Exporter")

judgeval/tracer/exporters/s3.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime
+from typing import Literal, Sequence, Optional, TYPE_CHECKING, cast
+import boto3
+from botocore.client import Config
+
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from opentelemetry.sdk.trace import ReadableSpan
+from judgeval.env import (
+    JUDGMENT_S3_ACCESS_KEY_ID,
+    JUDGMENT_S3_SECRET_ACCESS_KEY,
+    JUDGMENT_S3_REGION_NAME,
+    JUDGMENT_S3_BUCKET_NAME,
+    JUDGMENT_S3_PREFIX,
+    JUDGMENT_S3_ENDPOINT_URL,
+    JUDGMENT_S3_SIGNATURE_VERSION,
+    JUDGMENT_S3_ADDRESSING_STYLE,
+)
+from judgeval.exceptions import JudgmentRuntimeError
+from judgeval.logger import judgeval_logger
+
+if TYPE_CHECKING:
+    from mypy_boto3_s3.client import S3Client
+
+
+class S3Exporter(SpanExporter):
+    __slots__ = ("bucket_name", "prefix", "s3_client")
+
+    bucket_name: str
+    prefix: str
+    s3_client: S3Client
+
+    def __init__(
+        self,
+        bucket_name: Optional[str] = JUDGMENT_S3_BUCKET_NAME,
+        region_name: Optional[str] = JUDGMENT_S3_REGION_NAME,
+        prefix: str = JUDGMENT_S3_PREFIX,
+        s3_access_key_id: Optional[str] = JUDGMENT_S3_ACCESS_KEY_ID,
+        s3_secret_access_key: Optional[str] = JUDGMENT_S3_SECRET_ACCESS_KEY,
+        endpoint_url: Optional[str] = JUDGMENT_S3_ENDPOINT_URL,
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/guide/configuration.html
+        signature_version: str = JUDGMENT_S3_SIGNATURE_VERSION,
+        addressing_style: str = JUDGMENT_S3_ADDRESSING_STYLE,
+        batch_size: int = 8,
+    ):
+        if not bucket_name:
+            raise JudgmentRuntimeError("JUDGMENT_S3_BUCKET_NAME is not set")
+
+        if not region_name:
+            raise JudgmentRuntimeError("JUDGMENT_S3_REGION_NAME is not set")
+
+        if addressing_style not in ["auto", "virtual", "path"]:
+            raise JudgmentRuntimeError(f"Invalid addressing style: {addressing_style}")
+        addressing_style = cast(Literal["auto", "virtual", "path"], addressing_style)
+
+        self.bucket_name = bucket_name
+        self.prefix = prefix.rstrip("/")
+        self.batch_size = batch_size
+
+        self.s3_client = boto3.client(
+            "s3",
+            config=Config(
+                signature_version=signature_version,
+                s3={"addressing_style": addressing_style},
+            ),
+            aws_access_key_id=s3_access_key_id,
+            aws_secret_access_key=s3_secret_access_key,
+            endpoint_url=endpoint_url,
+            region_name=region_name,
+        )
+
+    def _upload_span(self, span: ReadableSpan) -> tuple[bool, str]:
+        """Upload a single span to S3. Returns (success, key)."""
+        try:
+            span_context = span.get_span_context()
+            if not span_context:
+                return False, ""
+
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S%f")
+            trace_id = format(span_context.trace_id, "032x")
+            span_id = format(span_context.span_id, "016x")
+            key = f"{self.prefix}/{trace_id}/{span_id}/{timestamp}.json"
+
+            span_json = span.to_json(indent=0)
+
+            self.s3_client.put_object(
+                Bucket=self.bucket_name,
+                Key=key,
+                Body=span_json,
+                ContentType="application/json",
+            )
+            return True, key
+        except Exception as e:
+            judgeval_logger.error(
+                f"Error uploading span {span_context.span_id if span_context else 'unknown'}: {e}"
+            )
+            return False, ""
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        if not spans:
+            return SpanExportResult.SUCCESS
+
+        try:
+            with ThreadPoolExecutor(
+                max_workers=min(len(spans), self.batch_size)
+            ) as executor:
+                futures = [executor.submit(self._upload_span, span) for span in spans]
+
+                for future in as_completed(futures):
+                    success, key = future.result()
+                    if not success:
+                        return SpanExportResult.FAILURE
+                return SpanExportResult.SUCCESS
+
+        except Exception as e:
+            judgeval_logger.error(f"Error exporting spans to S3: {e}")
+            return SpanExportResult.FAILURE

judgeval/tracer/exporters/store.py

@@ -0,0 +1,43 @@
+from abc import ABC, abstractmethod
+from typing import List
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+
+class ABCSpanStore(ABC):
+    @abstractmethod
+    def add(self, *spans: ReadableSpan): ...
+
+    @abstractmethod
+    def get(self, id: str) -> ReadableSpan: ...
+
+    @abstractmethod
+    def get_all(self) -> List[ReadableSpan]: ...
+
+
+class SpanStore(ABCSpanStore):
+    __slots__ = ("spans",)
+
+    spans: List[ReadableSpan]
+
+    def __init__(self):
+        self.spans = []
+
+    def add(self, *spans: ReadableSpan):
+        self.spans.extend(spans)
+
+    def get(self, id: str) -> ReadableSpan:
+        for span in self.spans:
+            context = span.get_span_context()
+            if context is None:
+                continue
+            if context.span_id == id:
+                return span
+
+        raise ValueError(f"Span with id {id} not found")
+
+    def get_all(self) -> List[ReadableSpan]:
+        return self.spans
+
+    def __repr__(self) -> str:
+        return f"SpanStore(spans={self.spans})"

judgeval/tracer/exporters/utils.py

@@ -0,0 +1,32 @@
+from typing import Sequence
+from opentelemetry.sdk.trace import ReadableSpan
+
+from judgeval.tracer.keys import AttributeKeys
+
+
+def deduplicate_spans(spans: Sequence[ReadableSpan]) -> Sequence[ReadableSpan]:
+    spans_by_key: dict[tuple[int, int], ReadableSpan] = {}
+    for span in spans:
+        if span.attributes and span.context:
+            update_id = span.attributes.get(AttributeKeys.JUDGMENT_UPDATE_ID)
+
+            if not isinstance(update_id, int):
+                continue
+
+            key = (span.context.trace_id, span.context.span_id)
+            if key not in spans_by_key:
+                spans_by_key[key] = span
+            else:
+                existing_attrs = spans_by_key[key].attributes
+                existing_update_id = (
+                    existing_attrs.get(AttributeKeys.JUDGMENT_UPDATE_ID, 0)
+                    if existing_attrs
+                    else 0
+                )
+                if (
+                    isinstance(existing_update_id, (int, float))
+                    and update_id > existing_update_id
+                ):
+                    spans_by_key[key] = span
+
+    return list(spans_by_key.values())

judgeval/tracer/keys.py

@@ -0,0 +1,67 @@
+"""
+Identifiers used by Judgeval to store specific types of data in the spans.
+"""
+
+from opentelemetry.semconv.resource import ResourceAttributes
+from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
+from enum import Enum
+
+
+class AttributeKeys(str, Enum):
+    # General function tracing attributes (custom namespace)
+    JUDGMENT_SPAN_KIND = "judgment.span_kind"
+    JUDGMENT_INPUT = "judgment.input"
+    JUDGMENT_OUTPUT = "judgment.output"
+    JUDGMENT_OFFLINE_MODE = "judgment.offline_mode"
+    JUDGMENT_UPDATE_ID = "judgment.update_id"
+
+    # Custom tracking attributes
+    JUDGMENT_CUSTOMER_ID = "judgment.customer_id"
+
+    # Agent specific attributes (custom namespace)
+    JUDGMENT_AGENT_ID = "judgment.agent_id"
+    JUDGMENT_PARENT_AGENT_ID = "judgment.parent_agent_id"
+    JUDGMENT_AGENT_CLASS_NAME = "judgment.agent_class_name"
+    JUDGMENT_AGENT_INSTANCE_NAME = "judgment.agent_instance_name"
+    JUDGMENT_IS_AGENT_ENTRY_POINT = "judgment.is_agent_entry_point"
+    JUDGMENT_CUMULATIVE_LLM_COST = "judgment.cumulative_llm_cost"
+    JUDGMENT_STATE_BEFORE = "judgment.state_before"
+    JUDGMENT_STATE_AFTER = "judgment.state_after"
+
+    # Evaluation-specific attributes (custom namespace)
+    PENDING_TRACE_EVAL = "judgment.pending_trace_eval"
+
+    # GenAI-specific attributes (semantic conventions)
+    GEN_AI_PROMPT = gen_ai_attributes.GEN_AI_PROMPT
+    GEN_AI_COMPLETION = gen_ai_attributes.GEN_AI_COMPLETION
+    GEN_AI_REQUEST_MODEL = gen_ai_attributes.GEN_AI_REQUEST_MODEL
+    GEN_AI_RESPONSE_MODEL = gen_ai_attributes.GEN_AI_RESPONSE_MODEL
+    GEN_AI_SYSTEM = gen_ai_attributes.GEN_AI_SYSTEM
+    GEN_AI_USAGE_INPUT_TOKENS = gen_ai_attributes.GEN_AI_USAGE_INPUT_TOKENS
+    GEN_AI_USAGE_OUTPUT_TOKENS = gen_ai_attributes.GEN_AI_USAGE_OUTPUT_TOKENS
+    GEN_AI_USAGE_COMPLETION_TOKENS = gen_ai_attributes.GEN_AI_USAGE_COMPLETION_TOKENS
+    GEN_AI_REQUEST_TEMPERATURE = gen_ai_attributes.GEN_AI_REQUEST_TEMPERATURE
+    GEN_AI_REQUEST_MAX_TOKENS = gen_ai_attributes.GEN_AI_REQUEST_MAX_TOKENS
+    GEN_AI_RESPONSE_FINISH_REASONS = gen_ai_attributes.GEN_AI_RESPONSE_FINISH_REASONS
+
+    # GenAI-specific attributes (custom namespace)
+    GEN_AI_USAGE_TOTAL_COST = "gen_ai.usage.total_cost_usd"
+
+
+class InternalAttributeKeys(str, Enum):
+    """
+    Internal attribute keys used for temporary state management in span processors.
+    These are NOT exported and are used only for internal span lifecycle management.
+    """
+
+    # Span control attributes
+    DISABLE_PARTIAL_EMIT = "disable_partial_emit"
+    CANCELLED = "cancelled"
+
+
+class ResourceKeys(str, Enum):
+    SERVICE_NAME = ResourceAttributes.SERVICE_NAME
+    TELEMETRY_SDK_LANGUAGE = ResourceAttributes.TELEMETRY_SDK_LANGUAGE
+    TELEMETRY_SDK_NAME = ResourceAttributes.TELEMETRY_SDK_NAME
+    TELEMETRY_SDK_VERSION = ResourceAttributes.TELEMETRY_SDK_VERSION
+    JUDGMENT_PROJECT_ID = "judgment.project_id"