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

judgeval/v1/trainers/config.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Dict, Any
+import json
+
+
+@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
+    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/v1/trainers/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")

judgeval/v1/trainers/fireworks_trainer.py

@@ -0,0 +1,392 @@
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING
+
+from fireworks import Dataset  # type: ignore[import-not-found,import-untyped]
+
+if TYPE_CHECKING:
+    from judgeval.v1.trainers.config import TrainerConfig, ModelConfig
+    from judgeval.v1.trainers.trainable_model import TrainableModel
+    from judgeval.v1.tracer.tracer import Tracer
+    from judgeval.v1.scorers.base_scorer import BaseScorer
+    from judgeval.v1.internal.api import JudgmentSyncClient
+
+from judgeval.v1.trainers.base_trainer import BaseTrainer
+from judgeval.v1.tracer.exporters import SpanStore, InMemorySpanExporter
+from judgeval.judgment_attribute_keys import AttributeKeys
+from judgeval.v1.data.example import Example
+from judgeval.v1.data.scoring_result import ScoringResult
+from judgeval.v1.internal.api.api_types import ExampleEvaluationRun
+from judgeval.v1.trainers.console import (
+    _spinner_progress,
+    _print_progress,
+    _print_progress_update,
+)
+from judgeval.exceptions import JudgmentRuntimeError
+from opentelemetry import trace
+
+
+class FireworksTrainer(BaseTrainer):
+    __slots__ = ("_client", "span_store", "span_exporter")
+
+    def __init__(
+        self,
+        config: "TrainerConfig",
+        trainable_model: "TrainableModel",
+        tracer: "Tracer",
+        project_name: Optional[str] = None,
+        client: Optional["JudgmentSyncClient"] = None,
+    ):
+        super().__init__(config, trainable_model, tracer, project_name)
+        if client is None:
+            raise ValueError("client is required")
+        self._client = client
+        self.span_store = SpanStore()
+        self.span_exporter = InMemorySpanExporter(self.span_store)
+
+    def _extract_message_history_from_spans(
+        self, trace_id: str
+    ) -> List[Dict[str, str]]:
+        spans = self.span_store.get_by_trace_id(trace_id)
+        if not spans:
+            return []
+
+        messages = []
+        first_found = False
+
+        for span in sorted(spans, key=lambda s: getattr(s, "start_time", 0)):
+            span_attributes = span.attributes or {}
+            span_type = span_attributes.get(AttributeKeys.JUDGMENT_SPAN_KIND, "span")
+
+            if (
+                not span_attributes.get(AttributeKeys.JUDGMENT_OUTPUT)
+                and span_type != "llm"
+            ):
+                continue
+
+            if span_type == "llm":
+                if not first_found and span_attributes.get(
+                    AttributeKeys.JUDGMENT_INPUT
+                ):
+                    input_data: Any = span_attributes.get(
+                        AttributeKeys.JUDGMENT_INPUT, {}
+                    )
+                    if isinstance(input_data, dict) and "messages" in input_data:
+                        input_messages = input_data["messages"]
+                        if input_messages:
+                            first_found = True
+                            for msg in input_messages:
+                                if (
+                                    isinstance(msg, dict)
+                                    and "role" in msg
+                                    and "content" in msg
+                                ):
+                                    messages.append(
+                                        {"role": msg["role"], "content": msg["content"]}
+                                    )
+
+                output = span_attributes.get(AttributeKeys.JUDGMENT_OUTPUT)
+                if output is not None:
+                    content = str(output)
+                    try:
+                        parsed = json.loads(content)
+                        if isinstance(parsed, dict) and "messages" in parsed:
+                            for msg in parsed["messages"]:
+                                if (
+                                    isinstance(msg, dict)
+                                    and msg.get("role") == "assistant"
+                                ):
+                                    content = msg.get("content", content)
+                                    break
+                    except (json.JSONDecodeError, KeyError):
+                        pass
+                    messages.append({"role": "assistant", "content": content})
+
+            elif span_type in ("user", "tool"):
+                output = span_attributes.get(AttributeKeys.JUDGMENT_OUTPUT)
+                if output is not None:
+                    content = str(output)
+                    try:
+                        parsed = json.loads(content)
+                        if isinstance(parsed, dict) and "messages" in parsed:
+                            for msg in parsed["messages"]:
+                                if isinstance(msg, dict) and msg.get("role") == "user":
+                                    content = msg.get("content", content)
+                                    break
+                    except (json.JSONDecodeError, KeyError):
+                        pass
+                    messages.append({"role": "user", "content": content})
+
+        return messages
+
+    async def generate_rollouts_and_rewards(
+        self,
+        agent_function: Callable[..., 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,
+    ):
+        num_prompts_per_step = min(
+            num_prompts_per_step or self.config.num_prompts_per_step, len(prompts)
+        )
+        num_generations_per_prompt = (
+            num_generations_per_prompt or self.config.num_generations_per_prompt
+        )
+        concurrency = concurrency or self.config.concurrency
+
+        semaphore = asyncio.Semaphore(concurrency)
+
+        @self.tracer.observe(span_type="function")
+        async def generate_single_response(prompt_id: int, generation_id: int):
+            async with semaphore:
+                prompt_input = prompts[prompt_id]
+                response_data = await agent_function(**prompt_input)
+                messages = response_data.get("messages", [])
+
+                current_span = trace.get_current_span()
+                trace_id = None
+                if current_span and current_span.is_recording():
+                    trace_id = format(current_span.get_span_context().trace_id, "032x")
+
+                try:
+                    if trace_id is not None:
+                        traced_messages = self._extract_message_history_from_spans(
+                            trace_id
+                        )
+                        if traced_messages:
+                            messages = traced_messages
+                except Exception as e:
+                    print(f"Warning: Failed to get message history from trace: {e}")
+                finally:
+                    if trace_id is not None:
+                        self.span_store.clear_trace(trace_id)
+
+                example = Example()
+                example.set_property("input", prompt_input)
+                example.set_property("messages", messages)
+                example.set_property("actual_output", response_data)
+
+                evaluation_run: ExampleEvaluationRun = {
+                    "project_name": self.project_name,
+                    "eval_name": f"training_step_{self.trainable_model.current_step}_prompt_{prompt_id}_gen_{generation_id}",
+                    "examples": [example.to_dict()],
+                    "judgment_scorers": [
+                        scorer.get_scorer_config() for scorer in scorers
+                    ],
+                }
+
+                response = self._client.add_to_run_eval_queue_examples(evaluation_run)
+                if not response.get("success", False):
+                    raise JudgmentRuntimeError(
+                        f"Failed to queue evaluation: {response.get('error', 'Unknown error')}"
+                    )
+
+                results = await self._poll_evaluation_until_complete(
+                    evaluation_run["eval_name"], len(scorers)
+                )
+
+                if results and results[0].scorers_data:
+                    scores = [
+                        scorer_data.score
+                        for scorer_data in results[0].scorers_data
+                        if scorer_data.score is not None
+                    ]
+                    reward = sum(scores) / len(scores) if scores else 0.0
+                else:
+                    reward = 0.0
+
+            return {
+                "prompt_id": prompt_id,
+                "generation_id": generation_id,
+                "messages": messages,
+                "evals": {"score": reward},
+            }
+
+        coros = []
+        for prompt_id in range(num_prompts_per_step):
+            for generation_id in range(num_generations_per_prompt):
+                coro = generate_single_response(prompt_id, generation_id)
+                coros.append(coro)
+
+        with _spinner_progress(f"Generating {len(coros)} rollouts..."):
+            num_completed = 0
+            results = []
+
+            for future in asyncio.as_completed(coros):
+                result = await future
+                results.append(result)
+                num_completed += 1
+
+        _print_progress(f"Generated {len(results)} rollouts successfully")
+
+        dataset_rows = []
+        for prompt_id in range(num_prompts_per_step):
+            prompt_generations = [r for r in results if r["prompt_id"] == prompt_id]
+            sample_generations = [
+                {"messages": gen["messages"], "evals": gen["evals"]}
+                for gen in prompt_generations
+            ]
+            dataset_rows.append({"samples": sample_generations})
+
+        return dataset_rows
+
+    async def _poll_evaluation_until_complete(
+        self, eval_name: str, expected_scorers: int
+    ) -> List[ScoringResult]:
+        import time
+
+        max_wait_time = 300
+        poll_interval = 2
+        start_time = time.time()
+
+        while time.time() - start_time < max_wait_time:
+            await asyncio.sleep(poll_interval)
+
+            from judgeval.v1.internal.api.api_types import EvalResultsFetch
+
+            fetch_request: EvalResultsFetch = {
+                "experiment_run_id": eval_name,
+                "project_name": self.project_name,
+            }
+
+            try:
+                response = self._client.fetch_experiment_run(fetch_request)
+                if response and response.get("results"):
+                    results_data = response.get("results")
+                    if results_data is not None and len(results_data) > 0:
+                        scoring_results = []
+                        for result_data in results_data:
+                            from judgeval.v1.data.scorer_data import ScorerData
+
+                            scorers_data = []
+                            for scorer_result in result_data.get("scorers_data", []):
+                                scorers_data.append(
+                                    ScorerData(
+                                        name=scorer_result.get("name", ""),
+                                        threshold=scorer_result.get("threshold", 0.0),
+                                        success=scorer_result.get("success", False),
+                                        score=scorer_result.get("score"),
+                                        reason=scorer_result.get("reason"),
+                                    )
+                                )
+
+                            scoring_results.append(
+                                ScoringResult(
+                                    success=result_data.get("success", False),
+                                    scorers_data=scorers_data,
+                                    name=result_data.get("name"),
+                                    trace_id=result_data.get("trace_id"),
+                                    run_duration=result_data.get("run_duration"),
+                                    evaluation_cost=result_data.get("evaluation_cost"),
+                                )
+                            )
+                        return scoring_results
+            except Exception:
+                pass
+
+        raise JudgmentRuntimeError(
+            f"Evaluation {eval_name} did not complete within {max_wait_time} seconds"
+        )
+
+    async def run_reinforcement_learning(
+        self,
+        agent_function: Callable[[Any], Any],
+        scorers: List["BaseScorer"],
+        prompts: dict[int, dict[Any, Any]],
+    ) -> "ModelConfig":
+        _print_progress("Starting reinforcement learning training")
+
+        training_params = {
+            "num_steps": self.config.num_steps,
+            "num_prompts_per_step": self.config.num_prompts_per_step,
+            "num_generations_per_prompt": self.config.num_generations_per_prompt,
+            "epochs": self.config.epochs,
+            "learning_rate": self.config.learning_rate,
+            "temperature": self.config.temperature,
+            "max_tokens": self.config.max_tokens,
+        }
+
+        start_step = self.trainable_model.current_step
+
+        for step in range(start_step, self.config.num_steps):
+            step_num = step + 1
+            _print_progress(
+                f"Starting training step {step_num}", step_num, self.config.num_steps
+            )
+
+            self.trainable_model.advance_to_next_step(step)
+
+            dataset_rows = await self.generate_rollouts_and_rewards(
+                agent_function, scorers, prompts
+            )
+
+            with _spinner_progress(
+                "Preparing training dataset", step_num, self.config.num_steps
+            ):
+                dataset = Dataset.from_list(dataset_rows)
+                dataset.sync()
+
+            _print_progress(
+                "Starting reinforcement training", step_num, self.config.num_steps
+            )
+            job = self.trainable_model.perform_reinforcement_step(dataset, step)
+
+            last_state = None
+            with _spinner_progress(
+                "Training job in progress", step_num, self.config.num_steps
+            ):
+                while not job.is_completed:
+                    job.raise_if_bad_state()
+                    current_state = job.state
+
+                    if current_state != last_state:
+                        if current_state in ["uploading", "validating"]:
+                            _print_progress_update(
+                                f"Training job: {current_state} data"
+                            )
+                        elif current_state == "training":
+                            _print_progress_update(
+                                "Training job: model training in progress"
+                            )
+                        else:
+                            _print_progress_update(f"Training job: {current_state}")
+                        last_state = current_state
+
+                    await asyncio.sleep(10)
+                    job = job.get()
+                    if job is None:
+                        raise JudgmentRuntimeError(
+                            "Training job was deleted while waiting for completion"
+                        )
+
+            _print_progress(
+                f"Training completed! New model: {job.output_model}",
+                step_num,
+                self.config.num_steps,
+            )
+
+        _print_progress("All training steps completed!")
+
+        with _spinner_progress("Deploying final trained model"):
+            self.trainable_model.advance_to_next_step(self.config.num_steps)
+
+        return self.trainable_model.get_model_config(training_params)
+
+    async def train(
+        self,
+        agent_function: Callable[[Any], Any],
+        scorers: List["BaseScorer"],
+        prompts: dict[int, dict[Any, Any]],
+    ) -> "ModelConfig":
+        try:
+            return await self.run_reinforcement_learning(
+                agent_function, scorers, prompts
+            )
+        except JudgmentRuntimeError:
+            raise
+        except Exception as e:
+            raise JudgmentRuntimeError(f"Training process failed: {str(e)}") from e