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

judgeval/scorers/judgeval_scorers/api_scorers/prompt_scorer.py

@@ -40,12 +40,6 @@ def push_prompt_scorer(
             }
         )
     except JudgmentAPIError as e:
-        if e.status_code == 500:
-            raise JudgmentAPIError(
-                status_code=e.status_code,
-                detail=f"The server is temporarily unavailable. Please try your request again in a few moments. Error details: {e.detail}",
-                response=e.response,
-            )
         raise JudgmentAPIError(
             status_code=e.status_code,
             detail=f"Failed to save prompt scorer: {e.detail}",
@@ -75,12 +69,6 @@ def fetch_prompt_scorer(
             scorer_config.pop("updated_at")
             return scorer_config
     except JudgmentAPIError as e:
-        if e.status_code == 500:
-            raise JudgmentAPIError(
-                status_code=e.status_code,
-                detail=f"The server is temporarily unavailable. Please try your request again in a few moments. Error details: {e.detail}",
-                response=e.response,
-            )
         raise JudgmentAPIError(
             status_code=e.status_code,
             detail=f"Failed to fetch prompt scorer '{name}': {e.detail}",

judgeval/tracer/__init__.py

@@ -71,6 +71,7 @@ from judgeval.tracer.processors import (
     NoOpJudgmentSpanProcessor,
 )
 from judgeval.tracer.utils import set_span_attribute, TraceScorerConfig
+from judgeval.utils.project import _resolve_project_id
 
 C = TypeVar("C", bound=Callable)
 Cls = TypeVar("Cls", bound=Type)
@@ -101,6 +102,7 @@ class Tracer(metaclass=SingletonMeta):
         "judgment_processor",
         "tracer",
         "agent_context",
+        "customer_id",
         "_initialized",
     )
 
@@ -114,6 +116,7 @@ class Tracer(metaclass=SingletonMeta):
     judgment_processor: JudgmentSpanProcessor
     tracer: ABCTracer
     agent_context: ContextVar[Optional[AgentContext]]
+    customer_id: ContextVar[Optional[str]]
     _initialized: bool
 
     def __init__(
@@ -131,6 +134,7 @@ class Tracer(metaclass=SingletonMeta):
         if not hasattr(self, "_initialized"):
             self._initialized = False
             self.agent_context = ContextVar("current_agent_context", default=None)
+            self.customer_id = ContextVar("current_customer_id", default=None)
 
             self.project_name = project_name
             self.api_key = expect_api_key(api_key or JUDGMENT_API_KEY)
@@ -155,7 +159,7 @@ class Tracer(metaclass=SingletonMeta):
 
         self.judgment_processor = NoOpJudgmentSpanProcessor()
         if self.enable_monitoring:
-            project_id = Tracer._resolve_project_id(
+            project_id = _resolve_project_id(
                 self.project_name, self.api_key, self.organization_id
             )
             if project_id:
@@ -224,20 +228,6 @@ class Tracer(metaclass=SingletonMeta):
             resource_attributes=resource_attributes,
         )
 
-    @dont_throw
-    @functools.lru_cache(maxsize=64)
-    @staticmethod
-    def _resolve_project_id(
-        project_name: str, api_key: str, organization_id: str
-    ) -> str:
-        """Resolve project_id from project_name using the API."""
-        client = JudgmentSyncClient(
-            api_key=api_key,
-            organization_id=organization_id,
-        )
-        response = client.projects_resolve({"project_name": project_name})
-        return response["project_id"]
-
     def get_current_span(self):
         return get_current_span()
 
@@ -247,17 +237,50 @@ class Tracer(metaclass=SingletonMeta):
     def get_current_agent_context(self):
         return self.agent_context
 
+    def get_current_customer_context(self):
+        return self.customer_id
+
     def get_span_processor(self) -> JudgmentSpanProcessor:
         """Get the internal span processor of this tracer instance."""
         return self.judgment_processor
 
     def set_customer_id(self, customer_id: str) -> None:
+        if not customer_id:
+            judgeval_logger.warning("Customer ID is empty, skipping.")
+            return
+
         span = self.get_current_span()
+
+        if not span or not span.is_recording():
+            judgeval_logger.warning(
+                "No active span found. Customer ID will not be set."
+            )
+            return
+
+        if self.get_current_customer_context().get():
+            judgeval_logger.warning("Customer ID is already set, skipping.")
+            return
+
         if span and span.is_recording():
             set_span_attribute(span, AttributeKeys.JUDGMENT_CUSTOMER_ID, customer_id)
+            self.get_current_customer_context().set(customer_id)
+
+            self.get_span_processor().set_internal_attribute(
+                span_context=span.get_span_context(),
+                key=InternalAttributeKeys.IS_CUSTOMER_CONTEXT_OWNER,
+                value=True,
+            )
+
+    def _maybe_clear_customer_context(self, span: Span) -> None:
+        if self.get_span_processor().get_internal_attribute(
+            span_context=span.get_span_context(),
+            key=InternalAttributeKeys.IS_CUSTOMER_CONTEXT_OWNER,
+            default=False,
+        ):
+            self.get_current_customer_context().set(None)
 
     @dont_throw
-    def add_agent_attributes_to_span(self, span):
+    def _add_agent_attributes_to_span(self, span):
         """Add agent ID, class name, and instance name to span if they exist in context"""
         current_agent_context = self.agent_context.get()
         if not current_agent_context:
@@ -289,7 +312,7 @@ class Tracer(metaclass=SingletonMeta):
         current_agent_context["is_agent_entry_point"] = False
 
     @dont_throw
-    def record_instance_state(self, record_point: Literal["before", "after"], span):
+    def _record_instance_state(self, record_point: Literal["before", "after"], span):
         current_agent_context = self.agent_context.get()
 
         if current_agent_context and current_agent_context.get("track_state"):
@@ -318,6 +341,17 @@ class Tracer(metaclass=SingletonMeta):
                 safe_serialize(attributes),
             )
 
+    @dont_throw
+    def _add_customer_id_to_span(self, span):
+        customer_id = self.get_current_customer_context().get()
+        if customer_id:
+            set_span_attribute(span, AttributeKeys.JUDGMENT_CUSTOMER_ID, customer_id)
+
+    @dont_throw
+    def _inject_judgment_context(self, span):
+        self._add_agent_attributes_to_span(span)
+        self._add_customer_id_to_span(span)
+
     def _set_pending_trace_eval(
         self,
         span: Span,
@@ -398,7 +432,7 @@ class Tracer(metaclass=SingletonMeta):
                 with sync_span_context(
                     self, yield_span_name, yield_attributes, disable_partial_emit=True
                 ) as yield_span:
-                    self.add_agent_attributes_to_span(yield_span)
+                    self._inject_judgment_context(yield_span)
 
                     try:
                         value = next(generator)
@@ -442,7 +476,7 @@ class Tracer(metaclass=SingletonMeta):
                 async with async_span_context(
                     self, yield_span_name, yield_attributes, disable_partial_emit=True
                 ) as yield_span:
-                    self.add_agent_attributes_to_span(yield_span)
+                    self._inject_judgment_context(yield_span)
 
                     try:
                         value = await async_generator.__anext__()
@@ -484,8 +518,8 @@ class Tracer(metaclass=SingletonMeta):
         def wrapper(*args, **kwargs):
             n = name or f.__qualname__
             with sync_span_context(self, n, attributes) as span:
-                self.add_agent_attributes_to_span(span)
-                self.record_instance_state("before", span)
+                self._inject_judgment_context(span)
+                self._record_instance_state("before", span)
                 try:
                     set_span_attribute(
                         span,
@@ -502,13 +536,14 @@ class Tracer(metaclass=SingletonMeta):
                 except Exception as user_exc:
                     span.record_exception(user_exc)
                     span.set_status(Status(StatusCode.ERROR, str(user_exc)))
+                    self._maybe_clear_customer_context(span)
                     raise
 
                 if inspect.isgenerator(result):
                     set_span_attribute(
                         span, AttributeKeys.JUDGMENT_OUTPUT, "<generator>"
                     )
-                    self.record_instance_state("after", span)
+                    self._record_instance_state("after", span)
                     return self._create_traced_sync_generator(
                         result, span, n, attributes
                     )
@@ -516,7 +551,8 @@ class Tracer(metaclass=SingletonMeta):
                     set_span_attribute(
                         span, AttributeKeys.JUDGMENT_OUTPUT, safe_serialize(result)
                     )
-                    self.record_instance_state("after", span)
+                    self._record_instance_state("after", span)
+                    self._maybe_clear_customer_context(span)
                     return result
 
         return wrapper
@@ -535,8 +571,8 @@ class Tracer(metaclass=SingletonMeta):
             n = name or f.__qualname__
 
             with sync_span_context(self, n, attributes) as main_span:
-                self.add_agent_attributes_to_span(main_span)
-                self.record_instance_state("before", main_span)
+                self._inject_judgment_context(main_span)
+                self._record_instance_state("before", main_span)
 
                 try:
                     set_span_attribute(
@@ -556,7 +592,7 @@ class Tracer(metaclass=SingletonMeta):
                     set_span_attribute(
                         main_span, AttributeKeys.JUDGMENT_OUTPUT, "<generator>"
                     )
-                    self.record_instance_state("after", main_span)
+                    self._record_instance_state("after", main_span)
 
                     return self._create_traced_sync_generator(
                         generator, main_span, n, attributes
@@ -586,8 +622,8 @@ class Tracer(metaclass=SingletonMeta):
         async def wrapper(*args, **kwargs):
             n = name or f.__qualname__
             async with async_span_context(self, n, attributes) as span:
-                self.add_agent_attributes_to_span(span)
-                self.record_instance_state("before", span)
+                self._inject_judgment_context(span)
+                self._record_instance_state("before", span)
                 try:
                     set_span_attribute(
                         span,
@@ -604,13 +640,14 @@ class Tracer(metaclass=SingletonMeta):
                 except Exception as user_exc:
                     span.record_exception(user_exc)
                     span.set_status(Status(StatusCode.ERROR, str(user_exc)))
+                    self._maybe_clear_customer_context(span)
                     raise
 
                 if inspect.isasyncgen(result):
                     set_span_attribute(
                         span, AttributeKeys.JUDGMENT_OUTPUT, "<async_generator>"
                     )
-                    self.record_instance_state("after", span)
+                    self._record_instance_state("after", span)
                     return self._create_traced_async_generator(
                         result, span, n, attributes
                     )
@@ -618,7 +655,8 @@ class Tracer(metaclass=SingletonMeta):
                     set_span_attribute(
                         span, AttributeKeys.JUDGMENT_OUTPUT, safe_serialize(result)
                     )
-                    self.record_instance_state("after", span)
+                    self._record_instance_state("after", span)
+                    self._maybe_clear_customer_context(span)
                     return result
 
         return wrapper
@@ -637,8 +675,8 @@ class Tracer(metaclass=SingletonMeta):
             n = name or f.__qualname__
 
             with sync_span_context(self, n, attributes) as main_span:
-                self.add_agent_attributes_to_span(main_span)
-                self.record_instance_state("before", main_span)
+                self._inject_judgment_context(main_span)
+                self._record_instance_state("before", main_span)
 
                 try:
                     set_span_attribute(
@@ -658,7 +696,7 @@ class Tracer(metaclass=SingletonMeta):
                     set_span_attribute(
                         main_span, AttributeKeys.JUDGMENT_OUTPUT, "<async_generator>"
                     )
-                    self.record_instance_state("after", main_span)
+                    self._record_instance_state("after", main_span)
 
                     return self._create_traced_async_generator(
                         async_generator, main_span, n, attributes

judgeval/tracer/exporters/store.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Dict
 
 from opentelemetry.sdk.trace import ReadableSpan
 
@@ -9,35 +9,51 @@ class ABCSpanStore(ABC):
     def add(self, *spans: ReadableSpan): ...
 
     @abstractmethod
-    def get(self, id: str) -> ReadableSpan: ...
+    def get_all(self) -> List[ReadableSpan]: ...
 
     @abstractmethod
-    def get_all(self) -> List[ReadableSpan]: ...
+    def get_by_trace_id(self, trace_id: str) -> List[ReadableSpan]: ...
+
+    @abstractmethod
+    def clear_trace(self, trace_id: str): ...
 
 
 class SpanStore(ABCSpanStore):
-    __slots__ = ("spans",)
+    __slots__ = ("_spans_by_trace",)
 
-    spans: List[ReadableSpan]
+    _spans_by_trace: Dict[str, List[ReadableSpan]]
 
     def __init__(self):
-        self.spans = []
+        self._spans_by_trace = {}
 
     def add(self, *spans: ReadableSpan):
-        self.spans.extend(spans)
-
-    def get(self, id: str) -> ReadableSpan:
-        for span in self.spans:
+        for span in spans:
             context = span.get_span_context()
             if context is None:
                 continue
-            if context.span_id == id:
-                return span
-
-        raise ValueError(f"Span with id {id} not found")
+            # Convert trace_id to hex string per OTEL spec
+            trace_id = format(context.trace_id, "032x")
+            if trace_id not in self._spans_by_trace:
+                self._spans_by_trace[trace_id] = []
+            self._spans_by_trace[trace_id].append(span)
 
     def get_all(self) -> List[ReadableSpan]:
-        return self.spans
+        all_spans = []
+        for spans in self._spans_by_trace.values():
+            all_spans.extend(spans)
+        return all_spans
+
+    def get_by_trace_id(self, trace_id: str) -> List[ReadableSpan]:
+        """Get all spans for a specific trace ID (32-char hex string)."""
+        return self._spans_by_trace.get(trace_id, [])
+
+    def clear_trace(self, trace_id: str):
+        """Clear all spans for a specific trace ID (32-char hex string)."""
+        if trace_id in self._spans_by_trace:
+            del self._spans_by_trace[trace_id]
 
     def __repr__(self) -> str:
-        return f"SpanStore(spans={self.spans})"
+        total_spans = sum(len(spans) for spans in self._spans_by_trace.values())
+        return (
+            f"SpanStore(traces={len(self._spans_by_trace)}, total_spans={total_spans})"
+        )

judgeval/tracer/keys.py

@@ -51,6 +51,7 @@ class InternalAttributeKeys(str, Enum):
 
     DISABLE_PARTIAL_EMIT = "disable_partial_emit"
     CANCELLED = "cancelled"
+    IS_CUSTOMER_CONTEXT_OWNER = "is_customer_context_owner"
 
 
 class ResourceKeys(str, Enum):

judgeval/tracer/llm/llm_anthropic/messages.py

@@ -89,7 +89,7 @@ def _wrap_non_streaming_sync(
         ctx["span"] = tracer.get_tracer().start_span(
             "ANTHROPIC_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)
         )
@@ -163,7 +163,7 @@ def _wrap_streaming_sync(
         ctx["span"] = tracer.get_tracer().start_span(
             "ANTHROPIC_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)
         )
@@ -273,7 +273,7 @@ def _wrap_non_streaming_async(
         ctx["span"] = tracer.get_tracer().start_span(
             "ANTHROPIC_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)
         )
@@ -348,7 +348,7 @@ def _wrap_streaming_async(
         ctx["span"] = tracer.get_tracer().start_span(
             "ANTHROPIC_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)
         )

judgeval/tracer/llm/llm_anthropic/messages_stream.py

@@ -37,7 +37,7 @@ def wrap_messages_stream_sync(tracer: Tracer, client: Anthropic) -> None:
         ctx["span"] = tracer.get_tracer().start_span(
             "ANTHROPIC_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)
         )
@@ -183,7 +183,7 @@ def wrap_messages_stream_async(tracer: Tracer, client: AsyncAnthropic) -> None:
         ctx["span"] = tracer.get_tracer().start_span(
             "ANTHROPIC_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)
         )

judgeval/tracer/llm/llm_google/generate_content.py

@@ -57,7 +57,7 @@ def wrap_generate_content_sync(tracer: Tracer, client: Client) -> None:
         ctx["span"] = tracer.get_tracer().start_span(
             "GOOGLE_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)
         )

judgeval/tracer/llm/llm_openai/beta_chat_completions.py

@@ -39,7 +39,7 @@ def _wrap_beta_non_streaming_sync(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )
@@ -122,7 +122,7 @@ def _wrap_beta_non_streaming_async(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )

judgeval/tracer/llm/llm_openai/chat_completions.py

@@ -62,7 +62,7 @@ def _wrap_non_streaming_sync(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )
@@ -139,7 +139,7 @@ def _wrap_streaming_sync(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )
@@ -258,7 +258,7 @@ def _wrap_non_streaming_async(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )
@@ -336,7 +336,7 @@ def _wrap_streaming_async(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )

judgeval/tracer/llm/llm_openai/responses.py

@@ -56,7 +56,7 @@ def _wrap_responses_non_streaming_sync(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )
@@ -131,7 +131,7 @@ def _wrap_responses_streaming_sync(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )
@@ -260,7 +260,7 @@ def _wrap_responses_non_streaming_async(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )
@@ -335,7 +335,7 @@ def _wrap_responses_streaming_async(
         ctx["span"] = tracer.get_tracer().start_span(
             "OPENAI_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)
         )

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)
         )
@@ -133,7 +133,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)
         )
@@ -239,7 +239,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)
         )
@@ -310,7 +310,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)
         )

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",
+]