diff --git a/src/langtrace_python_sdk/utils/llm.py b/src/langtrace_python_sdk/utils/llm.py index 0e2f5630..8edf8b9c 100644 --- a/src/langtrace_python_sdk/utils/llm.py +++ b/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, diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py index fd73702c..7c77efd6 100644 --- a/src/langtrace_python_sdk/version.py +++ b/src/langtrace_python_sdk/version.py @@ -1 +1 @@ -__version__ = "2.2.26" +__version__ = "2.2.27"