add token counts for when they are present

Author: rohit-kadhe

src/langtrace_python_sdk/instrumentation/langchain_community/patch.py

@@ -71,8 +71,15 @@ def traced_method(wrapped, instance, args, kwargs):
                 result = wrapped(*args, **kwargs)
                 if trace_output:
                     span.set_attribute("langchain.outputs", to_json_string(result))
-                    span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, instance.get_num_tokens(result))
-                    span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0]))
+
+                    prompt_tokens = instance.get_num_tokens(args[0])
+                    completion_tokens = instance.get_num_tokens(result)
+                    if hasattr(result, 'usage'):
+                        prompt_tokens = result.usage.prompt_tokens
+                        completion_tokens = result.usage.completion_tokens
+
+                    span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, prompt_tokens)
+                    span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, completion_tokens)
 
 
                 span.set_status(StatusCode.OK)

src/langtrace_python_sdk/instrumentation/langchain_core/patch.py

@@ -82,9 +82,15 @@ def traced_method(wrapped, instance, args, kwargs):
 
                 if trace_output:
                     span.set_attribute("langchain.outputs", to_json_string(result))
-                    if result.generations[0][0].text:
+                    if hasattr(result, 'usage'):
+                        prompt_tokens = result.usage.prompt_tokens
+                        completion_tokens = result.usage.completion_tokens
+                        span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, prompt_tokens)
+                        span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, completion_tokens)
+
+                    elif result.generations[0][0].text:
                         span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, instance.get_num_tokens(result.generations[0][0].text))
-                    if isinstance(args[0][0], str):
+                    elif isinstance(args[0][0], str):
                         span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0][0]))
 
                     else: