openlit 1.34.16__py3-none-any.whl → 1.34.17__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (8) hide show
  1. openlit/instrumentation/mistral/__init__.py +48 -40
  2. openlit/instrumentation/mistral/async_mistral.py +145 -556
  3. openlit/instrumentation/mistral/mistral.py +142 -553
  4. openlit/instrumentation/mistral/utils.py +298 -0
  5. {openlit-1.34.16.dist-info → openlit-1.34.17.dist-info}/METADATA +1 -1
  6. {openlit-1.34.16.dist-info → openlit-1.34.17.dist-info}/RECORD +8 -7
  7. {openlit-1.34.16.dist-info → openlit-1.34.17.dist-info}/LICENSE +0 -0
  8. {openlit-1.34.16.dist-info → openlit-1.34.17.dist-info}/WHEEL +0 -0
openlit/instrumentation/mistral/mistral.py CHANGED
@@ -2,610 +2,199 @@
2
2
  Module for monitoring Mistral API calls.
3
3
  """
4
4
 
5
- import logging
6
5
  import time
7
- from opentelemetry.trace import SpanKind, Status, StatusCode
8
- from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
6
+ from opentelemetry.trace import SpanKind
9
7
  from openlit.__helpers import (
10
- get_chat_model_cost,
11
- get_embed_model_cost,
12
8
  handle_exception,
13
- response_as_dict,
14
- calculate_ttft,
15
- calculate_tbt,
16
- create_metrics_attributes,
17
- set_server_address_and_port
9
+ set_server_address_and_port,
10
+ )
11
+ from openlit.instrumentation.mistral.utils import (
12
+ process_chunk,
13
+ process_chat_response,
14
+ process_streaming_chat_response,
15
+ process_embedding_response,
18
16
  )
19
17
  from openlit.semcov import SemanticConvention
20
18
 
21
- # Initialize logger for logging potential issues and operations
22
- logger = logging.getLogger(__name__)
23
-
24
- def chat(version, environment, application_name, tracer,
25
- pricing_info, capture_message_content, metrics, disable_metrics):
19
+ def complete(version, environment, application_name,
20
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics):
26
21
  """
27
- Generates a telemetry wrapper for chat to collect metrics.
28
-
29
- Args:
30
- version: Version of the monitoring package.
31
- environment: Deployment environment (e.g., production, staging).
32
- application_name: Name of the application using the Mistral API.
33
- tracer: OpenTelemetry tracer for creating spans.
34
- pricing_info: Information used for calculating the cost of Mistral usage.
35
- capture_message_content: Flag indicating whether to trace the actual content.
36
-
37
- Returns:
38
- A function that wraps the chat method to add telemetry.
22
+ Generates a telemetry wrapper for GenAI complete function call
39
23
  """
40
24
 
41
25
  def wrapper(wrapped, instance, args, kwargs):
42
26
  """
43
- Wraps the 'chat' API call to add telemetry.
44
-
45
- This collects metrics such as execution time, cost, and token usage, and handles errors
46
- gracefully, adding details to the trace for observability.
47
-
48
- Args:
49
- wrapped: The original 'chat' method to be wrapped.
50
- instance: The instance of the class where the original method is defined.
51
- args: Positional arguments for the 'chat' method.
52
- kwargs: Keyword arguments for the 'chat' method.
53
-
54
- Returns:
55
- The response from the original 'chat' method.
27
+ Wraps the GenAI complete function call.
56
28
  """
57
29
 
58
- server_address, server_port = set_server_address_and_port(instance, 'api.mistral.ai', 443)
59
- request_model = kwargs.get('model', 'mistral-small-latest')
30
+ server_address, server_port = set_server_address_and_port(instance, "api.mistral.ai", 443)
31
+ request_model = kwargs.get("model", "mistral-small-latest")
60
32
 
61
- span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}'
33
+ span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
62
34
 
63
35
  with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
64
36
  start_time = time.time()
65
37
  response = wrapped(*args, **kwargs)
66
- end_time = time.time()
38
+ response = process_chat_response(
39
+ response=response,
40
+ request_model=request_model,
41
+ pricing_info=pricing_info,
42
+ server_port=server_port,
43
+ server_address=server_address,
44
+ environment=environment,
45
+ application_name=application_name,
46
+ metrics=metrics,
47
+ start_time=start_time,
48
+ span=span,
49
+ capture_message_content=capture_message_content,
50
+ disable_metrics=disable_metrics,
51
+ version=version,
52
+ **kwargs
53
+ )
54
+
55
+ return response
67
56
 
68
- response_dict = response_as_dict(response)
57
+ return wrapper
69
58
 
70
- try:
71
- # Format 'messages' into a single string
72
- message_prompt = kwargs.get('messages', '')
73
- formatted_messages = []
74
- for message in message_prompt:
75
- role = message['role']
76
- content = message['content']
77
-
78
- if isinstance(content, list):
79
- content_str = ", ".join(
80
- f'{item["type"]}: {item["text"] if "text" in item else item["image_url"]}'
81
- if "type" in item else f'text: {item["text"]}'
82
- for item in content
83
- )
84
- formatted_messages.append(f'{role}: {content_str}')
85
- else:
86
- formatted_messages.append(f'{role}: {content}')
87
- prompt = '\n'.join(formatted_messages)
88
-
89
- input_tokens = response_dict.get('usage').get('prompt_tokens')
90
- output_tokens = response_dict.get('usage').get('completion_tokens')
91
-
92
- # Calculate cost of the operation
93
- cost = get_chat_model_cost(request_model,
94
- pricing_info, input_tokens,
95
- output_tokens)
96
-
97
- # Set base span attribues (OTel Semconv)
98
- span.set_attribute(TELEMETRY_SDK_NAME, 'openlit')
99
- span.set_attribute(SemanticConvention.GEN_AI_OPERATION,
100
- SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT)
101
- span.set_attribute(SemanticConvention.GEN_AI_SYSTEM,
102
- SemanticConvention.GEN_AI_SYSTEM_MISTRAL)
103
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_MODEL,
104
- request_model)
105
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_SEED,
106
- kwargs.get('seed', ''))
107
- span.set_attribute(SemanticConvention.SERVER_PORT,
108
- server_port)
109
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_FREQUENCY_PENALTY,
110
- kwargs.get('frequency_penalty', 0.0))
111
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_MAX_TOKENS,
112
- kwargs.get('max_tokens', -1))
113
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_PRESENCE_PENALTY,
114
- kwargs.get('presence_penalty', 0.0))
115
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_STOP_SEQUENCES,
116
- kwargs.get('stop', []))
117
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_TEMPERATURE,
118
- kwargs.get('temperature', 1.0))
119
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_TOP_P,
120
- kwargs.get('top_p', 1.0))
121
- span.set_attribute(SemanticConvention.GEN_AI_RESPONSE_ID,
122
- response_dict.get('id'))
123
- span.set_attribute(SemanticConvention.GEN_AI_RESPONSE_MODEL,
124
- response_dict.get('model'))
125
- span.set_attribute(SemanticConvention.GEN_AI_USAGE_INPUT_TOKENS,
126
- input_tokens)
127
- span.set_attribute(SemanticConvention.GEN_AI_USAGE_OUTPUT_TOKENS,
128
- output_tokens)
129
- span.set_attribute(SemanticConvention.SERVER_ADDRESS,
130
- server_address)
131
-
132
- # Set base span attribues (Extras)
133
- span.set_attribute(DEPLOYMENT_ENVIRONMENT,
134
- environment)
135
- span.set_attribute(SERVICE_NAME,
136
- application_name)
137
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_IS_STREAM,
138
- False)
139
- span.set_attribute(SemanticConvention.GEN_AI_USAGE_TOTAL_TOKENS,
140
- input_tokens + output_tokens)
141
- span.set_attribute(SemanticConvention.GEN_AI_USAGE_COST,
142
- cost)
143
- span.set_attribute(SemanticConvention.GEN_AI_SERVER_TTFT,
144
- end_time - start_time)
145
- span.set_attribute(SemanticConvention.GEN_AI_SDK_VERSION,
146
- version)
147
- if capture_message_content:
148
- span.add_event(
149
- name=SemanticConvention.GEN_AI_CONTENT_PROMPT_EVENT,
150
- attributes={
151
- SemanticConvention.GEN_AI_CONTENT_PROMPT: prompt,
152
- },
153
- )
154
-
155
- for i in range(kwargs.get('n',1)):
156
- span.set_attribute(SemanticConvention.GEN_AI_RESPONSE_FINISH_REASON,
157
- [response_dict.get('choices')[i].get('finish_reason')])
158
- if capture_message_content:
159
- span.add_event(
160
- name=SemanticConvention.GEN_AI_CONTENT_COMPLETION_EVENT,
161
- attributes={
162
- # pylint: disable=line-too-long
163
- SemanticConvention.GEN_AI_CONTENT_COMPLETION: str(response_dict.get('choices')[i].get('message').get('content')),
164
- },
165
- )
166
- if kwargs.get('tools'):
167
- span.set_attribute(SemanticConvention.GEN_AI_TOOL_CALLS,
168
- str(response_dict.get('choices')[i].get('message').get('tool_calls')))
169
-
170
- if isinstance(response_dict.get('choices')[i].get('message').get('content'), str):
171
- span.set_attribute(SemanticConvention.GEN_AI_OUTPUT_TYPE,
172
- 'text')
173
- elif response_dict.get('choices')[i].get('message').get('content') is not None:
174
- span.set_attribute(SemanticConvention.GEN_AI_OUTPUT_TYPE,
175
- 'json')
176
-
177
- span.set_status(Status(StatusCode.OK))
178
-
179
- if disable_metrics is False:
180
- attributes = create_metrics_attributes(
181
- service_name=application_name,
182
- deployment_environment=environment,
183
- operation=SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT,
184
- system=SemanticConvention.GEN_AI_SYSTEM_MISTRAL,
185
- request_model=request_model,
186
- server_address=server_address,
187
- server_port=server_port,
188
- response_model=response_dict.get('model'),
189
- )
190
-
191
- metrics['genai_client_usage_tokens'].record(
192
- input_tokens + output_tokens, attributes
193
- )
194
- metrics['genai_client_operation_duration'].record(
195
- end_time - start_time, attributes
196
- )
197
- metrics['genai_server_ttft'].record(
198
- end_time - start_time, attributes
199
- )
200
- metrics['genai_requests'].add(1, attributes)
201
- metrics['genai_completion_tokens'].add(output_tokens, attributes)
202
- metrics['genai_prompt_tokens'].add(input_tokens, attributes)
203
- metrics['genai_cost'].record(cost, attributes)
204
-
205
- # Return original response
206
- return response
59
+ def stream(version, environment, application_name,
60
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics):
61
+ """
62
+ Generates a telemetry wrapper for GenAI stream function call
63
+ """
207
64
 
208
- except Exception as e:
209
- handle_exception(span, e)
210
- logger.error('Error in trace creation: %s', e)
65
+ class TracedSyncStream:
66
+ """
67
+ Wrapper for streaming responses to collect telemetry.
68
+ """
211
69
 
212
- # Return original response
213
- return response
70
+ def __init__(
71
+ self,
72
+ wrapped,
73
+ span,
74
+ span_name,
75
+ kwargs,
76
+ server_address,
77
+ server_port,
78
+ **args,
79
+ ):
80
+ self.__wrapped__ = wrapped
81
+ self._span = span
82
+ self._span_name = span_name
83
+ self._llmresponse = ""
84
+ self._response_id = ""
85
+ self._response_model = ""
86
+ self._finish_reason = ""
87
+ self._tools = None
88
+ self._input_tokens = 0
89
+ self._output_tokens = 0
90
+
91
+ self._args = args
92
+ self._kwargs = kwargs
93
+ self._start_time = time.time()
94
+ self._end_time = None
95
+ self._timestamps = []
96
+ self._ttft = 0
97
+ self._tbt = 0
98
+ self._server_address = server_address
99
+ self._server_port = server_port
100
+
101
+ def __enter__(self):
102
+ self.__wrapped__.__enter__()
103
+ return self
104
+
105
+ def __exit__(self, exc_type, exc_value, traceback):
106
+ self.__wrapped__.__exit__(exc_type, exc_value, traceback)
107
+
108
+ def __iter__(self):
109
+ return self
110
+
111
+ def __getattr__(self, name):
112
+ """Delegate attribute access to the wrapped object."""
113
+ return getattr(self.__wrapped__, name)
114
+
115
+ def __next__(self):
116
+ try:
117
+ chunk = self.__wrapped__.__next__()
118
+ process_chunk(self, chunk)
119
+ return chunk
120
+ except StopIteration:
121
+ try:
122
+ with tracer.start_as_current_span(self._span_name, kind= SpanKind.CLIENT) as self._span:
123
+ process_streaming_chat_response(
124
+ self,
125
+ pricing_info=pricing_info,
126
+ environment=environment,
127
+ application_name=application_name,
128
+ metrics=metrics,
129
+ capture_message_content=capture_message_content,
130
+ disable_metrics=disable_metrics,
131
+ version=version
132
+ )
214
133
 
215
- return wrapper
134
+ except Exception as e:
135
+ handle_exception(self._span, e)
216
136
 
217
- def chat_stream(version, environment, application_name,
218
- tracer, pricing_info, capture_message_content, metrics, disable_metrics):
219
- """
220
- Generates a telemetry wrapper for chat_stream to collect metrics.
221
-
222
- Args:
223
- version: Version of the monitoring package.
224
- environment: Deployment environment (e.g., production, staging).
225
- application_name: Name of the application using the Mistral API.
226
- tracer: OpenTelemetry tracer for creating spans.
227
- pricing_info: Information used for calculating the cost of Mistral usage.
228
- capture_message_content: Flag indicating whether to trace the actual content.
229
-
230
- Returns:
231
- A function that wraps the chat method to add telemetry.
232
- """
137
+ raise
233
138
 
234
139
  def wrapper(wrapped, instance, args, kwargs):
235
140
  """
236
- Wraps the 'chat_stream' API call to add telemetry.
237
-
238
- This collects metrics such as execution time, cost, and token usage, and handles errors
239
- gracefully, adding details to the trace for observability.
240
-
241
- Args:
242
- wrapped: The original 'chat_stream' method to be wrapped.
243
- instance: The instance of the class where the original method is defined.
244
- args: Positional arguments for the 'chat_stream' method.
245
- kwargs: Keyword arguments for the 'chat_stream' method.
246
-
247
- Returns:
248
- The response from the original 'chat_stream' method.
141
+ Wraps the GenAI stream function call.
249
142
  """
250
143
 
251
- class TracedSyncStream:
252
- """
253
- Wrapper for streaming responses to collect metrics and trace data.
254
- Wraps the 'mistral.syncStream' response to collect message IDs and aggregated response.
255
-
256
- This class implements the '__aiter__' and '__anext__' methods that
257
- handle asynchronous streaming responses.
258
-
259
- This class also implements '__aenter__' and '__aexit__' methods that
260
- handle asynchronous context management protocol.
261
- """
262
- def __init__(
263
- self,
264
- wrapped,
265
- span,
266
- kwargs,
267
- server_address,
268
- server_port,
269
- **args,
270
- ):
271
- self.__wrapped__ = wrapped
272
- self._span = span
273
- # Placeholder for aggregating streaming response
274
- self._llmresponse = ''
275
- self._response_id = ''
276
- self._response_model = ''
277
- self._finish_reason = ''
278
- self._input_tokens = ''
279
- self._output_tokens = ''
280
-
281
- self._args = args
282
- self._kwargs = kwargs
283
- self._start_time = time.time()
284
- self._end_time = None
285
- self._timestamps = []
286
- self._ttft = 0
287
- self._tbt = 0
288
- self._server_address = server_address
289
- self._server_port = server_port
290
-
291
- def __enter__(self):
292
- self.__wrapped__.__enter__()
293
- return self
294
-
295
- def __exit__(self, exc_type, exc_value, traceback):
296
- self.__wrapped__.__exit__(exc_type, exc_value, traceback)
297
-
298
- def __iter__(self):
299
- return self
300
-
301
- def __getattr__(self, name):
302
- """Delegate attribute access to the wrapped object."""
303
- return getattr(self.__wrapped__, name)
304
-
305
- def __next__(self):
306
- try:
307
- chunk = self.__wrapped__.__next__()
308
- end_time = time.time()
309
- # Record the timestamp for the current chunk
310
- self._timestamps.append(end_time)
311
-
312
- if len(self._timestamps) == 1:
313
- # Calculate time to first chunk
314
- self._ttft = calculate_ttft(self._timestamps, self._start_time)
315
-
316
- chunked = response_as_dict(chunk)
317
-
318
- self._llmresponse += chunked.get('data').get('choices')[0].get('delta').get('content')
319
- if chunked.get('data').get('usage') is not None:
320
- self._response_id = chunked.get('data').get('id')
321
- self._response_model = chunked.get('data').get('model')
322
- self._input_tokens = chunked.get('data').get('usage').get('prompt_tokens')
323
- self._output_tokens = chunked.get('data').get('usage').get('completion_tokens')
324
- self._finish_reason = chunked.get('data').get('choices')[0].get('finish_reason')
325
-
326
- return chunk
327
- except StopIteration:
328
- # Handling exception ensure observability without disrupting operation
329
- try:
330
- self._end_time = time.time()
331
- if len(self._timestamps) > 1:
332
- self._tbt = calculate_tbt(self._timestamps)
333
-
334
- # Format 'messages' into a single string
335
- message_prompt = self._kwargs.get('messages', '')
336
- formatted_messages = []
337
- for message in message_prompt:
338
- role = message['role']
339
- content = message['content']
340
-
341
- if isinstance(content, list):
342
- content_str_list = []
343
- for item in content:
344
- if item['type'] == 'text':
345
- content_str_list.append(f'text: {item["text"]}')
346
- elif (item['type'] == 'image_url' and
347
- not item['image_url']['url'].startswith('data:')):
348
- content_str_list.append(f'image_url: {item["image_url"]["url"]}')
349
- content_str = ", ".join(content_str_list)
350
- formatted_messages.append(f'{role}: {content_str}')
351
- else:
352
- formatted_messages.append(f'{role}: {content}')
353
- prompt = '\n'.join(formatted_messages)
354
-
355
- request_model = self._kwargs.get('model', 'mistral-small-latest')
356
-
357
- # Calculate cost of the operation
358
- cost = get_chat_model_cost(request_model,
359
- pricing_info, self._input_tokens,
360
- self._output_tokens)
361
-
362
- # Set Span attributes (OTel Semconv)
363
- self._span.set_attribute(TELEMETRY_SDK_NAME, 'openlit')
364
- self._span.set_attribute(SemanticConvention.GEN_AI_OPERATION,
365
- SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT)
366
- self._span.set_attribute(SemanticConvention.GEN_AI_SYSTEM,
367
- SemanticConvention.GEN_AI_SYSTEM_MISTRAL)
368
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_MODEL,
369
- request_model)
370
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_SEED,
371
- self._kwargs.get('seed', ''))
372
- self._span.set_attribute(SemanticConvention.SERVER_PORT,
373
- self._server_port)
374
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_FREQUENCY_PENALTY,
375
- self._kwargs.get('frequency_penalty', 0.0))
376
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_MAX_TOKENS,
377
- self._kwargs.get('max_tokens', -1))
378
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_PRESENCE_PENALTY,
379
- self._kwargs.get('presence_penalty', 0.0))
380
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_STOP_SEQUENCES,
381
- self._kwargs.get('stop_sequences', []))
382
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_TEMPERATURE,
383
- self._kwargs.get('temperature', 0.3))
384
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_TOP_K,
385
- self._kwargs.get('k', 1.0))
386
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_TOP_P,
387
- self._kwargs.get('p', 1.0))
388
- self._span.set_attribute(SemanticConvention.GEN_AI_RESPONSE_FINISH_REASON,
389
- [self._finish_reason])
390
- self._span.set_attribute(SemanticConvention.GEN_AI_RESPONSE_ID,
391
- self._response_id)
392
- self._span.set_attribute(SemanticConvention.GEN_AI_RESPONSE_MODEL,
393
- self._response_model)
394
- self._span.set_attribute(SemanticConvention.GEN_AI_USAGE_INPUT_TOKENS,
395
- self._input_tokens)
396
- self._span.set_attribute(SemanticConvention.GEN_AI_USAGE_OUTPUT_TOKENS,
397
- self._output_tokens)
398
- self._span.set_attribute(SemanticConvention.SERVER_ADDRESS,
399
- self._server_address)
400
-
401
- if isinstance(self._llmresponse, str):
402
- self._span.set_attribute(SemanticConvention.GEN_AI_OUTPUT_TYPE,
403
- 'text')
404
- else:
405
- self._span.set_attribute(SemanticConvention.GEN_AI_OUTPUT_TYPE,
406
- 'json')
407
-
408
- # Set Span attributes (Extra)
409
- self._span.set_attribute(DEPLOYMENT_ENVIRONMENT,
410
- environment)
411
- self._span.set_attribute(SERVICE_NAME,
412
- application_name)
413
- self._span.set_attribute(SemanticConvention.GEN_AI_REQUEST_IS_STREAM,
414
- True)
415
- self._span.set_attribute(SemanticConvention.GEN_AI_USAGE_TOTAL_TOKENS,
416
- self._input_tokens + self._output_tokens)
417
- self._span.set_attribute(SemanticConvention.GEN_AI_USAGE_COST,
418
- cost)
419
- self._span.set_attribute(SemanticConvention.GEN_AI_SERVER_TBT,
420
- self._tbt)
421
- self._span.set_attribute(SemanticConvention.GEN_AI_SERVER_TTFT,
422
- self._ttft)
423
- self._span.set_attribute(SemanticConvention.GEN_AI_SDK_VERSION,
424
- version)
425
- if capture_message_content:
426
- self._span.add_event(
427
- name=SemanticConvention.GEN_AI_CONTENT_PROMPT_EVENT,
428
- attributes={
429
- SemanticConvention.GEN_AI_CONTENT_PROMPT: prompt,
430
- },
431
- )
432
- self._span.add_event(
433
- name=SemanticConvention.GEN_AI_CONTENT_COMPLETION_EVENT,
434
- attributes={
435
- SemanticConvention.GEN_AI_CONTENT_COMPLETION: self._llmresponse,
436
- },
437
- )
438
- self._span.set_status(Status(StatusCode.OK))
439
-
440
- if disable_metrics is False:
441
- attributes = create_metrics_attributes(
442
- service_name=application_name,
443
- deployment_environment=environment,
444
- operation=SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT,
445
- system=SemanticConvention.GEN_AI_SYSTEM_MISTRAL,
446
- request_model=request_model,
447
- server_address=self._server_address,
448
- server_port=self._server_port,
449
- response_model=self._response_model,
450
- )
451
-
452
- metrics['genai_client_usage_tokens'].record(
453
- self._input_tokens + self._output_tokens, attributes
454
- )
455
- metrics['genai_client_operation_duration'].record(
456
- self._end_time - self._start_time, attributes
457
- )
458
- metrics['genai_server_tbt'].record(
459
- self._tbt, attributes
460
- )
461
- metrics['genai_server_ttft'].record(
462
- self._ttft, attributes
463
- )
464
- metrics['genai_requests'].add(1, attributes)
465
- metrics['genai_completion_tokens'].add(self._output_tokens, attributes)
466
- metrics['genai_prompt_tokens'].add(self._input_tokens, attributes)
467
- metrics['genai_cost'].record(cost, attributes)
468
-
469
- except Exception as e:
470
- handle_exception(self._span, e)
471
- logger.error('Error in trace creation: %s', e)
472
- finally:
473
- self._span.end()
474
- raise
475
-
476
- server_address, server_port = set_server_address_and_port(instance, 'api.mistral.ai', 443)
477
- request_model = kwargs.get('model', 'mistral-small-latest')
478
-
479
- span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}'
144
+ server_address, server_port = set_server_address_and_port(instance, "api.mistral.ai", 443)
145
+ request_model = kwargs.get("model", "mistral-small-latest")
146
+
147
+ span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
480
148
 
149
+ # Stream endpoint is always streaming
481
150
  awaited_wrapped = wrapped(*args, **kwargs)
482
151
  span = tracer.start_span(span_name, kind=SpanKind.CLIENT)
483
- return TracedSyncStream(awaited_wrapped, span, kwargs, server_address, server_port)
152
+
153
+ return TracedSyncStream(awaited_wrapped, span, span_name, kwargs, server_address, server_port)
484
154
 
485
155
  return wrapper
486
156
 
487
- def embeddings(version, environment, application_name,
488
- tracer, pricing_info, capture_message_content, metrics, disable_metrics):
157
+ def embed(version, environment, application_name,
158
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics):
489
159
  """
490
- Generates a telemetry wrapper for embeddings to collect metrics.
491
-
492
- Args:
493
- version: Version of the monitoring package.
494
- environment: Deployment environment (e.g., production, staging).
495
- application_name: Name of the application using the Mistral API.
496
- tracer: OpenTelemetry tracer for creating spans.
497
- pricing_info: Information used for calculating the cost of Mistral usage.
498
- capture_message_content: Flag indicating whether to trace the actual content.
499
-
500
- Returns:
501
- A function that wraps the embeddings method to add telemetry.
160
+ Generates a telemetry wrapper for GenAI embedding function call
502
161
  """
503
162
 
504
163
  def wrapper(wrapped, instance, args, kwargs):
505
164
  """
506
- Wraps the 'embeddings' API call to add telemetry.
507
-
508
- This collects metrics such as execution time, cost, and token usage, and handles errors
509
- gracefully, adding details to the trace for observability.
510
-
511
- Args:
512
- wrapped: The original 'embeddings' method to be wrapped.
513
- instance: The instance of the class where the original method is defined.
514
- args: Positional arguments for the 'embeddings' method.
515
- kwargs: Keyword arguments for the 'embeddings' method.
516
-
517
- Returns:
518
- The response from the original 'embeddings' method.
165
+ Wraps the GenAI embedding function call.
519
166
  """
520
167
 
521
- server_address, server_port = set_server_address_and_port(instance, 'api.mistral.ai', 443)
522
- request_model = kwargs.get('model', 'mistral-embed')
168
+ server_address, server_port = set_server_address_and_port(instance, "api.mistral.ai", 443)
169
+ request_model = kwargs.get("model", "mistral-embed")
523
170
 
524
- span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_EMBEDDING} {request_model}'
171
+ span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_EMBEDDING} {request_model}"
525
172
 
526
- with tracer.start_as_current_span(span_name, kind= SpanKind.CLIENT) as span:
173
+ with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
527
174
  start_time = time.time()
528
175
  response = wrapped(*args, **kwargs)
529
- end_time = time.time()
530
176
 
531
- response_dict = response_as_dict(response)
532
177
  try:
533
- input_tokens = response_dict.get('usage').get('prompt_tokens')
534
-
535
- # Calculate cost of the operation
536
- cost = get_embed_model_cost(request_model,
537
- pricing_info, input_tokens)
538
-
539
- # Set Span attributes (OTel Semconv)
540
- span.set_attribute(TELEMETRY_SDK_NAME, 'openlit')
541
- span.set_attribute(SemanticConvention.GEN_AI_OPERATION,
542
- SemanticConvention.GEN_AI_OPERATION_TYPE_EMBEDDING)
543
- span.set_attribute(SemanticConvention.GEN_AI_SYSTEM,
544
- SemanticConvention.GEN_AI_SYSTEM_MISTRAL)
545
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_MODEL,
546
- request_model)
547
- span.set_attribute(SemanticConvention.GEN_AI_REQUEST_ENCODING_FORMATS,
548
- [kwargs.get('encoding_format', 'float')])
549
- span.set_attribute(SemanticConvention.GEN_AI_RESPONSE_MODEL,
550
- response_dict.get('model'))
551
- span.set_attribute(SemanticConvention.SERVER_ADDRESS,
552
- server_address)
553
- span.set_attribute(SemanticConvention.SERVER_PORT,
554
- server_port)
555
- span.set_attribute(SemanticConvention.GEN_AI_USAGE_INPUT_TOKENS,
556
- input_tokens)
557
-
558
- # Set Span attributes (Extras)
559
- span.set_attribute(DEPLOYMENT_ENVIRONMENT,
560
- environment)
561
- span.set_attribute(SERVICE_NAME,
562
- application_name)
563
- span.set_attribute(SemanticConvention.GEN_AI_USAGE_TOTAL_TOKENS,
564
- input_tokens)
565
- span.set_attribute(SemanticConvention.GEN_AI_USAGE_COST,
566
- cost)
567
- span.set_attribute(SemanticConvention.GEN_AI_SDK_VERSION,
568
- version)
569
-
570
- if capture_message_content:
571
- span.add_event(
572
- name=SemanticConvention.GEN_AI_CONTENT_PROMPT_EVENT,
573
- attributes={
574
- SemanticConvention.GEN_AI_CONTENT_PROMPT: str(kwargs.get('inputs', '')),
575
- },
576
- )
577
-
578
- span.set_status(Status(StatusCode.OK))
579
-
580
- if disable_metrics is False:
581
- attributes = create_metrics_attributes(
582
- service_name=application_name,
583
- deployment_environment=environment,
584
- operation=SemanticConvention.GEN_AI_OPERATION_TYPE_EMBEDDING,
585
- system=SemanticConvention.GEN_AI_SYSTEM_MISTRAL,
586
- request_model=request_model,
587
- server_address=server_address,
588
- server_port=server_port,
589
- response_model=response_dict.get('model'),
590
- )
591
- metrics['genai_client_usage_tokens'].record(
592
- input_tokens, attributes
593
- )
594
- metrics['genai_client_operation_duration'].record(
595
- end_time - start_time, attributes
596
- )
597
- metrics['genai_requests'].add(1, attributes)
598
- metrics['genai_prompt_tokens'].add(input_tokens, attributes)
599
- metrics['genai_cost'].record(cost, attributes)
600
-
601
- # Return original response
602
- return response
178
+ response = process_embedding_response(
179
+ response=response,
180
+ request_model=request_model,
181
+ pricing_info=pricing_info,
182
+ server_port=server_port,
183
+ server_address=server_address,
184
+ environment=environment,
185
+ application_name=application_name,
186
+ metrics=metrics,
187
+ start_time=start_time,
188
+ span=span,
189
+ capture_message_content=capture_message_content,
190
+ disable_metrics=disable_metrics,
191
+ version=version,
192
+ **kwargs
193
+ )
603
194
 
604
195
  except Exception as e:
605
196
  handle_exception(span, e)
606
- logger.error('Error in trace creation: %s', e)
607
197
 
608
- # Return original response
609
- return response
198
+ return response
610
199
 
611
200
  return wrapper