pr comments

Author: justinsheu

src/judgeval/tracer/__init__.py

@@ -50,6 +50,7 @@
 from judgeval.utils.serialize import safe_serialize
 from judgeval.version import get_version
 from judgeval.warnings import JudgmentWarning
+from judgeval.exceptions import JudgmentRuntimeError
 
 from judgeval.tracer.exporters import JudgmentSpanExporter
 from judgeval.tracer.keys import AttributeKeys, ResourceKeys
@@ -490,7 +491,7 @@ def _create_agent_context(args) -> Token:
 
             if identifier:
                 if not class_name or not instance or not isinstance(instance, object):
-                    raise Exception(
+                    raise JudgmentRuntimeError(
                         "'identifier' is set but no class name or instance is available. 'identifier' can only be specified when using the agent() decorator on a class method."
                     )
                 if (
@@ -501,7 +502,7 @@ def _create_agent_context(args) -> Token:
                     instance_name = str(getattr(instance, identifier))
                     agent_context["instance_name"] = instance_name
                 else:
-                    raise Exception(
+                    raise JudgmentRuntimeError(
                         f"Attribute {identifier} does not exist for {class_name}. Check your agent() decorator."
                     )
 

src/judgeval/tracer/keys.py

@@ -38,7 +38,6 @@ class AttributeKeys:
     GEN_AI_RESPONSE_FINISH_REASONS = gen_ai_attributes.GEN_AI_RESPONSE_FINISH_REASONS
 
     # GenAI-specific attributes (custom namespace)
-    GEN_AI_USAGE_TOTAL_TOKENS = "gen_ai.usage.total_tokens"
     GEN_AI_USAGE_TOTAL_COST = "gen_ai.usage.total_cost"
 
 

src/judgeval/tracer/llm/__init__.py

@@ -76,11 +76,6 @@ def wrapper(*args, **kwargs):
                                 AttributeKeys.GEN_AI_USAGE_COMPLETION_TOKENS,
                                 usage.completion_tokens,
                             )
-                        if usage.total_tokens:
-                            span.set_attribute(
-                                AttributeKeys.GEN_AI_USAGE_TOTAL_TOKENS,
-                                usage.total_tokens,
-                            )
                         if usage.total_cost_usd:
                             span.set_attribute(
                                 AttributeKeys.GEN_AI_USAGE_TOTAL_COST,
@@ -123,11 +118,6 @@ async def wrapper(*args, **kwargs):
                                 AttributeKeys.GEN_AI_USAGE_COMPLETION_TOKENS,
                                 usage.completion_tokens,
                             )
-                        if usage.total_tokens:
-                            span.set_attribute(
-                                AttributeKeys.GEN_AI_USAGE_TOTAL_TOKENS,
-                                usage.total_tokens,
-                            )
                         if usage.total_cost_usd:
                             span.set_attribute(
                                 AttributeKeys.GEN_AI_USAGE_TOTAL_COST,