judgeval 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

judgeval/common/trainer/trainer.py

@@ -0,0 +1,301 @@
+import asyncio
+import time
+from typing import Optional, Callable, Any, List, Union
+from fireworks import Dataset
+from .config import TrainerConfig, ModelConfig
+from .trainable_model import TrainableModel
+from judgeval.tracer import Tracer
+from judgeval.judgment_client import JudgmentClient
+from judgeval.scorers import BaseScorer, APIScorerConfig
+from judgeval.data import Example
+from .console import _spinner_progress, _print_progress, _print_progress_update
+from judgeval.common.exceptions import JudgmentAPIError
+
+
+class JudgmentTrainer:
+    """
+    A reinforcement learning trainer for Judgment models using Fine-Tuning.
+
+    This class handles the iterative training process where models are improved
+    through reinforcement learning fine-tuning based on generated rollouts and rewards.
+    """
+
+    def __init__(
+        self,
+        config: TrainerConfig,
+        trainable_model: TrainableModel,
+        tracer: Tracer,
+        project_name: Optional[str] = None,
+    ):
+        """
+        Initialize the JudgmentTrainer.
+
+        Args:
+            config: TrainerConfig instance with training parameters. If None, uses default config.
+            tracer: Optional tracer for observability
+            trainable_model: Optional trainable model instance
+            project_name: Project name for organizing training runs and evaluations
+        """
+        try:
+            self.config = config
+            self.tracer = tracer
+            self.tracer.show_trace_urls = False
+            self.project_name = project_name or "judgment_training"
+
+            if trainable_model is None:
+                self.trainable_model = TrainableModel(self.config)
+            else:
+                self.trainable_model = trainable_model
+
+            self.judgment_client = JudgmentClient()
+        except Exception as e:
+            raise JudgmentAPIError(
+                f"Failed to initialize JudgmentTrainer: {str(e)}"
+            ) from e
+
+    async def generate_rollouts_and_rewards(
+        self,
+        agent_function: Callable[[Any], Any],
+        scorers: List[Union[APIScorerConfig, BaseScorer]],
+        prompts: List[Any],
+        num_prompts_per_step: Optional[int] = None,
+        num_generations_per_prompt: Optional[int] = None,
+        concurrency: Optional[int] = None,
+    ):
+        """
+        Generate rollouts and compute rewards using the current model snapshot.
+        Each sample contains multiple generations for reinforcement learning optimization.
+
+        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 (defaults to config value, limited by prompts list length)
+            num_generations_per_prompt: Generations per prompt (defaults to config value)
+            concurrency: Concurrency limit (defaults to config value)
+
+        Returns:
+            List of dataset rows containing samples with messages and evaluations
+        """
+        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, generation_id):
+            async with semaphore:
+                prompt_input = prompts[prompt_id]
+                response_data = await agent_function(**prompt_input)
+                messages = response_data.get("messages", [])
+
+                try:
+                    traced_messages = self.tracer.get_current_message_history()
+                    if traced_messages:
+                        messages = traced_messages
+                except Exception as e:
+                    print(f"Warning: Failed to get message history from trace: {e}")
+                    pass
+
+                example = Example(
+                    input=prompt_input,
+                    messages=messages,
+                    actual_output=response_data,
+                )
+
+                scoring_results = self.judgment_client.run_evaluation(
+                    examples=[example],
+                    scorers=scorers,
+                    project_name=self.project_name,
+                    eval_run_name=f"training_step_{self.trainable_model.current_step}_prompt_{prompt_id}_gen_{generation_id}",
+                    show_url=False,
+                )
+
+                if scoring_results and scoring_results[0].scorers_data:
+                    reward = sum(
+                        scorer_data.score
+                        for scorer_data in scoring_results[0].scorers_data
+                    ) / len(scoring_results[0].scorers_data)
+                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 coro in asyncio.as_completed(coros):
+                result = await coro
+                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 run_reinforcement_learning(
+        self,
+        agent_function: Callable[[Any], Any],
+        scorers: List[Union[APIScorerConfig, BaseScorer]],
+        prompts: List[Any],
+    ) -> ModelConfig:
+        """
+        Run the iterative reinforcement learning fine-tuning loop.
+
+        This method performs multiple steps of reinforcement learning, where each step:
+        1. Advances to the appropriate model snapshot
+        2. Generates rollouts and computes rewards using scorers
+        3. Trains a new model using reinforcement learning
+        4. Waits for training completion
+
+        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 for inference and future training
+        """
+
+        _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,
+            "accelerator_count": self.config.accelerator_count,
+            "accelerator_type": self.config.accelerator_type,
+            "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
+
+                    time.sleep(10)
+                    job = job.get()
+                    if job is None:
+                        raise JudgmentAPIError(
+                            "Training job was deleted while waiting for completion"
+                        )
+
+            _print_progress(
+                f"Training completed! New model: {job.output_model}",
+                step_num,
+                self.config.num_steps,
+            )
+
+            dataset.delete()
+
+        _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[Union[APIScorerConfig, BaseScorer]],
+        prompts: List[Any],
+        rft_provider: Optional[str] = None,
+    ) -> ModelConfig:
+        """
+        Start the reinforcement learning fine-tuning process.
+
+        This is the main entry point for running the reinforcement learning 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
+            rft_provider: RFT provider to use for training. Currently only "fireworks" is supported.
+                         Support for other providers is planned for future releases.
+
+        Returns:
+            ModelConfig: Configuration of the trained model for future loading
+        """
+        try:
+            if rft_provider is not None:
+                self.config.rft_provider = rft_provider
+
+            return await self.run_reinforcement_learning(
+                agent_function, scorers, prompts
+            )
+        except JudgmentAPIError:
+            # Re-raise JudgmentAPIError as-is
+            raise
+        except Exception as e:
+            raise JudgmentAPIError(f"Training process failed: {str(e)}") from e

judgeval/data/evaluation_run.py

@@ -0,0 +1,104 @@
+from typing import List, Optional, Union
+from pydantic import field_validator, model_validator, Field
+from datetime import datetime, timezone
+import uuid
+
+from judgeval.data import Example
+from judgeval.scorers import BaseScorer, APIScorerConfig
+from judgeval.constants import ACCEPTABLE_MODELS
+from judgeval.data.judgment_types import EvaluationRunJudgmentType
+
+
+class EvaluationRun(EvaluationRunJudgmentType):
+    """
+    Stores example and evaluation scorers together for running an eval task
+
+    Args:
+        project_name (str): The name of the project the evaluation results belong to
+        eval_name (str): A name for this evaluation run
+        examples (List[Example]): The examples to evaluate
+        scorers (List[Union[BaseScorer, APIScorerConfig]]): A list of scorers to use for evaluation
+        model (str): The model used as a judge when using LLM as a Judge
+        metadata (Optional[Dict[str, Any]]): Additional metadata to include for this evaluation run, e.g. comments, dataset name, purpose, etc.
+    """
+
+    id: Optional[str] = Field(default_factory=lambda: str(uuid.uuid4()))
+    created_at: Optional[str] = Field(
+        default_factory=lambda: datetime.now(timezone.utc).isoformat()
+    )
+    custom_scorers: Optional[List[BaseScorer]] = None
+    judgment_scorers: Optional[List[APIScorerConfig]] = None
+    organization_id: Optional[str] = None
+
+    def __init__(
+        self,
+        scorers: Optional[List[Union[BaseScorer, APIScorerConfig]]] = None,
+        **kwargs,
+    ):
+        """
+        Initialize EvaluationRun with automatic scorer classification.
+
+        Args:
+            scorers: List of scorers that will be automatically sorted into custom_scorers or judgment_scorers
+            **kwargs: Other initialization arguments
+        """
+        if scorers is not None:
+            # Automatically sort scorers into appropriate fields
+            custom_scorers = [s for s in scorers if isinstance(s, BaseScorer)]
+            judgment_scorers = [s for s in scorers if isinstance(s, APIScorerConfig)]
+
+            # Always set both fields as lists (even if empty) to satisfy validation
+            kwargs["custom_scorers"] = custom_scorers
+            kwargs["judgment_scorers"] = judgment_scorers
+
+        super().__init__(**kwargs)
+
+    def model_dump(self, **kwargs):
+        data = super().model_dump(**kwargs)
+        data["custom_scorers"] = [s.model_dump() for s in self.custom_scorers]
+        data["judgment_scorers"] = [s.model_dump() for s in self.judgment_scorers]
+        data["examples"] = [example.model_dump() for example in self.examples]
+
+        return data
+
+    @field_validator("examples")
+    def validate_examples(cls, v):
+        if not v:
+            raise ValueError("Examples cannot be empty.")
+        for item in v:
+            if not isinstance(item, Example):
+                raise ValueError(f"Item of type {type(item)} is not a Example")
+        return v
+
+    @model_validator(mode="after")
+    @classmethod
+    def validate_scorer_lists(cls, values):
+        custom_scorers = values.custom_scorers
+        judgment_scorers = values.judgment_scorers
+
+        # Check that both lists are not empty
+        if not custom_scorers and not judgment_scorers:
+            raise ValueError(
+                "At least one of custom_scorers or judgment_scorers must be provided."
+            )
+
+        # Check that only one list is filled
+        if custom_scorers and judgment_scorers:
+            raise ValueError(
+                "Only one of custom_scorers or judgment_scorers can be provided, not both."
+            )
+
+        return values
+
+    @field_validator("model")
+    def validate_model(cls, v, values):
+        if not v:
+            raise ValueError("Model cannot be empty.")
+
+        # Check if model is string or list of strings
+        if isinstance(v, str):
+            if v not in ACCEPTABLE_MODELS:
+                raise ValueError(
+                    f"Model name {v} not recognized. Please select a valid model name.)"
+                )
+            return v

judgeval/data/judgment_types.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  openapi_new.json
-#   timestamp: 2025-08-01T22:19:19+00:00
+#   timestamp: 2025-08-08T18:50:51+00:00
 
 from __future__ import annotations
 
@@ -51,6 +51,31 @@ class ScorerConfigJudgmentType(BaseModel):
     kwargs: Annotated[Optional[Dict[str, Any]], Field(title="Kwargs")] = None
 
 
+class BaseScorerJudgmentType(BaseModel):
+    score_type: Annotated[str, Field(title="Score Type")]
+    threshold: Annotated[Optional[float], Field(title="Threshold")] = 0.5
+    name: Annotated[Optional[str], Field(title="Name")] = None
+    class_name: Annotated[Optional[str], Field(title="Class Name")] = None
+    score: Annotated[Optional[float], Field(title="Score")] = None
+    score_breakdown: Annotated[
+        Optional[Dict[str, Any]], Field(title="Score Breakdown")
+    ] = None
+    reason: Annotated[Optional[str], Field(title="Reason")] = ""
+    using_native_model: Annotated[Optional[bool], Field(title="Using Native Model")] = (
+        None
+    )
+    success: Annotated[Optional[bool], Field(title="Success")] = None
+    model: Annotated[Optional[str], Field(title="Model")] = None
+    model_client: Annotated[Any, Field(title="Model Client")] = None
+    strict_mode: Annotated[Optional[bool], Field(title="Strict Mode")] = False
+    error: Annotated[Optional[str], Field(title="Error")] = None
+    additional_metadata: Annotated[
+        Optional[Dict[str, Any]], Field(title="Additional Metadata")
+    ] = None
+    user: Annotated[Optional[str], Field(title="User")] = None
+    server_hosted: Annotated[Optional[bool], Field(title="Server Hosted")] = False
+
+
 class TraceUsageJudgmentType(BaseModel):
     prompt_tokens: Annotated[Optional[int], Field(title="Prompt Tokens")] = None
     completion_tokens: Annotated[Optional[int], Field(title="Completion Tokens")] = None
@@ -90,16 +115,21 @@ class HTTPValidationErrorJudgmentType(BaseModel):
     ] = None
 
 
-class JudgmentEvalJudgmentType(BaseModel):
+class EvaluationRunJudgmentType(BaseModel):
+    id: Annotated[Optional[str], Field(title="Id")] = None
     project_name: Annotated[Optional[str], Field(title="Project Name")] = None
     eval_name: Annotated[Optional[str], Field(title="Eval Name")] = None
     examples: Annotated[List[ExampleJudgmentType], Field(title="Examples")]
-    scorers: Annotated[List[ScorerConfigJudgmentType], Field(title="Scorers")]
+    custom_scorers: Annotated[
+        Optional[List[BaseScorerJudgmentType]], Field(title="Custom Scorers")
+    ] = Field(default_factory=list)
+    judgment_scorers: Annotated[
+        Optional[List[ScorerConfigJudgmentType]], Field(title="Judgment Scorers")
+    ] = Field(default_factory=list)
     model: Annotated[str, Field(title="Model")]
-    append: Annotated[Optional[bool], Field(title="Append")] = False
-    override: Annotated[Optional[bool], Field(title="Override")] = False
     trace_span_id: Annotated[Optional[str], Field(title="Trace Span Id")] = None
     trace_id: Annotated[Optional[str], Field(title="Trace Id")] = None
+    created_at: Annotated[Optional[str], Field(title="Created At")] = None
 
 
 class TraceSpanJudgmentType(BaseModel):
@@ -123,6 +153,7 @@ class TraceSpanJudgmentType(BaseModel):
     ] = None
     has_evaluation: Annotated[Optional[bool], Field(title="Has Evaluation")] = False
     agent_name: Annotated[Optional[str], Field(title="Agent Name")] = None
+    class_name: Annotated[Optional[str], Field(title="Class Name")] = None
     state_before: Annotated[Optional[Dict[str, Any]], Field(title="State Before")] = (
         None
     )
@@ -172,8 +203,6 @@ class TraceRunJudgmentType(BaseModel):
     traces: Annotated[List[TraceJudgmentType], Field(title="Traces")]
     scorers: Annotated[List[ScorerConfigJudgmentType], Field(title="Scorers")]
     model: Annotated[str, Field(title="Model")]
-    append: Annotated[Optional[bool], Field(title="Append")] = False
-    override: Annotated[Optional[bool], Field(title="Override")] = False
     trace_span_id: Annotated[Optional[str], Field(title="Trace Span Id")] = None
     tools: Annotated[Optional[List[Dict[str, Any]]], Field(title="Tools")] = None
 
@@ -181,5 +210,5 @@ class TraceRunJudgmentType(BaseModel):
 class EvalResultsJudgmentType(BaseModel):
     results: Annotated[List[ScoringResultJudgmentType], Field(title="Results")]
     run: Annotated[
-        Union[TraceRunJudgmentType, JudgmentEvalJudgmentType], Field(title="Run")
+        Union[TraceRunJudgmentType, EvaluationRunJudgmentType], Field(title="Run")
     ]

judgeval/data/trace.py

@@ -32,6 +32,7 @@ class TraceSpan(TraceSpanJudgmentType):
             "usage": self.usage.model_dump() if self.usage else None,
             "has_evaluation": self.has_evaluation,
             "agent_name": self.agent_name,
+            "class_name": self.class_name,
             "state_before": self.state_before,
             "state_after": self.state_after,
             "additional_metadata": json_encoder(self.additional_metadata),

judgeval/data/trace_run.py

@@ -29,8 +29,6 @@ class TraceRun(BaseModel):
     scorers: List[Union[APIScorerConfig, BaseScorer]]
     model: Optional[str] = DEFAULT_GPT_MODEL
     trace_span_id: Optional[str] = None
-    append: Optional[bool] = False
-    override: Optional[bool] = False
     rules: Optional[List[Rule]] = None
     tools: Optional[List[Dict[str, Any]]] = None
 

judgeval/integrations/langgraph.py

@@ -133,7 +133,8 @@ class JudgevalCallbackHandler(BaseCallbackHandler):
         inputs: Optional[Dict[str, Any]] = None,
     ) -> None:
         """Start tracking a span, ensuring trace client exists"""
-
+        if name.startswith("__") and name.endswith("__"):
+            return
         start_time = time.time()
         span_id = str(uuid.uuid4())
         parent_span_id: Optional[str] = None