judgeval 0.16.8__tar.gz → 0.17.0__tar.gz

{judgeval-0.16.8 → judgeval-0.17.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: judgeval
-Version: 0.16.8
+Version: 0.17.0
 Summary: Judgeval Package
 Project-URL: Homepage, https://github.com/JudgmentLabs/judgeval
 Project-URL: Issues, https://github.com/JudgmentLabs/judgeval/issues
@@ -63,8 +63,7 @@ Judgeval's agent monitoring infra provides a simple harness for integrating GRPO
 await trainer.train(
     agent_function=your_agent_function,  # entry point to your agent
     scorers=[RewardScorer()],  # Custom scorer you define based on task criteria, acts as reward
-    prompts=training_prompts,  # Tasks
-    rft_provider="fireworks"
+    prompts=training_prompts  # Tasks
 )
 ```
 

{judgeval-0.16.8 → judgeval-0.17.0}/README.md

@@ -36,8 +36,7 @@ Judgeval's agent monitoring infra provides a simple harness for integrating GRPO
 await trainer.train(
     agent_function=your_agent_function,  # entry point to your agent
     scorers=[RewardScorer()],  # Custom scorer you define based on task criteria, acts as reward
-    prompts=training_prompts,  # Tasks
-    rft_provider="fireworks"
+    prompts=training_prompts  # Tasks
 )
 ```
 

{judgeval-0.16.8 → judgeval-0.17.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "judgeval"
-version = "0.16.8"
+version = "0.17.0"
 authors = [
     { name = "Andrew Li", email = "andrew@judgmentlabs.ai" },
     { name = "Alex Shan", email = "alex@judgmentlabs.ai" },

judgeval-0.17.0/src/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-0.17.0/src/judgeval/trainer/base_trainer.py

@@ -0,0 +1,117 @@
+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) -> List[Dict[str, str]]:
+        """
+        Extract message history from spans for training purposes.
+
+        Returns:
+            List of message dictionaries with 'role' and 'content' keys
+        """
+        pass

{judgeval-0.16.8 → judgeval-0.17.0}/src/judgeval/trainer/config.py

@@ -16,7 +16,7 @@ class TrainerConfig:
     user_id: str
     model_id: str
     base_model_name: str = "qwen2p5-7b-instruct"
-    rft_provider: str = "fireworks"
+    rft_provider: str = "fireworks"  # Supported: "fireworks", "verifiers" (future)
     num_steps: int = 5
     num_generations_per_prompt: int = 4
     num_prompts_per_step: int = 4

judgeval-0.16.8/src/judgeval/trainer/trainer.py → judgeval-0.17.0/src/judgeval/trainer/fireworks_trainer.py

@@ -1,9 +1,9 @@
 import asyncio
 import json
-import time
 from typing import Optional, Callable, Any, List, Union, Dict
 from fireworks import Dataset  # type: ignore[import-not-found]
 from .config import TrainerConfig, ModelConfig
+from .base_trainer import BaseTrainer
 from .trainable_model import TrainableModel
 from judgeval.tracer import Tracer
 from judgeval.tracer.exporters.store import SpanStore
@@ -16,12 +16,12 @@ from .console import _spinner_progress, _print_progress, _print_progress_update
 from judgeval.exceptions import JudgmentRuntimeError
 
 
-class JudgmentTrainer:
+class FireworksTrainer(BaseTrainer):
     """
-    A reinforcement learning trainer for Judgment models using Fine-Tuning.
+    Fireworks AI implementation of the training provider.
 
-    This class handles the iterative training process where models are improved
-    through reinforcement learning fine-tuning based on generated rollouts and rewards.
+    This trainer uses Fireworks AI's infrastructure for reinforcement learning
+    fine-tuning (RFT) of language models.
     """
 
     def __init__(
@@ -32,26 +32,23 @@ class JudgmentTrainer:
         project_name: Optional[str] = None,
     ):
         """
-        Initialize the JudgmentTrainer.
+        Initialize the FireworksTrainer.
 
         Args:
-            config: TrainerConfig instance with training parameters. If None, uses default config.
-            tracer: Optional tracer for observability
-            trainable_model: Optional trainable model instance
+            config: TrainerConfig instance with training parameters
+            trainable_model: TrainableModel instance for Fireworks training
+            tracer: Tracer for observability
             project_name: Project name for organizing training runs and evaluations
         """
         try:
-            self.config = config
-            self.tracer = tracer
-            self.project_name = project_name or "judgment_training"
-            self.trainable_model = trainable_model
+            super().__init__(config, trainable_model, tracer, project_name)
 
             self.judgment_client = JudgmentClient()
             self.span_store = SpanStore()
             self.span_exporter = InMemorySpanExporter(self.span_store)
         except Exception as e:
             raise JudgmentRuntimeError(
-                f"Failed to initialize JudgmentTrainer: {str(e)}"
+                f"Failed to initialize FireworksTrainer: {str(e)}"
             ) from e
 
     def _extract_message_history_from_spans(self) -> List[Dict[str, str]]:
@@ -121,22 +118,7 @@ class JudgmentTrainer:
                         pass
                     messages.append({"role": "assistant", "content": content})
 
-            elif span_type == "user":
-                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})
-
-            elif span_type == "tool":
+            elif span_type in ("user", "tool"):
                 output = span_attributes.get(AttributeKeys.JUDGMENT_OUTPUT)
                 if output is not None:
                     content = str(output)
@@ -347,7 +329,7 @@ class JudgmentTrainer:
                             _print_progress_update(f"Training job: {current_state}")
                         last_state = current_state
 
-                    time.sleep(10)
+                    await asyncio.sleep(10)
                     job = job.get()
                     if job is None:
                         raise JudgmentRuntimeError(
@@ -374,7 +356,6 @@ class JudgmentTrainer:
         agent_function: Callable[[Any], Any],
         scorers: List[Union[ExampleAPIScorerConfig, ExampleScorer]],
         prompts: List[Any],
-        rft_provider: Optional[str] = None,
     ) -> ModelConfig:
         """
         Start the reinforcement learning fine-tuning process.
@@ -385,21 +366,16 @@ class JudgmentTrainer:
             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 JudgmentRuntimeError:
-            # Re-raise JudgmentAPIError as-is
+            # Re-raise JudgmentRuntimeError as-is
             raise
         except Exception as e:
             raise JudgmentRuntimeError(f"Training process failed: {str(e)}") from e

judgeval-0.17.0/src/judgeval/trainer/trainer.py

@@ -0,0 +1,70 @@
+from typing import Optional
+from .config import TrainerConfig
+from .base_trainer import BaseTrainer
+from .fireworks_trainer import FireworksTrainer
+from .trainable_model import TrainableModel
+from judgeval.tracer import Tracer
+from judgeval.exceptions import JudgmentRuntimeError
+
+
+def JudgmentTrainer(
+    config: TrainerConfig,
+    trainable_model: TrainableModel,
+    tracer: Tracer,
+    project_name: Optional[str] = None,
+) -> BaseTrainer:
+    """
+    Factory function for creating reinforcement learning trainers.
+
+    This factory creates and returns provider-specific trainer implementations
+    (FireworksTrainer, VerifiersTrainer, etc.) based on the configured RFT provider.
+
+    The factory pattern allows for easy extension to support multiple training
+    providers without changing the client-facing API.
+
+    Example:
+        config = TrainerConfig(
+            deployment_id="my-deployment",
+            user_id="my-user",
+            model_id="my-model",
+            rft_provider="fireworks"  # or "verifiers" in the future
+        )
+
+        # User creates and configures the trainable model
+        trainable_model = TrainableModel(config)
+        tracer = Tracer()
+
+        # JudgmentTrainer automatically creates the appropriate provider-specific trainer
+        trainer = JudgmentTrainer(config, trainable_model, tracer)
+
+        # The returned trainer implements the BaseTrainer interface
+        model_config = await trainer.train(agent_function, scorers, prompts)
+
+        Args:
+        config: TrainerConfig instance with training parameters including rft_provider
+        trainable_model: Provider-specific trainable model instance (e.g., TrainableModel for Fireworks)
+        tracer: Tracer for observability
+        project_name: Project name for organizing training runs and evaluations
+
+        Returns:
+        Provider-specific trainer instance (FireworksTrainer, etc.) that implements
+        the BaseTrainer interface
+
+    Raises:
+        JudgmentRuntimeError: If the specified provider is not supported
+    """
+    provider = config.rft_provider.lower()
+
+    if provider == "fireworks":
+        return FireworksTrainer(config, trainable_model, tracer, project_name)
+    elif provider == "verifiers":
+        # Placeholder for future implementation
+        raise JudgmentRuntimeError(
+            "Verifiers provider is not yet implemented. "
+            "Currently supported providers: 'fireworks'"
+        )
+    else:
+        raise JudgmentRuntimeError(
+            f"Unsupported RFT provider: '{config.rft_provider}'. "
+            f"Currently supported providers: 'fireworks'"
+        )

{judgeval-0.16.8 → judgeval-0.17.0}/src/judgeval/version.py

@@ -1,4 +1,4 @@
-__version__ = "0.16.8"
+__version__ = "0.17.0"
 
 
 def get_version() -> str:

judgeval-0.16.8/src/judgeval/trainer/__init__.py

@@ -1,5 +0,0 @@
-from judgeval.trainer.trainer import JudgmentTrainer
-from judgeval.trainer.config import TrainerConfig, ModelConfig
-from judgeval.trainer.trainable_model import TrainableModel
-
-__all__ = ["JudgmentTrainer", "TrainerConfig", "ModelConfig", "TrainableModel"]