Merge pull request #303 from Scale3-Labs/rohit/S3EN-2755-bug-fix

Fix output tokens count 0 issue

Author: rohit-kadhe

src/langtrace_python_sdk/utils/llm.py

@@ -18,6 +18,7 @@
 from langtrace_python_sdk.utils import set_span_attribute
 from openai import NOT_GIVEN
 from tiktoken import get_encoding
+from tiktoken import get_encoding, list_encoding_names
 
 from langtrace_python_sdk.constants.instrumentation.common import (
     LANGTRACE_ADDITIONAL_SPAN_ATTRIBUTES_KEY,
@@ -30,6 +31,7 @@
 from opentelemetry import baggage
 from opentelemetry.trace import Span
 from opentelemetry.trace.status import StatusCode
+
 import os
 
 
@@ -253,6 +255,11 @@ def setup(self):
             self._span_started = True
 
     def cleanup(self):
+        if self.completion_tokens==0:
+            response_model = 'cl100k_base'
+            if self._response_model in list_encoding_names():
+                response_model = self._response_model
+            self.completion_tokens = estimate_tokens_using_tiktoken("".join(self.result_content), response_model)
         if self._span_started:
             set_span_attribute(
                 self.span,

src/langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "2.2.26"
+__version__ = "2.2.27"