token estimation langchain

Author: rohit-kadhe

src/langtrace_python_sdk/instrumentation/langchain_community/patch.py

@@ -23,6 +23,7 @@
 
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
+from langtrace.trace_attributes import SpanAttributes
 
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
@@ -70,6 +71,9 @@ 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]))
+
 
                 span.set_status(StatusCode.OK)
                 return result

src/langtrace_python_sdk/instrumentation/langchain_core/patch.py

@@ -30,6 +30,7 @@
 from importlib_metadata import version as v
 
 from langtrace_python_sdk.constants import LANGTRACE_SDK_NAME
+from langtrace.trace_attributes import SpanAttributes
 
 
 def generic_patch(
@@ -78,8 +79,16 @@ def traced_method(wrapped, instance, args, kwargs):
             try:
                 # Attempt to call the original method
                 result = wrapped(*args, **kwargs)
+
                 if trace_output:
                     span.set_attribute("langchain.outputs", to_json_string(result))
+                    if 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):
+                        span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0][0]))
+
+                    else:
+                        span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, instance.get_num_tokens(args[0][0].text))
 
                 span.set_status(StatusCode.OK)
                 return result
@@ -156,6 +165,7 @@ def traced_method(wrapped, instance, args, kwargs):
             try:
                 # Attempt to call the original method
                 result = wrapped(*args, **kwargs)
+
                 if trace_output:
                     outputs = {}
                     if isinstance(result, dict):