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

judgeval/v1/tracer/base_tracer.py

@@ -0,0 +1,520 @@
+from __future__ import annotations
+
+import datetime
+import functools
+import inspect
+import time
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, Optional, Tuple, TypeVar, overload
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace.export import SpanExporter
+from opentelemetry.trace import Span, SpanContext, Status, StatusCode
+
+from judgeval.logger import judgeval_logger
+from judgeval.utils.decorators.dont_throw import dont_throw
+from judgeval.v1.data.example import Example
+from judgeval.v1.instrumentation import wrap_provider
+from judgeval.v1.instrumentation.llm.providers import ApiClient
+from judgeval.v1.internal.api import JudgmentSyncClient
+from judgeval.v1.utils import resolve_project_id
+from judgeval.v1.internal.api.api_types import (
+    ExampleEvaluationRun,
+    TraceEvaluationRun,
+)
+from judgeval.v1.scorers.base_scorer import BaseScorer
+from judgeval.judgment_attribute_keys import AttributeKeys
+from judgeval.v1.scorers.custom_scorer.custom_scorer import CustomScorer
+from judgeval.v1.tracer.exporters.judgment_span_exporter import JudgmentSpanExporter
+from judgeval.v1.tracer.processors.judgment_span_processor import JudgmentSpanProcessor
+from uuid import uuid4
+from opentelemetry.context import attach, detach, get_value, set_value
+from judgeval.v1.tracer.processors._lifecycles import (
+    AGENT_ID_KEY,
+    PARENT_AGENT_ID_KEY,
+    CUSTOMER_ID_KEY,
+    AGENT_CLASS_NAME_KEY,
+    AGENT_INSTANCE_NAME_KEY,
+)
+
+C = TypeVar("C", bound=Callable[..., Any])
+
+
+class BaseTracer(ABC):
+    __slots__ = (
+        "project_name",
+        "enable_evaluation",
+        "api_client",
+        "serializer",
+        "project_id",
+    )
+
+    TRACER_NAME = "judgeval"
+
+    def __init__(
+        self,
+        project_name: str,
+        enable_evaluation: bool,
+        api_client: JudgmentSyncClient,
+        serializer: Callable[[Any], str],
+    ):
+        self.project_name = project_name
+        self.enable_evaluation = enable_evaluation
+        self.api_client = api_client
+        self.serializer = serializer
+        self.project_id = resolve_project_id(api_client, project_name)
+
+        if self.project_id is None:
+            judgeval_logger.error(
+                f"Failed to resolve project {project_name}, "
+                f"please create it first at https://app.judgmentlabs.ai/org/{self.api_client.organization_id}/projects. "
+                "Skipping Judgment export."
+            )
+
+    @abstractmethod
+    def initialize(self) -> None:
+        pass
+
+    @abstractmethod
+    def force_flush(self, timeout_millis: int) -> bool:
+        pass
+
+    @abstractmethod
+    def shutdown(self, timeout_millis: int) -> None:
+        pass
+
+    def get_span_exporter(self) -> SpanExporter:
+        if self.project_id is not None:
+            return JudgmentSpanExporter(
+                endpoint=self._build_endpoint(self.api_client.base_url),
+                api_key=self.api_client.api_key,
+                organization_id=self.api_client.organization_id,
+                project_id=self.project_id,
+            )
+        else:
+            judgeval_logger.error(
+                "Project not resolved; cannot create exporter, returning NoOpSpanExporter"
+            )
+            from judgeval.v1.tracer.exporters.noop_span_exporter import NoOpSpanExporter
+
+            return NoOpSpanExporter()
+
+    def get_span_processor(self) -> JudgmentSpanProcessor:
+        if self.project_id is not None:
+            return JudgmentSpanProcessor(
+                self,
+                self.get_span_exporter(),
+            )
+        else:
+            judgeval_logger.error(
+                "Project not resolved; cannot create processor, returning NoOpSpanProcessor"
+            )
+            from judgeval.v1.tracer.processors.noop_span_processor import (
+                NoOpJudgmentSpanProcessor,
+            )
+
+            return NoOpJudgmentSpanProcessor()
+
+    def get_tracer(self) -> trace.Tracer:
+        return trace.get_tracer(self.TRACER_NAME)
+
+    def set_span_kind(self, kind: str) -> None:
+        if kind is None:
+            return
+        current_span = trace.get_current_span()
+        if current_span is not None:
+            current_span.set_attribute(AttributeKeys.JUDGMENT_SPAN_KIND, kind)
+
+    @dont_throw
+    def set_attribute(self, key: str, value: Any) -> None:
+        if not self._is_valid_key(key):
+            return
+        if value is None:
+            return
+        current_span = trace.get_current_span()
+        if current_span is not None:
+            serialized_value = (
+                self.serializer(value)
+                if not isinstance(value, (str, int, float, bool))
+                else value
+            )
+            current_span.set_attribute(key, serialized_value)
+
+    def set_attributes(self, attributes: Dict[str, Any]) -> None:
+        if attributes is None:
+            return
+        for key, value in attributes.items():
+            self.set_attribute(key, value)
+
+    def set_customer_id(self, customer_id: str) -> None:
+        ctx = set_value(CUSTOMER_ID_KEY, customer_id)
+        attach(ctx)
+
+    def set_llm_span(self) -> None:
+        self.set_span_kind("llm")
+
+    def set_tool_span(self) -> None:
+        self.set_span_kind("tool")
+
+    def set_general_span(self) -> None:
+        self.set_span_kind("span")
+
+    def set_input(self, input_data: Any) -> None:
+        self.set_attribute(AttributeKeys.JUDGMENT_INPUT, input_data)
+
+    def set_output(self, output_data: Any) -> None:
+        self.set_attribute(AttributeKeys.JUDGMENT_OUTPUT, output_data)
+
+    def span(self, span_name: str, callable_func: Callable[[], Any]) -> Any:
+        tracer = self.get_tracer()
+        with tracer.start_as_current_span(span_name) as span:
+            try:
+                return callable_func()
+            except Exception as e:
+                span.set_status(trace.Status(trace.StatusCode.ERROR))
+                span.record_exception(e)
+                raise
+
+    @staticmethod
+    def start_span(span_name: str) -> Span:
+        tracer = trace.get_tracer(BaseTracer.TRACER_NAME)
+        return tracer.start_span(span_name)
+
+    @dont_throw
+    def async_evaluate(
+        self,
+        scorer: BaseScorer,
+        example: Example,
+    ) -> None:
+        if not self.enable_evaluation:
+            return
+
+        span_context = self._get_sampled_span_context()
+        if span_context is None:
+            return
+
+        trace_id = span_context.trace_id
+        span_id = span_context.span_id
+        trace_id_hex = format(trace_id, "032x")
+        span_id_hex = format(span_id, "016x")
+
+        self._log_evaluation_info(
+            "asyncEvaluate", trace_id_hex, span_id_hex, scorer.get_name()
+        )
+
+        evaluation_run = self._create_evaluation_run(
+            scorer, example, trace_id_hex, span_id_hex
+        )
+        self._enqueue_evaluation(evaluation_run)
+
+    @dont_throw
+    def async_trace_evaluate(
+        self,
+        scorer: BaseScorer,
+    ) -> None:
+        if not self.enable_evaluation:
+            return
+
+        current_span = self._get_sampled_span()
+        if current_span is None:
+            return
+
+        span_context = current_span.get_span_context()
+        trace_id = span_context.trace_id
+        span_id = span_context.span_id
+        trace_id_hex = format(trace_id, "032x")
+        span_id_hex = format(span_id, "016x")
+
+        self._log_evaluation_info(
+            "asyncTraceEvaluate", trace_id_hex, span_id_hex, scorer.get_name()
+        )
+
+        evaluation_run = self._create_trace_evaluation_run(
+            scorer, trace_id_hex, span_id_hex
+        )
+        try:
+            trace_eval_json = self.serializer(evaluation_run)
+            current_span.set_attribute(
+                AttributeKeys.JUDGMENT_PENDING_TRACE_EVAL, trace_eval_json
+            )
+        except Exception as e:
+            judgeval_logger.error(f"Failed to serialize trace evaluation: {e}")
+
+    def _build_endpoint(self, base_url: str) -> str:
+        return (
+            base_url + "otel/v1/traces"
+            if base_url.endswith("/")
+            else base_url + "/otel/v1/traces"
+        )
+
+    def _generate_run_id(self, prefix: str, span_id: Optional[str]) -> str:
+        return prefix + (
+            span_id if span_id is not None else str(int(time.time() * 1000))
+        )
+
+    def _create_evaluation_run(
+        self,
+        scorer: BaseScorer,
+        example: Example,
+        trace_id: str,
+        span_id: str,
+    ) -> ExampleEvaluationRun:
+        run_id = self._generate_run_id("async_evaluate_", span_id)
+
+        judgment_scorers = (
+            [] if isinstance(scorer, CustomScorer) else [scorer.get_scorer_config()]
+        )
+        custom_scorers = [scorer.to_dict()] if isinstance(scorer, CustomScorer) else []
+
+        return ExampleEvaluationRun(
+            project_name=self.project_name,
+            eval_name=run_id,
+            trace_id=trace_id,
+            trace_span_id=span_id,
+            examples=[example.to_dict()],
+            judgment_scorers=judgment_scorers,
+            custom_scorers=custom_scorers,
+            created_at=datetime.datetime.now(datetime.timezone.utc).isoformat(),
+        )
+
+    def _create_trace_evaluation_run(
+        self,
+        scorer: BaseScorer,
+        trace_id: str,
+        span_id: str,
+    ) -> TraceEvaluationRun:
+        eval_name = self._generate_run_id("async_trace_evaluate_", span_id)
+
+        judgment_scorers = (
+            [] if isinstance(scorer, CustomScorer) else [scorer.get_scorer_config()]
+        )
+        custom_scorers = [scorer.to_dict()] if isinstance(scorer, CustomScorer) else []
+
+        return TraceEvaluationRun(
+            project_name=self.project_name,
+            eval_name=eval_name,
+            trace_and_span_ids=[[trace_id, span_id]],
+            judgment_scorers=judgment_scorers,
+            custom_scorers=custom_scorers,
+            is_offline=False,
+            is_bucket_run=False,
+            created_at=datetime.datetime.now(datetime.timezone.utc).isoformat(),
+        )
+
+    def _enqueue_evaluation(self, evaluation_run: ExampleEvaluationRun) -> None:
+        try:
+            self.api_client.add_to_run_eval_queue_examples(evaluation_run)
+        except Exception as e:
+            judgeval_logger.error(f"Failed to enqueue evaluation run: {e}")
+
+    def _get_sampled_span_context(self) -> Optional[SpanContext]:
+        current_span = trace.get_current_span()
+        if current_span is None:
+            return None
+        span_context = current_span.get_span_context()
+        if not span_context.is_valid or not span_context.trace_flags.sampled:
+            return None
+        return span_context
+
+    def _get_sampled_span(self) -> Optional[Span]:
+        current_span = trace.get_current_span()
+        if current_span is None:
+            return None
+        span_context = current_span.get_span_context()
+        if not span_context.is_valid or not span_context.trace_flags.sampled:
+            return None
+        return current_span
+
+    def _log_evaluation_info(
+        self, method: str, trace_id: str, span_id: str, scorer_name: str
+    ) -> None:
+        judgeval_logger.info(
+            f"{method}: project={self.project_name}, traceId={trace_id}, spanId={span_id}, scorer={scorer_name}"
+        )
+
+    @staticmethod
+    def _is_valid_key(key: str) -> bool:
+        return key is not None and len(key) > 0
+
+    @overload
+    def observe(
+        self,
+        func: C,
+        span_type: Optional[str] = "span",
+        span_name: Optional[str] = None,
+    ) -> C: ...
+
+    @overload
+    def observe(
+        self,
+        func: None = None,
+        span_type: Optional[str] = "span",
+        span_name: Optional[str] = None,
+    ) -> Callable[[C], C]: ...
+
+    def observe(
+        self,
+        func: Optional[C] = None,
+        span_type: Optional[str] = "span",
+        span_name: Optional[str] = None,
+    ) -> C | Callable[[C], C]:
+        if func is None:
+            return lambda f: self.observe(f, span_type, span_name)  # type: ignore[return-value]
+
+        tracer = self.get_tracer()
+        name = span_name or func.__name__
+
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with tracer.start_as_current_span(name) as span:
+                    if span_type:
+                        span.set_attribute(AttributeKeys.JUDGMENT_SPAN_KIND, span_type)
+
+                    try:
+                        input_data = _format_inputs(func, args, kwargs)
+                        span.set_attribute(
+                            AttributeKeys.JUDGMENT_INPUT, self.serializer(input_data)
+                        )
+
+                        self.get_span_processor().emit_partial()
+
+                        result = await func(*args, **kwargs)
+
+                        span.set_attribute(
+                            AttributeKeys.JUDGMENT_OUTPUT, self.serializer(result)
+                        )
+                        return result
+                    except Exception as e:
+                        span.record_exception(e)
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                        raise
+
+            return async_wrapper  # type: ignore[return-value]
+        else:
+
+            @functools.wraps(func)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with tracer.start_as_current_span(name) as span:
+                    if span_type:
+                        span.set_attribute(AttributeKeys.JUDGMENT_SPAN_KIND, span_type)
+
+                    try:
+                        input_data = _format_inputs(func, args, kwargs)
+                        span.set_attribute(
+                            AttributeKeys.JUDGMENT_INPUT, self.serializer(input_data)
+                        )
+
+                        self.get_span_processor().emit_partial()
+
+                        result = func(*args, **kwargs)
+
+                        span.set_attribute(
+                            AttributeKeys.JUDGMENT_OUTPUT, self.serializer(result)
+                        )
+                        return result
+                    except Exception as e:
+                        span.record_exception(e)
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                        raise
+
+            return sync_wrapper  # type: ignore[return-value]
+
+    @overload
+    def agent(self, func: C, /, *, identifier: Optional[str] = None) -> C: ...
+
+    @overload
+    def agent(
+        self, func: None = None, /, *, identifier: Optional[str] = None
+    ) -> Callable[[C], C]: ...
+
+    def agent(
+        self, func: Optional[C] = None, /, *, identifier: Optional[str] = None
+    ) -> C | Callable[[C], C]:
+        if func is None:
+            return lambda f: self.agent(f, identifier=identifier)  # type: ignore[return-value]
+
+        class_name = None
+        if hasattr(func, "__qualname__") and "." in func.__qualname__:
+            parts = func.__qualname__.split(".")
+            if len(parts) >= 2:
+                class_name = parts[-2]
+
+        if inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                agent_id = str(uuid4())
+                parent_agent_id = get_value(AGENT_ID_KEY)
+                ctx = set_value(AGENT_ID_KEY, agent_id)
+                if parent_agent_id:
+                    ctx = set_value(PARENT_AGENT_ID_KEY, parent_agent_id, context=ctx)
+                if class_name:
+                    ctx = set_value(AGENT_CLASS_NAME_KEY, class_name, context=ctx)
+                if identifier and args:
+                    instance = args[0]
+                    if hasattr(instance, identifier):
+                        instance_name = str(getattr(instance, identifier))
+                        ctx = set_value(
+                            AGENT_INSTANCE_NAME_KEY, instance_name, context=ctx
+                        )
+                token = attach(ctx)
+                try:
+                    return await func(*args, **kwargs)
+                finally:
+                    detach(token)
+
+            return async_wrapper  # type: ignore[return-value]
+        else:
+
+            @functools.wraps(func)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                agent_id = str(uuid4())
+                parent_agent_id = get_value(AGENT_ID_KEY)
+                ctx = set_value(AGENT_ID_KEY, agent_id)
+                if parent_agent_id:
+                    ctx = set_value(PARENT_AGENT_ID_KEY, parent_agent_id, context=ctx)
+                if class_name:
+                    ctx = set_value(AGENT_CLASS_NAME_KEY, class_name, context=ctx)
+                if identifier and args:
+                    instance = args[0]
+                    if hasattr(instance, identifier):
+                        instance_name = str(getattr(instance, identifier))
+                        ctx = set_value(
+                            AGENT_INSTANCE_NAME_KEY, instance_name, context=ctx
+                        )
+                token = attach(ctx)
+                try:
+                    return func(*args, **kwargs)
+                finally:
+                    detach(token)
+
+            return sync_wrapper  # type: ignore[return-value]
+
+    def wrap(self, client: ApiClient) -> ApiClient:
+        return wrap_provider(self, client)
+
+
+def _format_inputs(
+    f: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any]
+) -> Dict[str, Any]:
+    try:
+        params = list(inspect.signature(f).parameters.values())
+        inputs: Dict[str, Any] = {}
+        arg_i = 0
+        for param in params:
+            if param.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD:
+                if arg_i < len(args):
+                    inputs[param.name] = args[arg_i]
+                    arg_i += 1
+                elif param.name in kwargs:
+                    inputs[param.name] = kwargs[param.name]
+            elif param.kind == inspect.Parameter.VAR_POSITIONAL:
+                inputs[param.name] = args[arg_i:]
+                arg_i = len(args)
+            elif param.kind == inspect.Parameter.VAR_KEYWORD:
+                inputs[param.name] = kwargs
+        return inputs
+    except Exception:
+        return {}

judgeval/v1/tracer/exporters/__init__.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from judgeval.v1.tracer.exporters.judgment_span_exporter import JudgmentSpanExporter
+from judgeval.v1.tracer.exporters.noop_span_exporter import NoOpSpanExporter
+from judgeval.v1.tracer.exporters.span_store import ABCSpanStore, SpanStore
+from judgeval.v1.tracer.exporters.in_memory_span_exporter import InMemorySpanExporter
+
+__all__ = [
+    "JudgmentSpanExporter",
+    "NoOpSpanExporter",
+    "ABCSpanStore",
+    "SpanStore",
+    "InMemorySpanExporter",
+]

judgeval/v1/tracer/exporters/in_memory_span_exporter.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+from judgeval.v1.tracer.exporters.span_store import ABCSpanStore
+
+
+class InMemorySpanExporter(SpanExporter):
+    __slots__ = ("_store",)
+
+    def __init__(self, store: ABCSpanStore) -> None:
+        self._store = store
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        self._store.add(*spans)
+        return SpanExportResult.SUCCESS
+
+    def shutdown(self) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True

judgeval/v1/tracer/exporters/judgment_span_exporter.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from typing import Sequence
+
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+from judgeval.logger import judgeval_logger
+
+
+class JudgmentSpanExporter(SpanExporter):
+    __slots__ = ("_delegate",)
+
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: str,
+        organization_id: str,
+        project_id: str,
+    ):
+        if not project_id or len(project_id.strip()) == 0:
+            raise ValueError("project_id is required for JudgmentSpanExporter")
+
+        self._delegate = OTLPSpanExporter(
+            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:
+        judgeval_logger.info(f"Exported {len(spans)} spans")
+        return self._delegate.export(spans)
+
+    def shutdown(self) -> None:
+        self._delegate.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return self._delegate.force_flush(timeout_millis)

judgeval/v1/tracer/exporters/noop_span_exporter.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from typing import Sequence
+
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+
+
+class NoOpSpanExporter(SpanExporter):
+    __slots__ = ()
+
+    def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult:
+        return SpanExportResult.SUCCESS
+
+    def shutdown(self) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True

judgeval/v1/tracer/exporters/span_store.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Dict, List
+
+from opentelemetry.sdk.trace import ReadableSpan
+
+
+class ABCSpanStore(ABC):
+    @abstractmethod
+    def add(self, *spans: ReadableSpan) -> None: ...
+
+    @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) -> None: ...
+
+
+class SpanStore(ABCSpanStore):
+    __slots__ = ("_spans_by_trace",)
+
+    def __init__(self) -> None:
+        self._spans_by_trace: Dict[str, List[ReadableSpan]] = {}
+
+    def add(self, *spans: ReadableSpan) -> None:
+        for span in spans:
+            context = span.get_span_context()
+            if context is None:
+                continue
+            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]:
+        return self._spans_by_trace.get(trace_id, [])
+
+    def clear_trace(self, trace_id: str) -> None:
+        if trace_id in self._spans_by_trace:
+            del self._spans_by_trace[trace_id]