judgeval 0.1.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

judgeval/tracer/constants.py

@@ -0,0 +1 @@
+JUDGEVAL_TRACER_INSTRUMENTING_MODULE_NAME = "judgeval"

judgeval/tracer/exporters/__init__.py

@@ -0,0 +1,40 @@
+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, project_id: str
+    ):
+        super().__init__(
+            endpoint=endpoint,
+            headers={
+                "Authorization": f"Bearer {api_key}",
+                "X-Organization-Id": organization_id,
+                "X-Project-Id": project_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,59 @@
+from abc import ABC, abstractmethod
+from typing import List, Dict
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+
+class ABCSpanStore(ABC):
+    @abstractmethod
+    def add(self, *spans: ReadableSpan): ...
+
+    @abstractmethod
+    def get_all(self) -> List[ReadableSpan]: ...
+
+    @abstractmethod
+    def get_by_trace_id(self, trace_id: str) -> List[ReadableSpan]: ...
+
+    @abstractmethod
+    def clear_trace(self, trace_id: str): ...
+
+
+class SpanStore(ABCSpanStore):
+    __slots__ = ("_spans_by_trace",)
+
+    _spans_by_trace: Dict[str, List[ReadableSpan]]
+
+    def __init__(self):
+        self._spans_by_trace = {}
+
+    def add(self, *spans: ReadableSpan):
+        for span in spans:
+            context = span.get_span_context()
+            if context is None:
+                continue
+            # Convert trace_id to hex string per OTEL spec
+            trace_id = format(context.trace_id, "032x")
+            if trace_id not in self._spans_by_trace:
+                self._spans_by_trace[trace_id] = []
+            self._spans_by_trace[trace_id].append(span)
+
+    def get_all(self) -> List[ReadableSpan]:
+        all_spans = []
+        for spans in self._spans_by_trace.values():
+            all_spans.extend(spans)
+        return all_spans
+
+    def get_by_trace_id(self, trace_id: str) -> List[ReadableSpan]:
+        """Get all spans for a specific trace ID (32-char hex string)."""
+        return self._spans_by_trace.get(trace_id, [])
+
+    def clear_trace(self, trace_id: str):
+        """Clear all spans for a specific trace ID (32-char hex string)."""
+        if trace_id in self._spans_by_trace:
+            del self._spans_by_trace[trace_id]
+
+    def __repr__(self) -> str:
+        total_spans = sum(len(spans) for spans in self._spans_by_trace.values())
+        return (
+            f"SpanStore(traces={len(self._spans_by_trace)}, total_spans={total_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,63 @@
+"""
+Identifiers used by Judgeval to store specific types of data in the spans.
+"""
+
+from enum import Enum
+
+
+class AttributeKeys(str, Enum):
+    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"
+
+    JUDGMENT_USAGE_METADATA = "judgment.usage.metadata"
+
+    JUDGMENT_CUSTOMER_ID = "judgment.customer_id"
+
+    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_STATE_BEFORE = "judgment.state_before"
+    JUDGMENT_STATE_AFTER = "judgment.state_after"
+
+    PENDING_TRACE_EVAL = "judgment.pending_trace_eval"
+
+    JUDGMENT_LLM_PROVIDER = "judgment.llm.provider"
+    JUDGMENT_LLM_MODEL_NAME = "judgment.llm.model"
+    JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS = "judgment.usage.non_cached_input_tokens"
+    JUDGMENT_USAGE_CACHE_CREATION_INPUT_TOKENS = (
+        "judgment.usage.cache_creation_input_tokens"
+    )
+    JUDGMENT_USAGE_CACHE_READ_INPUT_TOKENS = "judgment.usage.cache_read_input_tokens"
+    JUDGMENT_USAGE_OUTPUT_TOKENS = "judgment.usage.output_tokens"
+    JUDGMENT_USAGE_TOTAL_COST_USD = "judgment.usage.total_cost_usd"
+
+    GEN_AI_PROMPT = "gen_ai.prompt"
+    GEN_AI_COMPLETION = "gen_ai.completion"
+    GEN_AI_SYSTEM = "gen_ai.system"
+    GEN_AI_REQUEST_TEMPERATURE = "gen_ai.request.temperature"
+    GEN_AI_REQUEST_MAX_TOKENS = "gen_ai.request.max_tokens"
+    GEN_AI_RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons"
+
+
+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.
+    """
+
+    DISABLE_PARTIAL_EMIT = "disable_partial_emit"
+    CANCELLED = "cancelled"
+    IS_CUSTOMER_CONTEXT_OWNER = "is_customer_context_owner"
+
+
+class ResourceKeys(str, Enum):
+    SERVICE_NAME = "service.name"
+    TELEMETRY_SDK_LANGUAGE = "telemetry.sdk.language"
+    TELEMETRY_SDK_NAME = "telemetry.sdk.name"
+    TELEMETRY_SDK_VERSION = "telemetry.sdk.version"
+    JUDGMENT_PROJECT_ID = "judgment.project_id"

judgeval/tracer/llm/__init__.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+
+from .config import _detect_provider, wrap_provider
+
+
+__all__ = ["_detect_provider", "wrap_provider"]

judgeval/tracer/llm/config.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+from judgeval.logger import judgeval_logger
+
+from judgeval.tracer.llm.constants import ProviderType
+from judgeval.tracer.llm.providers import (
+    HAS_OPENAI,
+    HAS_TOGETHER,
+    HAS_ANTHROPIC,
+    HAS_GOOGLE_GENAI,
+    ApiClient,
+)
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+
+
+def _detect_provider(client: ApiClient) -> ProviderType:
+    if HAS_OPENAI:
+        from openai import OpenAI, AsyncOpenAI
+
+        if isinstance(client, (OpenAI, AsyncOpenAI)):
+            return ProviderType.OPENAI
+
+    if HAS_ANTHROPIC:
+        from anthropic import Anthropic, AsyncAnthropic
+
+        if isinstance(client, (Anthropic, AsyncAnthropic)):
+            return ProviderType.ANTHROPIC
+
+    if HAS_TOGETHER:
+        from together import Together, AsyncTogether  # type: ignore[import-untyped]
+
+        if isinstance(client, (Together, AsyncTogether)):
+            return ProviderType.TOGETHER
+
+    if HAS_GOOGLE_GENAI:
+        from google.genai import Client as GoogleClient
+
+        if isinstance(client, GoogleClient):
+            return ProviderType.GOOGLE
+
+    judgeval_logger.warning(
+        f"Unknown client type {type(client)}, Trying to wrap as OpenAI-compatible. "
+        "If this is a mistake or you think we should support this client, please file an issue at https://github.com/JudgmentLabs/judgeval/issues!"
+    )
+
+    return ProviderType.DEFAULT
+
+
+def wrap_provider(tracer: Tracer, client: ApiClient) -> ApiClient:
+    """
+    Wraps an API client to add tracing capabilities.
+    Supports OpenAI, Together, Anthropic, and Google GenAI clients.
+    """
+    provider_type = _detect_provider(client)
+
+    if provider_type == ProviderType.OPENAI:
+        from .llm_openai.wrapper import wrap_openai_client
+
+        return wrap_openai_client(tracer, client)
+    elif provider_type == ProviderType.ANTHROPIC:
+        from .llm_anthropic.wrapper import wrap_anthropic_client
+
+        return wrap_anthropic_client(tracer, client)
+    elif provider_type == ProviderType.TOGETHER:
+        from .llm_together.wrapper import wrap_together_client
+
+        return wrap_together_client(tracer, client)
+    elif provider_type == ProviderType.GOOGLE:
+        from .llm_google.wrapper import wrap_google_client
+
+        return wrap_google_client(tracer, client)
+    else:
+        # Default to OpenAI-compatible wrapping for unknown clients
+        from .llm_openai.wrapper import wrap_openai_client
+
+        return wrap_openai_client(tracer, client)

judgeval/tracer/llm/constants.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class ProviderType(Enum):
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+    TOGETHER = "together"
+    GOOGLE = "google"
+    DEFAULT = "default"

judgeval/tracer/llm/llm_anthropic/__init__.py

@@ -0,0 +1,3 @@
+from .wrapper import wrap_anthropic_client
+
+__all__ = ["wrap_anthropic_client"]

judgeval/tracer/llm/llm_anthropic/config.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+import importlib.util
+
+HAS_ANTHROPIC = importlib.util.find_spec("anthropic") is not None
+
+__all__ = ["HAS_ANTHROPIC"]