Skip to content

Commit

Permalink
Merge pull request #415 from Scale3-Labs/development
Browse files Browse the repository at this point in the history
Add `usage` to mistral instrumentation
  • Loading branch information
alizenhom authored Nov 21, 2024
2 parents 33f36e0 + 8672c8f commit 4740f7c
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 6 deletions.
7 changes: 2 additions & 5 deletions src/langtrace_python_sdk/instrumentation/mistral/patch.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@


def chat_complete(original_method, version, tracer, is_async=False, is_streaming=False):

def traced_method(wrapped, instance, args, kwargs):
service_provider = SERVICE_PROVIDERS["MISTRAL"]
llm_prompts = []
Expand All @@ -65,7 +65,6 @@ def traced_method(wrapped, instance, args, kwargs):

attributes = LLMSpanAttributes(**span_attributes)


span = tracer.start_span(
name=get_span_name(APIS[api]["METHOD"]),
kind=SpanKind.CLIENT,
Expand All @@ -87,7 +86,6 @@ def traced_method(wrapped, instance, args, kwargs):
span.set_status(StatusCode.OK)
span.end()
return result


except Exception as error:
span.record_exception(error)
Expand Down Expand Up @@ -186,7 +184,6 @@ def _set_response_attributes(span, kwargs, result):
for choice in result.choices
]
set_event_completion(span, responses)

# Get the usage
if hasattr(result, "usage") and result.usage is not None:
set_usage_attributes(span, result.usage)
set_usage_attributes(span, dict(result.usage))
2 changes: 1 addition & 1 deletion src/langtrace_python_sdk/version.py
Original file line number Diff line number Diff line change
@@ -1 +1 @@
__version__ = "3.3.8"
__version__ = "3.3.9"

0 comments on commit 4740f7c

Please sign in to comment.