diff --git a/src/langtrace_python_sdk/instrumentation/cohere/patch.py b/src/langtrace_python_sdk/instrumentation/cohere/patch.py index 8b9a8e53..524b14d6 100644 --- a/src/langtrace_python_sdk/instrumentation/cohere/patch.py +++ b/src/langtrace_python_sdk/instrumentation/cohere/patch.py @@ -87,25 +87,25 @@ def traced_method(wrapped, instance, args, kwargs): and result.meta.billed_units is not None ): usage = result.meta.billed_units - if usage is not None: - span.set_attribute( - SpanAttributes.LLM_USAGE_PROMPT_TOKENS, - usage.input_tokens or 0, - ) - span.set_attribute( - SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, - usage.output_tokens or 0, - ) - - span.set_attribute( - SpanAttributes.LLM_USAGE_TOTAL_TOKENS, - (usage.input_tokens or 0) + (usage.output_tokens or 0), - ) - - span.set_attribute( - "search_units", - usage.search_units or 0, - ) + input_tokens = int(usage.input_tokens) if usage.input_tokens else 0 + output_tokens = int(usage.output_tokens) if usage.output_tokens else 0 + span.set_attribute( + SpanAttributes.LLM_USAGE_PROMPT_TOKENS, + input_tokens, + ) + span.set_attribute( + SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, + output_tokens, + ) + span.set_attribute( + SpanAttributes.LLM_USAGE_TOTAL_TOKENS, + input_tokens + output_tokens, + ) + span.set_attribute( + "search_units", + int(usage.search_units) if usage.search_units else 0, + ) + span.set_status(StatusCode.OK) span.end() @@ -309,25 +309,25 @@ def traced_method(wrapped, instance, args, kwargs): and result.meta.billed_units is not None ): usage = result.meta.billed_units - if usage is not None: - span.set_attribute( - SpanAttributes.LLM_USAGE_PROMPT_TOKENS, - usage.input_tokens or 0, - ) - span.set_attribute( - SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, - usage.output_tokens or 0, - ) - - span.set_attribute( - SpanAttributes.LLM_USAGE_TOTAL_TOKENS, - (usage.input_tokens or 0) + (usage.output_tokens or 0), - ) - - span.set_attribute( - "search_units", - usage.search_units or 0, - ) + input_tokens = int(usage.input_tokens) if usage.input_tokens else 0 + output_tokens = int(usage.output_tokens) if usage.output_tokens else 0 + span.set_attribute( + SpanAttributes.LLM_USAGE_PROMPT_TOKENS, + input_tokens, + ) + span.set_attribute( + SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, + output_tokens, + ) + span.set_attribute( + SpanAttributes.LLM_USAGE_TOTAL_TOKENS, + input_tokens + output_tokens, + ) + span.set_attribute( + "search_units", + int(usage.search_units) if usage.search_units else 0, + ) + span.set_status(StatusCode.OK) span.end() return result @@ -419,10 +419,12 @@ def traced_method(wrapped, instance, args, kwargs): if (hasattr(result.usage, "billed_units") and result.usage.billed_units is not None): usage = result.usage.billed_units + input_tokens = int(usage.input_tokens) if usage.input_tokens else 0 + output_tokens = int(usage.output_tokens) if usage.output_tokens else 0 for metric, value in { - "input": usage.input_tokens or 0, - "output": usage.output_tokens or 0, - "total": (usage.input_tokens or 0) + (usage.output_tokens or 0), + "input": input_tokens, + "output": output_tokens, + "total": input_tokens + output_tokens, }.items(): span.set_attribute( f"gen_ai.usage.{metric}_tokens", @@ -571,26 +573,27 @@ def traced_method(wrapped, instance, args, kwargs): and response.meta.billed_units is not None ): usage = response.meta.billed_units - if usage is not None: - span.set_attribute( - SpanAttributes.LLM_USAGE_PROMPT_TOKENS, - usage.input_tokens or 0, - ) - span.set_attribute( - SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, - usage.output_tokens or 0, - ) - + input_tokens = int(usage.input_tokens) if usage.input_tokens else 0 + output_tokens = int(usage.output_tokens) if usage.output_tokens else 0 + span.set_attribute( + SpanAttributes.LLM_USAGE_PROMPT_TOKENS, + input_tokens, + ) + span.set_attribute( + SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, + output_tokens, + ) + span.set_attribute( + SpanAttributes.LLM_USAGE_TOTAL_TOKENS, + input_tokens + output_tokens, + ) + + if usage.search_units is not None: span.set_attribute( - SpanAttributes.LLM_USAGE_TOTAL_TOKENS, - (usage.input_tokens or 0) - + (usage.output_tokens or 0), + "search_units", + int(usage.search_units) if usage.search_units else 0, ) - if usage.search_units is not None: - span.set_attribute( - "search_units", - usage.search_units or 0, - ) + yield event finally: diff --git a/src/langtrace_python_sdk/instrumentation/openai/patch.py b/src/langtrace_python_sdk/instrumentation/openai/patch.py index 5d642a71..ed2592ad 100644 --- a/src/langtrace_python_sdk/instrumentation/openai/patch.py +++ b/src/langtrace_python_sdk/instrumentation/openai/patch.py @@ -114,6 +114,9 @@ def traced_method( return StreamWrapper(response, span) else: _set_openai_agentic_response_attributes(span, response) + + span.set_status(StatusCode.OK) + span.end() return response except Exception as err: span.record_exception(err) @@ -738,7 +741,7 @@ def _set_openai_agentic_response_attributes(span: Span, response) -> None: "input_tokens": response.usage.input_tokens, "output_tokens": response.usage.output_tokens, "total_tokens": response.usage.total_tokens, - "cached_tokens": response.usage.input_tokens_details["cached_tokens"], + "cached_tokens": response.usage.input_tokens_details.cached_tokens, }, ) diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py index d01aeae4..9d978f31 100644 --- a/src/langtrace_python_sdk/langtrace.py +++ b/src/langtrace_python_sdk/langtrace.py @@ -299,7 +299,8 @@ def init( init_instrumentations(config.disable_instrumentations, all_instrumentations) add_span_processor(provider, config, exporter) - sys.stdout = sys.__stdout__ + if config.disable_logging: + sys.stdout = sys.__stdout__ init_sentry(config, host) diff --git a/src/langtrace_python_sdk/utils/llm.py b/src/langtrace_python_sdk/utils/llm.py index d2dcfb7f..f2261405 100644 --- a/src/langtrace_python_sdk/utils/llm.py +++ b/src/langtrace_python_sdk/utils/llm.py @@ -347,19 +347,19 @@ def set_usage_attributes(span, usage): set_span_attribute( span, SpanAttributes.LLM_USAGE_PROMPT_TOKENS, - input_tokens, + int(input_tokens), ) set_span_attribute( span, SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, - output_tokens, + int(output_tokens), ) set_span_attribute( span, SpanAttributes.LLM_USAGE_TOTAL_TOKENS, - input_tokens + output_tokens, + int(input_tokens) + int(output_tokens), ) if "search_units" in usage: diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py index 01d6f588..b01d1044 100644 --- a/src/langtrace_python_sdk/version.py +++ b/src/langtrace_python_sdk/version.py @@ -1 +1 @@ -__version__ = "3.8.10" +__version__ = "3.8.11"