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

judgeval/v1/tracer/judgment_tracer_provider.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from typing import Callable, Optional
+
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.trace import Tracer, NoOpTracer
+from opentelemetry.util.types import Attributes
+
+from judgeval.logger import judgeval_logger
+from judgeval.v1.tracer.base_tracer import BaseTracer
+
+FilterTracerCallback = Callable[[str, Optional[str], Optional[str], Attributes], bool]
+
+
+class JudgmentTracerProvider(TracerProvider):
+    __slots__ = ("_filter_tracer",)
+
+    def __init__(
+        self,
+        filter_tracer: Optional[FilterTracerCallback] = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self._filter_tracer = (
+            filter_tracer if filter_tracer is not None else lambda *_: True
+        )
+
+    def get_tracer(
+        self,
+        instrumenting_module_name: str,
+        instrumenting_library_version: Optional[str] = None,
+        schema_url: Optional[str] = None,
+        attributes: Attributes = None,
+    ) -> Tracer:
+        if instrumenting_module_name == BaseTracer.TRACER_NAME:
+            return super().get_tracer(
+                instrumenting_module_name,
+                instrumenting_library_version,
+                schema_url,
+                attributes,
+            )
+
+        try:
+            if self._filter_tracer(
+                instrumenting_module_name,
+                instrumenting_library_version,
+                schema_url,
+                attributes,
+            ):
+                return super().get_tracer(
+                    instrumenting_module_name,
+                    instrumenting_library_version,
+                    schema_url,
+                    attributes,
+                )
+            else:
+                judgeval_logger.debug(
+                    f"[JudgmentTracerProvider] Returning NoOpTracer for tracer {instrumenting_module_name} as it is disallowed by the filterTracer callback."
+                )
+                return NoOpTracer()
+        except Exception as error:
+            judgeval_logger.error(
+                f"[JudgmentTracerProvider] Failed to filter tracer {instrumenting_module_name}: {error}."
+            )
+            return super().get_tracer(
+                instrumenting_module_name,
+                instrumenting_library_version,
+                schema_url,
+                attributes,
+            )

judgeval/v1/tracer/processors/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from judgeval.v1.tracer.processors.judgment_span_processor import JudgmentSpanProcessor
+from judgeval.v1.tracer.processors.noop_span_processor import NoOpJudgmentSpanProcessor
+
+__all__ = ["JudgmentSpanProcessor", "NoOpJudgmentSpanProcessor"]

judgeval/v1/tracer/processors/_lifecycles/__init__.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+from judgeval.v1.tracer.processors._lifecycles.customer_id_processor import (
+    CustomerIdProcessor,
+)
+from judgeval.v1.tracer.processors._lifecycles.agent_id_processor import (
+    AgentIdProcessor,
+)
+from judgeval.v1.tracer.processors._lifecycles.registry import get_all, register
+from judgeval.v1.tracer.processors._lifecycles.context_keys import (
+    CUSTOMER_ID_KEY,
+    AGENT_ID_KEY,
+    PARENT_AGENT_ID_KEY,
+    AGENT_CLASS_NAME_KEY,
+    AGENT_INSTANCE_NAME_KEY,
+)
+
+__all__ = [
+    "CustomerIdProcessor",
+    "AgentIdProcessor",
+    "get_all",
+    "register",
+    "CUSTOMER_ID_KEY",
+    "AGENT_ID_KEY",
+    "PARENT_AGENT_ID_KEY",
+    "AGENT_CLASS_NAME_KEY",
+    "AGENT_INSTANCE_NAME_KEY",
+]

judgeval/v1/tracer/processors/_lifecycles/agent_id_processor.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from opentelemetry.context import Context, get_value
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
+
+from judgeval.v1.tracer.processors._lifecycles.registry import register
+from judgeval.v1.tracer.processors._lifecycles.context_keys import (
+    AGENT_ID_KEY,
+    PARENT_AGENT_ID_KEY,
+    AGENT_CLASS_NAME_KEY,
+    AGENT_INSTANCE_NAME_KEY,
+)
+from judgeval.judgment_attribute_keys import AttributeKeys
+
+
+class AgentIdProcessor(SpanProcessor):
+    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+        agent_id = get_value(AGENT_ID_KEY, context=parent_context)
+        if agent_id is not None:
+            span.set_attribute(AttributeKeys.JUDGMENT_AGENT_ID, str(agent_id))
+
+        parent_agent_id = get_value(PARENT_AGENT_ID_KEY, context=parent_context)
+        if parent_agent_id is not None:
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_PARENT_AGENT_ID, str(parent_agent_id)
+            )
+
+        class_name = get_value(AGENT_CLASS_NAME_KEY, context=parent_context)
+        if class_name is not None:
+            span.set_attribute(AttributeKeys.JUDGMENT_AGENT_CLASS_NAME, str(class_name))
+
+        instance_name = get_value(AGENT_INSTANCE_NAME_KEY, context=parent_context)
+        if instance_name is not None:
+            span.set_attribute(
+                AttributeKeys.JUDGMENT_AGENT_INSTANCE_NAME, str(instance_name)
+            )
+
+        if agent_id is not None and agent_id != parent_agent_id:
+            span.set_attribute(AttributeKeys.JUDGMENT_IS_AGENT_ENTRY_POINT, True)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        pass
+
+    def shutdown(self) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+
+register(AgentIdProcessor)

judgeval/v1/tracer/processors/_lifecycles/context_keys.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from opentelemetry.context import create_key
+from judgeval.judgment_attribute_keys import AttributeKeys
+
+
+CUSTOMER_ID_KEY = create_key(AttributeKeys.JUDGMENT_CUSTOMER_ID)
+AGENT_ID_KEY = create_key(AttributeKeys.JUDGMENT_AGENT_ID)
+PARENT_AGENT_ID_KEY = create_key(AttributeKeys.JUDGMENT_PARENT_AGENT_ID)
+AGENT_CLASS_NAME_KEY = create_key(AttributeKeys.JUDGMENT_AGENT_CLASS_NAME)
+AGENT_INSTANCE_NAME_KEY = create_key(AttributeKeys.JUDGMENT_AGENT_INSTANCE_NAME)

judgeval/v1/tracer/processors/_lifecycles/customer_id_processor.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from opentelemetry.context import Context, get_value
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
+
+from judgeval.v1.tracer.processors._lifecycles.registry import register
+from judgeval.v1.tracer.processors._lifecycles.context_keys import CUSTOMER_ID_KEY
+from judgeval.judgment_attribute_keys import AttributeKeys
+
+
+class CustomerIdProcessor(SpanProcessor):
+    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+        customer_id = get_value(CUSTOMER_ID_KEY, context=parent_context)
+        if customer_id is not None:
+            span.set_attribute(AttributeKeys.JUDGMENT_CUSTOMER_ID, str(customer_id))
+
+    def on_end(self, span: ReadableSpan) -> None:
+        pass
+
+    def shutdown(self) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        return True
+
+
+register(CustomerIdProcessor)

judgeval/v1/tracer/processors/_lifecycles/registry.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import Callable, List
+
+from opentelemetry.sdk.trace import SpanProcessor
+
+
+ProcessorFactory = Callable[[], SpanProcessor]
+
+_lifecycle_processors: List[ProcessorFactory] = []
+
+
+def register(processor_class: ProcessorFactory) -> None:
+    _lifecycle_processors.append(processor_class)
+
+
+def get_all() -> List[SpanProcessor]:
+    return [factory() for factory in _lifecycle_processors]

judgeval/v1/tracer/processors/judgment_span_processor.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Optional
+
+from collections import defaultdict
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan, Span
+from opentelemetry.trace import get_current_span
+from opentelemetry.trace.span import SpanContext
+from opentelemetry.sdk.trace.export import (
+    BatchSpanProcessor,
+    SpanExporter,
+)
+
+from judgeval.judgment_attribute_keys import AttributeKeys
+from judgeval.tracer.keys import InternalAttributeKeys
+from judgeval.utils.decorators.dont_throw import dont_throw
+from judgeval.v1.tracer.processors._lifecycles import get_all
+
+
+if TYPE_CHECKING:
+    from judgeval.v1.tracer import BaseTracer
+
+
+class JudgmentSpanProcessor(BatchSpanProcessor):
+    __slots__ = ("tracer", "resource_attributes", "_internal_attributes")
+
+    def __init__(
+        self,
+        tracer: BaseTracer,
+        exporter: SpanExporter,
+        /,
+        *,
+        max_queue_size: int | None = None,
+        schedule_delay_millis: float | None = None,
+        max_export_batch_size: int | None = None,
+        export_timeout_millis: float | None = None,
+    ):
+        self.tracer = tracer
+
+        super().__init__(
+            exporter,
+            max_queue_size=max_queue_size,
+            schedule_delay_millis=schedule_delay_millis,
+            max_export_batch_size=max_export_batch_size,
+            export_timeout_millis=export_timeout_millis,
+        )
+        self._internal_attributes: defaultdict[tuple[int, int], dict[str, Any]] = (
+            defaultdict(dict)
+        )
+
+    def _get_span_key(self, span_context: SpanContext) -> tuple[int, int]:
+        return (span_context.trace_id, span_context.span_id)
+
+    def set_internal_attribute(
+        self, span_context: SpanContext, key: str, value: Any
+    ) -> None:
+        span_key = self._get_span_key(span_context)
+        self._internal_attributes[span_key][key] = value
+
+    def get_internal_attribute(
+        self, span_context: SpanContext, key: str, default: Any = None
+    ) -> Any:
+        span_key = self._get_span_key(span_context)
+        return self._internal_attributes[span_key].get(key, default)
+
+    def increment_update_id(self, span_context: SpanContext) -> int:
+        current_id = self.get_internal_attribute(
+            span_context=span_context, key=AttributeKeys.JUDGMENT_UPDATE_ID, default=0
+        )
+        new_id = current_id + 1
+        self.set_internal_attribute(
+            span_context=span_context,
+            key=AttributeKeys.JUDGMENT_UPDATE_ID,
+            value=new_id,
+        )
+        return current_id
+
+    def _cleanup_span_state(self, span_key: tuple[int, int]) -> None:
+        self._internal_attributes.pop(span_key, None)
+
+    @dont_throw
+    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+        for processor in get_all():
+            processor.on_start(span, parent_context)
+
+    @dont_throw
+    def emit_partial(self) -> None:
+        current_span = get_current_span()
+        if (
+            not current_span
+            or not current_span.is_recording()
+            or not isinstance(current_span, ReadableSpan)
+        ):
+            return
+
+        span_context = current_span.get_span_context()
+        if self.get_internal_attribute(
+            span_context, InternalAttributeKeys.DISABLE_PARTIAL_EMIT, False
+        ):
+            return
+
+        attributes = dict(current_span.attributes or {})
+        attributes[AttributeKeys.JUDGMENT_UPDATE_ID] = self.increment_update_id(
+            span_context
+        )
+
+        partial_span = ReadableSpan(
+            name=current_span.name,
+            context=span_context,
+            parent=current_span.parent,
+            resource=current_span.resource,
+            attributes=attributes,
+            events=current_span.events,
+            links=current_span.links,
+            status=current_span.status,
+            kind=current_span.kind,
+            start_time=current_span.start_time,
+            end_time=None,
+            instrumentation_scope=current_span.instrumentation_scope,
+        )
+
+        super().on_end(partial_span)
+
+    @dont_throw
+    def on_end(self, span: ReadableSpan) -> None:
+        for processor in get_all():
+            processor.on_end(span)
+
+        if not span.context:
+            super().on_end(span)
+            return
+
+        span_key = self._get_span_key(span.context)
+
+        if self.get_internal_attribute(
+            span.context, InternalAttributeKeys.CANCELLED, False
+        ):
+            self._cleanup_span_state(span_key)
+            return
+
+        if span.end_time is not None:
+            attributes = dict(span.attributes or {})
+            attributes[AttributeKeys.JUDGMENT_UPDATE_ID] = 20
+
+            final_span = ReadableSpan(
+                name=span.name,
+                context=span.context,
+                parent=span.parent,
+                resource=span.resource,
+                attributes=attributes,
+                events=span.events,
+                links=span.links,
+                status=span.status,
+                kind=span.kind,
+                start_time=span.start_time,
+                end_time=span.end_time,
+                instrumentation_scope=span.instrumentation_scope,
+            )
+
+            self._cleanup_span_state(span_key)
+            super().on_end(final_span)
+        else:
+            super().on_end(span)

judgeval/v1/tracer/processors/noop_span_processor.py

@@ -0,0 +1,42 @@
+from opentelemetry.sdk.trace import Span, ReadableSpan
+from opentelemetry.context import Context
+from typing import Any, Optional
+
+from opentelemetry.trace import SpanContext
+
+from judgeval.v1.tracer.processors.judgment_span_processor import JudgmentSpanProcessor
+
+
+class NoOpJudgmentSpanProcessor(JudgmentSpanProcessor):
+    __slots__ = ("resource_attributes",)
+
+    def __init__(self):
+        self.resource_attributes = {}
+
+    def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
+        pass
+
+    def on_end(self, span: ReadableSpan) -> None:
+        pass
+
+    def shutdown(self) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: int | None = 30000) -> bool:
+        return True
+
+    def emit_partial(self) -> None:
+        pass
+
+    def set_internal_attribute(
+        self, span_context: SpanContext, key: str, value: Any
+    ) -> None:
+        pass
+
+    def get_internal_attribute(
+        self, span_context: SpanContext, key: str, default: Any = None
+    ) -> Any:
+        return default
+
+    def increment_update_id(self, span_context: SpanContext) -> int:
+        return 0

judgeval/v1/tracer/tracer.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Optional
+
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+
+from judgeval.v1.internal.api import JudgmentSyncClient
+from judgeval.logger import judgeval_logger
+from judgeval.v1.tracer.judgment_tracer_provider import JudgmentTracerProvider
+from judgeval.version import get_version
+from judgeval.v1.tracer.base_tracer import BaseTracer
+from judgeval.v1.tracer.judgment_tracer_provider import FilterTracerCallback
+
+
+class Tracer(BaseTracer):
+    __slots__ = ("_tracer_provider", "_filter_tracer")
+
+    def __init__(
+        self,
+        project_name: str,
+        enable_evaluation: bool,
+        api_client: JudgmentSyncClient,
+        serializer: Callable[[Any], str],
+        initialize: bool,
+        filter_tracer: Optional[FilterTracerCallback] = None,
+    ):
+        super().__init__(
+            project_name=project_name,
+            enable_evaluation=enable_evaluation,
+            api_client=api_client,
+            serializer=serializer,
+        )
+        self._tracer_provider: Optional[TracerProvider] = None
+        self._filter_tracer = filter_tracer
+
+        if initialize:
+            self.initialize()
+
+    def initialize(self) -> None:
+        resource = Resource.create(
+            {
+                "service.name": self.project_name,
+                "telemetry.sdk.name": self.TRACER_NAME,
+                "telemetry.sdk.version": get_version(),
+            }
+        )
+
+        self._tracer_provider = JudgmentTracerProvider(
+            resource=resource, filter_tracer=self._filter_tracer
+        )
+        self._tracer_provider.add_span_processor(self.get_span_processor())
+
+        trace.set_tracer_provider(self._tracer_provider)
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        if self._tracer_provider is None:
+            judgeval_logger.error("Cannot forceFlush: tracer not initialized")
+            return False
+        return self._tracer_provider.force_flush(timeout_millis)
+
+    def shutdown(self, timeout_millis: int = 30000) -> None:
+        if self._tracer_provider is None:
+            judgeval_logger.error("Cannot shutdown: tracer not initialized")
+            return
+        self._tracer_provider.shutdown()

judgeval/v1/tracer/tracer_factory.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from typing import Any, Callable, Optional
+
+from judgeval.utils.serialize import safe_serialize
+from judgeval.v1.internal.api import JudgmentSyncClient
+from judgeval.v1.tracer.judgment_tracer_provider import FilterTracerCallback
+from judgeval.v1.tracer.tracer import Tracer
+
+
+class TracerFactory:
+    __slots__ = "_client"
+
+    def __init__(
+        self,
+        client: JudgmentSyncClient,
+    ):
+        self._client = client
+
+    def create(
+        self,
+        project_name: str,
+        enable_evaluation: bool = True,
+        serializer: Optional[Callable[[Any], str]] = None,
+        filter_tracer: Optional[FilterTracerCallback] = None,
+        initialize: bool = True,
+    ) -> Tracer:
+        if serializer is None:
+            serializer = safe_serialize
+
+        return Tracer(
+            project_name=project_name,
+            enable_evaluation=enable_evaluation,
+            api_client=self._client,
+            serializer=serializer,
+            initialize=initialize,
+            filter_tracer=filter_tracer,
+        )

judgeval/v1/trainers/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from judgeval.v1.trainers.trainers_factory import TrainersFactory
+
+__all__ = ["TrainersFactory"]

judgeval/v1/trainers/base_trainer.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from judgeval.v1.tracer.tracer import Tracer
+    from judgeval.v1.trainers.trainable_model import TrainableModel
+    from judgeval.v1.trainers.config import TrainerConfig, ModelConfig
+    from judgeval.v1.scorers.base_scorer import BaseScorer
+
+
+class BaseTrainer(ABC):
+    __slots__ = ("config", "trainable_model", "tracer", "project_name")
+
+    def __init__(
+        self,
+        config: TrainerConfig,
+        trainable_model: TrainableModel,
+        tracer: Tracer,
+        project_name: Optional[str] = None,
+    ):
+        self.config = config
+        self.trainable_model = trainable_model
+        self.tracer = tracer
+        self.project_name = project_name or "judgment_training"
+
+    @abstractmethod
+    async def generate_rollouts_and_rewards(
+        self,
+        agent_function: Callable[[Any], Any],
+        scorers: List[BaseScorer],
+        prompts: dict[int, dict[Any, Any]],
+        num_prompts_per_step: Optional[int] = None,
+        num_generations_per_prompt: Optional[int] = None,
+        concurrency: Optional[int] = None,
+    ) -> Any:
+        pass
+
+    @abstractmethod
+    async def run_reinforcement_learning(
+        self,
+        agent_function: Callable[[Any], Any],
+        scorers: List[BaseScorer],
+        prompts: dict[int, dict[Any, Any]],
+    ) -> "ModelConfig":
+        pass
+
+    @abstractmethod
+    async def train(
+        self,
+        agent_function: Callable[[Any], Any],
+        scorers: List[BaseScorer],
+        prompts: dict[int, dict[Any, Any]],
+    ) -> "ModelConfig":
+        pass
+
+    @abstractmethod
+    def _extract_message_history_from_spans(
+        self, trace_id: str
+    ) -> List[Dict[str, str]]:
+        pass