judgeval 0.0.11__py3-none-any.whl → 0.22.2__py3-none-any.whl

judgeval/tracer/processors/__init__.py

@@ -0,0 +1,220 @@
+from __future__ import annotations
+from typing import Optional, TYPE_CHECKING, Any
+from collections import defaultdict
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
+from opentelemetry.trace.span import SpanContext
+from opentelemetry.sdk.trace.export import (
+    BatchSpanProcessor,
+)
+from judgeval.tracer.exporters import JudgmentSpanExporter
+from judgeval.tracer.keys import AttributeKeys, InternalAttributeKeys, ResourceKeys
+from judgeval.utils.url import url_for
+from judgeval.utils.decorators.dont_throw import dont_throw
+from judgeval.version import get_version
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+
+
+class NoOpSpanProcessor(SpanProcessor):
+    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 = 30000) -> bool:
+        return True
+
+
+class JudgmentSpanProcessor(BatchSpanProcessor):
+    __slots__ = ("tracer", "resource_attributes", "_internal_attributes")
+
+    def __init__(
+        self,
+        tracer: Tracer,
+        project_name: str,
+        project_id: str,
+        api_key: str,
+        organization_id: str,
+        /,
+        *,
+        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,
+        resource_attributes: Optional[dict[str, Any]] = None,
+    ):
+        self.tracer = tracer
+
+        attrs = {
+            ResourceKeys.SERVICE_NAME: project_name,
+            ResourceKeys.TELEMETRY_SDK_NAME: "judgeval",
+            ResourceKeys.TELEMETRY_SDK_VERSION: get_version(),
+            ResourceKeys.JUDGMENT_PROJECT_ID: project_id,
+            **(resource_attributes or {}),
+        }
+        self.resource_attributes = attrs
+
+        super().__init__(
+            JudgmentSpanExporter(
+                endpoint=url_for("/otel/v1/traces"),
+                api_key=api_key,
+                organization_id=organization_id,
+                project_id=project_id,
+            ),
+            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 emit_partial(self) -> None:
+        current_span = self.tracer.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)
+
+    def on_end(self, span: ReadableSpan) -> None:
+        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)
+
+
+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
+
+
+__all__ = ["NoOpSpanProcessor", "JudgmentSpanProcessor", "NoOpJudgmentSpanProcessor"]

judgeval/tracer/utils.py

@@ -0,0 +1,19 @@
+from typing import Any
+from opentelemetry.trace import Span
+from pydantic import BaseModel
+from typing import Callable, Optional
+from judgeval.scorers.api_scorer import TraceAPIScorerConfig
+
+
+def set_span_attribute(span: Span, name: str, value: Any):
+    if value is None or value == "":
+        return
+
+    span.set_attribute(name, value)
+
+
+class TraceScorerConfig(BaseModel):
+    scorer: TraceAPIScorerConfig | None
+    model: Optional[str] = None
+    sampling_rate: float = 1.0
+    run_condition: Optional[Callable[..., bool]] = None

judgeval/trainer/__init__.py

@@ -0,0 +1,14 @@
+from judgeval.trainer.trainer import JudgmentTrainer
+from judgeval.trainer.config import TrainerConfig, ModelConfig
+from judgeval.trainer.trainable_model import TrainableModel
+from judgeval.trainer.base_trainer import BaseTrainer
+from judgeval.trainer.fireworks_trainer import FireworksTrainer
+
+__all__ = [
+    "JudgmentTrainer",
+    "TrainerConfig",
+    "ModelConfig",
+    "TrainableModel",
+    "BaseTrainer",
+    "FireworksTrainer",
+]

judgeval/trainer/base_trainer.py

@@ -0,0 +1,122 @@
+from abc import ABC, abstractmethod
+from typing import Any, Callable, List, Optional, Union, Dict, TYPE_CHECKING
+from .config import TrainerConfig, ModelConfig
+from judgeval.scorers import ExampleScorer, ExampleAPIScorerConfig
+
+if TYPE_CHECKING:
+    from judgeval.tracer import Tracer
+    from .trainable_model import TrainableModel
+
+
+class BaseTrainer(ABC):
+    """
+    Abstract base class for training providers.
+
+    This class defines the interface that all training provider implementations
+    must follow. Each provider (Fireworks, Verifiers, etc.) will have its own
+    concrete implementation of this interface.
+    """
+
+    def __init__(
+        self,
+        config: TrainerConfig,
+        trainable_model: "TrainableModel",
+        tracer: "Tracer",
+        project_name: Optional[str] = None,
+    ):
+        """
+        Initialize the base trainer.
+
+        Args:
+            config: TrainerConfig instance with training parameters
+            trainable_model: TrainableModel instance to use for training
+            tracer: Tracer for observability
+            project_name: Project name for organizing training runs
+        """
+        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[Union[ExampleAPIScorerConfig, ExampleScorer]],
+        prompts: List[Any],
+        num_prompts_per_step: Optional[int] = None,
+        num_generations_per_prompt: Optional[int] = None,
+        concurrency: Optional[int] = None,
+    ) -> Any:
+        """
+        Generate rollouts and compute rewards using the current model snapshot.
+
+        Args:
+            agent_function: Function/agent to call for generating responses
+            scorers: List of scorer objects to evaluate responses
+            prompts: List of prompts to use for training
+            num_prompts_per_step: Number of prompts to use per step
+            num_generations_per_prompt: Generations per prompt
+            concurrency: Concurrency limit
+
+        Returns:
+            Provider-specific dataset format for training
+        """
+        pass
+
+    @abstractmethod
+    async def run_reinforcement_learning(
+        self,
+        agent_function: Callable[[Any], Any],
+        scorers: List[Union[ExampleAPIScorerConfig, ExampleScorer]],
+        prompts: List[Any],
+    ) -> ModelConfig:
+        """
+        Run the iterative reinforcement learning fine-tuning loop.
+
+        Args:
+            agent_function: Function/agent to call for generating responses
+            scorers: List of scorer objects to evaluate responses
+            prompts: List of prompts to use for training
+
+        Returns:
+            ModelConfig: Configuration of the trained model
+        """
+        pass
+
+    @abstractmethod
+    async def train(
+        self,
+        agent_function: Callable[[Any], Any],
+        scorers: List[Union[ExampleAPIScorerConfig, ExampleScorer]],
+        prompts: List[Any],
+    ) -> ModelConfig:
+        """
+        Start the reinforcement learning fine-tuning process.
+
+        This is the main entry point for running the training.
+
+        Args:
+            agent_function: Function/agent to call for generating responses
+            scorers: List of scorer objects to evaluate responses
+            prompts: List of prompts to use for training
+
+        Returns:
+            ModelConfig: Configuration of the trained model
+        """
+        pass
+
+    @abstractmethod
+    def _extract_message_history_from_spans(
+        self, trace_id: str
+    ) -> List[Dict[str, str]]:
+        """
+        Extract message history from spans for training purposes.
+
+        Args:
+            trace_id: The trace ID (32-char hex string) to extract message history from
+
+        Returns:
+            List of message dictionaries with 'role' and 'content' keys
+        """
+        pass

judgeval/trainer/config.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Dict, Any, TYPE_CHECKING
+import json
+
+if TYPE_CHECKING:
+    from fireworks.llm.llm_reinforcement_step import ReinforcementAcceleratorTypeLiteral  # type: ignore[import-not-found]
+
+
+@dataclass
+class TrainerConfig:
+    """Configuration class for JudgmentTrainer parameters."""
+
+    deployment_id: str
+    user_id: str
+    model_id: str
+    base_model_name: str = "qwen2p5-7b-instruct"
+    rft_provider: str = "fireworks"  # Supported: "fireworks", "verifiers" (future)
+    num_steps: int = 5
+    num_generations_per_prompt: int = 4
+    num_prompts_per_step: int = 4
+    concurrency: int = 100
+    epochs: int = 1
+    learning_rate: float = 1e-5
+    accelerator_count: int = 1
+    accelerator_type: ReinforcementAcceleratorTypeLiteral = "NVIDIA_A100_80GB"
+    temperature: float = 1.5
+    max_tokens: int = 50
+    enable_addons: bool = True
+
+
+@dataclass
+class ModelConfig:
+    """
+    Configuration class for storing and loading trained model state.
+
+    This class enables persistence of trained models so they can be loaded
+    and used later without retraining.
+
+    Example usage:
+        trainer = JudgmentTrainer(config)
+        model_config = trainer.train(agent_function, scorers, prompts)
+
+        # Save the trained model configuration
+        model_config.save_to_file("my_trained_model.json")
+
+        # Later, load and use the trained model
+        loaded_config = ModelConfig.load_from_file("my_trained_model.json")
+        trained_model = TrainableModel.from_model_config(loaded_config)
+
+        # Use the trained model for inference
+        response = trained_model.chat.completions.create(
+            model="current",  # Uses the loaded trained model
+            messages=[{"role": "user", "content": "Hello!"}]
+        )
+    """
+
+    # Base model configuration
+    base_model_name: str
+    deployment_id: str
+    user_id: str
+    model_id: str
+    enable_addons: bool
+
+    # Training state
+    current_step: int
+    total_steps: int
+
+    # Current model information
+    current_model_name: Optional[str] = None
+    is_trained: bool = False
+
+    # Training parameters used (for reference)
+    training_params: Optional[Dict[str, Any]] = None
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert ModelConfig to dictionary for serialization."""
+        return {
+            "base_model_name": self.base_model_name,
+            "deployment_id": self.deployment_id,
+            "user_id": self.user_id,
+            "model_id": self.model_id,
+            "enable_addons": self.enable_addons,
+            "current_step": self.current_step,
+            "total_steps": self.total_steps,
+            "current_model_name": self.current_model_name,
+            "is_trained": self.is_trained,
+            "training_params": self.training_params,
+        }
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> ModelConfig:
+        """Create ModelConfig from dictionary."""
+        return cls(
+            base_model_name=data.get("base_model_name", "qwen2p5-7b-instruct"),
+            deployment_id=data.get("deployment_id", "my-base-deployment"),
+            user_id=data.get("user_id", ""),
+            model_id=data.get("model_id", ""),
+            enable_addons=data.get("enable_addons", True),
+            current_step=data.get("current_step", 0),
+            total_steps=data.get("total_steps", 0),
+            current_model_name=data.get("current_model_name"),
+            is_trained=data.get("is_trained", False),
+            training_params=data.get("training_params"),
+        )
+
+    def to_json(self) -> str:
+        """Convert ModelConfig to JSON string."""
+        return json.dumps(self.to_dict(), indent=2)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> ModelConfig:
+        """Create ModelConfig from JSON string."""
+        data = json.loads(json_str)
+        return cls.from_dict(data)
+
+    def save_to_file(self, filepath: str):
+        """Save ModelConfig to a JSON file."""
+        with open(filepath, "w") as f:
+            f.write(self.to_json())
+
+    @classmethod
+    def load_from_file(cls, filepath: str) -> ModelConfig:
+        """Load ModelConfig from a JSON file."""
+        with open(filepath, "r") as f:
+            json_str = f.read()
+        return cls.from_json(json_str)

judgeval/trainer/console.py

@@ -0,0 +1,144 @@
+from contextlib import contextmanager
+from typing import Optional
+import sys
+import os
+from judgeval.utils.decorators.use_once import use_once
+
+
+@use_once
+def _is_jupyter_environment():
+    """Check if we're running in a Jupyter notebook or similar environment."""
+    try:
+        # Check for IPython kernel
+        if "ipykernel" in sys.modules or "IPython" in sys.modules:
+            return True
+        # Check for Jupyter environment variables
+        if "JPY_PARENT_PID" in os.environ:
+            return True
+        # Check if we're in Google Colab
+        if "google.colab" in sys.modules:
+            return True
+        return False
+    except Exception:
+        return False
+
+
+IS_JUPYTER = _is_jupyter_environment()
+
+if not IS_JUPYTER:
+    try:
+        from rich.console import Console
+        from rich.spinner import Spinner
+        from rich.live import Live
+        from rich.text import Text
+
+        shared_console = Console()
+        RICH_AVAILABLE = True
+    except ImportError:
+        RICH_AVAILABLE = False
+else:
+    RICH_AVAILABLE = False
+
+
+class SimpleSpinner:
+    def __init__(self, name, text):
+        self.text = text
+
+
+class SimpleLive:
+    def __init__(self, spinner, console=None, refresh_per_second=None):
+        self.spinner = spinner
+
+    def __enter__(self):
+        print(f"🔄 {self.spinner.text}")
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+    def update(self, spinner):
+        print(f"🔄 {spinner.text}")
+
+
+def safe_print(message, style=None):
+    """Safe print function that works in all environments."""
+    if RICH_AVAILABLE and not IS_JUPYTER:
+        shared_console.print(message, style=style)
+    else:
+        if style == "green":
+            print(f"✅ {message}")
+        elif style == "yellow":
+            print(f"⚠️ {message}")
+        elif style == "blue":
+            print(f"🔵 {message}")
+        elif style == "cyan":
+            print(f"🔷 {message}")
+        else:
+            print(message)
+
+
+@contextmanager
+def _spinner_progress(
+    message: str, step: Optional[int] = None, total_steps: Optional[int] = None
+):
+    """Context manager for spinner-based progress display."""
+    if step is not None and total_steps is not None:
+        full_message = f"[Step {step}/{total_steps}] {message}"
+    else:
+        full_message = f"[Training] {message}"
+
+    if RICH_AVAILABLE and not IS_JUPYTER:
+        spinner = Spinner("dots", text=Text(full_message, style="cyan"))
+        with Live(spinner, console=shared_console, refresh_per_second=10):
+            yield
+    else:
+        print(f"🔄 {full_message}")
+        try:
+            yield
+        finally:
+            print(f"✅ {full_message} - Complete")
+
+
+@contextmanager
+def _model_spinner_progress(message: str):
+    """Context manager for model operation spinner-based progress display."""
+    if RICH_AVAILABLE and not IS_JUPYTER:
+        spinner = Spinner("dots", text=Text(f"[Model] {message}", style="blue"))
+        with Live(spinner, console=shared_console, refresh_per_second=10) as live:
+
+            def update_progress(progress_message: str):
+                """Update the spinner with a new progress message."""
+                new_text = f"[Model] {message}\n  └─ {progress_message}"
+                spinner.text = Text(new_text, style="blue")
+                live.update(spinner)
+
+            yield update_progress
+    else:
+        print(f"🔵 [Model] {message}")
+
+        def update_progress(progress_message: str):
+            print(f"  └─ {progress_message}")
+
+        yield update_progress
+
+
+def _print_progress(
+    message: str, step: Optional[int] = None, total_steps: Optional[int] = None
+):
+    """Print progress message with consistent formatting."""
+    if step is not None and total_steps is not None:
+        safe_print(f"[Step {step}/{total_steps}] {message}", style="green")
+    else:
+        safe_print(f"[Training] {message}", style="green")
+
+
+def _print_progress_update(
+    message: str, step: Optional[int] = None, total_steps: Optional[int] = None
+):
+    """Print progress update message (for status changes during long operations)."""
+    safe_print(f"  └─ {message}", style="yellow")
+
+
+def _print_model_progress(message: str):
+    """Print model progress message with consistent formatting."""
+    safe_print(f"[Model] {message}", style="blue")