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

judgeval/tracer/llm/llm_together/chat_completions.py

@@ -63,7 +63,7 @@ def _wrap_non_streaming_sync(
         ctx["span"] = tracer.get_tracer().start_span(
             "TOGETHER_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
         )
-        tracer.add_agent_attributes_to_span(ctx["span"])
+        tracer._inject_judgment_context(ctx["span"])
         set_span_attribute(
             ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
         )
@@ -73,7 +73,7 @@ def _wrap_non_streaming_sync(
         )
         ctx["model_name"] = prefixed_model_name
         set_span_attribute(
-            ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, prefixed_model_name
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, prefixed_model_name
         )
 
     def post_hook(ctx: Dict[str, Any], result: ChatCompletionResponse) -> None:
@@ -90,10 +90,12 @@ def _wrap_non_streaming_sync(
                 result.usage
             )
             set_span_attribute(
-                span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
+                span,
+                AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS,
+                prompt_tokens,
             )
             set_span_attribute(
-                span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
+                span, AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS, completion_tokens
             )
             set_span_attribute(
                 span,
@@ -103,7 +105,7 @@ def _wrap_non_streaming_sync(
 
         set_span_attribute(
             span,
-            AttributeKeys.GEN_AI_RESPONSE_MODEL,
+            AttributeKeys.JUDGMENT_LLM_MODEL_NAME,
             ctx["model_name"],
         )
 
@@ -133,7 +135,7 @@ def _wrap_streaming_sync(
         ctx["span"] = tracer.get_tracer().start_span(
             "TOGETHER_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
         )
-        tracer.add_agent_attributes_to_span(ctx["span"])
+        tracer._inject_judgment_context(ctx["span"])
         set_span_attribute(
             ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
         )
@@ -143,7 +145,7 @@ def _wrap_streaming_sync(
         )
         ctx["model_name"] = prefixed_model_name
         set_span_attribute(
-            ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, prefixed_model_name
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, prefixed_model_name
         )
         ctx["accumulated_content"] = ""
 
@@ -171,10 +173,12 @@ def _wrap_streaming_sync(
                     chunk.usage
                 )
                 set_span_attribute(
-                    span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
+                    span,
+                    AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS,
+                    prompt_tokens,
                 )
                 set_span_attribute(
-                    span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
+                    span, AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS, completion_tokens
                 )
                 set_span_attribute(
                     span,
@@ -239,7 +243,7 @@ def _wrap_non_streaming_async(
         ctx["span"] = tracer.get_tracer().start_span(
             "TOGETHER_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
         )
-        tracer.add_agent_attributes_to_span(ctx["span"])
+        tracer._inject_judgment_context(ctx["span"])
         set_span_attribute(
             ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
         )
@@ -249,7 +253,7 @@ def _wrap_non_streaming_async(
         )
         ctx["model_name"] = prefixed_model_name
         set_span_attribute(
-            ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, prefixed_model_name
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, prefixed_model_name
         )
 
     def post_hook(ctx: Dict[str, Any], result: ChatCompletionResponse) -> None:
@@ -266,10 +270,12 @@ def _wrap_non_streaming_async(
                 result.usage
             )
             set_span_attribute(
-                span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
+                span,
+                AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS,
+                prompt_tokens,
             )
             set_span_attribute(
-                span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
+                span, AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS, completion_tokens
             )
             set_span_attribute(
                 span,
@@ -279,7 +285,7 @@ def _wrap_non_streaming_async(
 
         set_span_attribute(
             span,
-            AttributeKeys.GEN_AI_RESPONSE_MODEL,
+            AttributeKeys.JUDGMENT_LLM_MODEL_NAME,
             ctx["model_name"],
         )
 
@@ -310,7 +316,7 @@ def _wrap_streaming_async(
         ctx["span"] = tracer.get_tracer().start_span(
             "TOGETHER_API_CALL", attributes={AttributeKeys.JUDGMENT_SPAN_KIND: "llm"}
         )
-        tracer.add_agent_attributes_to_span(ctx["span"])
+        tracer._inject_judgment_context(ctx["span"])
         set_span_attribute(
             ctx["span"], AttributeKeys.GEN_AI_PROMPT, safe_serialize(kwargs)
         )
@@ -320,7 +326,7 @@ def _wrap_streaming_async(
         )
         ctx["model_name"] = prefixed_model_name
         set_span_attribute(
-            ctx["span"], AttributeKeys.GEN_AI_REQUEST_MODEL, prefixed_model_name
+            ctx["span"], AttributeKeys.JUDGMENT_LLM_MODEL_NAME, prefixed_model_name
         )
         ctx["accumulated_content"] = ""
 
@@ -348,10 +354,12 @@ def _wrap_streaming_async(
                     chunk.usage
                 )
                 set_span_attribute(
-                    span, AttributeKeys.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens
+                    span,
+                    AttributeKeys.JUDGMENT_USAGE_NON_CACHED_INPUT_TOKENS,
+                    prompt_tokens,
                 )
                 set_span_attribute(
-                    span, AttributeKeys.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
+                    span, AttributeKeys.JUDGMENT_USAGE_OUTPUT_TOKENS, completion_tokens
                 )
                 set_span_attribute(
                     span,

judgeval/tracer/managers.py

@@ -16,6 +16,7 @@ def sync_span_context(
     name: str,
     span_attributes: Optional[Dict[str, str]] = None,
     disable_partial_emit: bool = False,
+    end_on_exit: bool = False,
 ):
     if span_attributes is None:
         span_attributes = {}
@@ -23,6 +24,7 @@ def sync_span_context(
     with tracer.get_tracer().start_as_current_span(
         name=name,
         attributes=span_attributes,
+        end_on_exit=end_on_exit,
     ) as span:
         if disable_partial_emit:
             tracer.judgment_processor.set_internal_attribute(
@@ -39,6 +41,7 @@ async def async_span_context(
     name: str,
     span_attributes: Optional[Dict[str, str]] = None,
     disable_partial_emit: bool = False,
+    end_on_exit: bool = False,
 ):
     if span_attributes is None:
         span_attributes = {}
@@ -46,6 +49,7 @@ async def async_span_context(
     with tracer.get_tracer().start_as_current_span(
         name=name,
         attributes=span_attributes,
+        end_on_exit=end_on_exit,
     ) as span:
         if disable_partial_emit:
             tracer.judgment_processor.set_internal_attribute(

judgeval/trainer/__init__.py

@@ -1,5 +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"]
+__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

@@ -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