openlit 1.31.1__py3-none-any.whl → 1.32.3__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.
@@ -0,0 +1,507 @@
1
+ # pylint: disable=duplicate-code, broad-exception-caught, too-many-statements, unused-argument, possibly-used-before-assignment, too-many-branches
2
+ """
3
+ Module for monitoring Prem AI API calls.
4
+ """
5
+
6
+ import logging
7
+ from opentelemetry.trace import SpanKind, Status, StatusCode
8
+ from opentelemetry.sdk.resources import TELEMETRY_SDK_NAME
9
+ from openlit.__helpers import (
10
+ handle_exception,
11
+ general_tokens,
12
+ get_chat_model_cost,
13
+ get_embed_model_cost,
14
+ response_as_dict
15
+ )
16
+ from openlit.semcov import SemanticConvetion
17
+
18
+ # Initialize logger for logging potential issues and operations
19
+ logger = logging.getLogger(__name__)
20
+
21
+ def chat(gen_ai_endpoint, version, environment, application_name,
22
+ tracer, pricing_info, trace_content, metrics, disable_metrics):
23
+ """
24
+ Generates a telemetry wrapper for chat completions to collect metrics.
25
+
26
+ Args:
27
+ gen_ai_endpoint: Endpoint identifier for logging and tracing.
28
+ version: Version of the monitoring package.
29
+ environment: Deployment environment (e.g., production, staging).
30
+ application_name: Name of the application using the PremAI API.
31
+ tracer: OpenTelemetry tracer for creating spans.
32
+ pricing_info: Information used for calculating the cost of PremAI usage.
33
+ trace_content: Flag indicating whether to trace the actual content.
34
+
35
+ Returns:
36
+ A function that wraps the chat completions method to add telemetry.
37
+ """
38
+
39
+ class TracedSyncStream:
40
+ """
41
+ Wrapper for streaming responses to collect metrics and trace data.
42
+ Wraps the response to collect message IDs and aggregated response.
43
+ """
44
+
45
+ def __init__(self, wrapped, span, kwargs, **args):
46
+ self.__wrapped__ = wrapped
47
+ self._span = span
48
+ self._llmresponse = ""
49
+ self._response_id = ""
50
+ self._args = args
51
+ self._kwargs = kwargs
52
+
53
+ def __enter__(self):
54
+ # Using context management protocols (if needed)
55
+ return self
56
+
57
+ def __exit__(self, exc_type, exc_value, traceback):
58
+ # Add any resource cleanup or finalization if required.
59
+ pass
60
+
61
+ def __getattr__(self, name):
62
+ """Delegate attribute access to the wrapped object."""
63
+ return getattr(self.__wrapped__, name)
64
+
65
+ def __iter__(self):
66
+ try:
67
+ for chunk in self.__wrapped__:
68
+ # Assuming `chunk` has similar structure as 'ChatCompletionResponseStream'
69
+ if chunk.choices:
70
+ first_choice = chunk.choices[0]
71
+
72
+ if first_choice.delta.get('content'):
73
+ self._llmresponse += first_choice.delta.get('content')
74
+
75
+ self._response_id = chunk.id
76
+ if not chunk:
77
+ # pylint: disable= stop-iteration-return
78
+ raise StopIteration
79
+ yield chunk
80
+
81
+ finally:
82
+ # Handling exception ensure observability without disrupting operation
83
+ try:
84
+ # Format 'messages' into a single string
85
+ message_prompt = self._kwargs.get("messages", "")
86
+ formatted_messages = []
87
+ for message in message_prompt:
88
+ role = message["role"]
89
+ content = message["content"]
90
+
91
+ if isinstance(content, list):
92
+ content_str = ", ".join(
93
+ # pylint: disable=line-too-long
94
+ f'{item["type"]}: {item["text"] if "text" in item else item["image_url"]}'
95
+ if "type" in item else f'text: {item["text"]}'
96
+ for item in content
97
+ )
98
+ formatted_messages.append(f"{role}: {content_str}")
99
+ else:
100
+ formatted_messages.append(f"{role}: {content}")
101
+ prompt = "\n".join(formatted_messages)
102
+
103
+ # Calculate tokens using input prompt and aggregated response
104
+ prompt_tokens = general_tokens(prompt,)
105
+ completion_tokens = general_tokens(self._llmresponse)
106
+
107
+ # Calculate cost of the operation
108
+ cost = get_chat_model_cost(self._kwargs.get("model", "gpt-4o-mini"),
109
+ pricing_info, prompt_tokens,
110
+ completion_tokens)
111
+ print(self._kwargs)
112
+ # Set Span attributes
113
+ self._span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
114
+ self._span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM,
115
+ SemanticConvetion.GEN_AI_SYSTEM_PREMAI)
116
+ self._span.set_attribute(SemanticConvetion.GEN_AI_TYPE,
117
+ SemanticConvetion.GEN_AI_TYPE_CHAT)
118
+ self._span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
119
+ gen_ai_endpoint)
120
+ self._span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_ID,
121
+ self._response_id)
122
+ self._span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
123
+ environment)
124
+ self._span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
125
+ application_name)
126
+ self._span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL,
127
+ self._kwargs.get("model", "gpt-4o-mini"))
128
+ self._span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_USER,
129
+ self._kwargs.get("user", ""))
130
+ self._span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TOP_P,
131
+ self._kwargs.get("top_p", 1.0))
132
+ self._span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MAX_TOKENS,
133
+ self._kwargs.get("max_tokens", -1))
134
+ self._span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE,
135
+ self._kwargs.get("temperature", 1.0))
136
+ self._span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_PRESENCE_PENALTY,
137
+ self._kwargs.get("presence_penalty", 0.0))
138
+ self._span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_FREQUENCY_PENALTY,
139
+ self._kwargs.get("frequency_penalty", 0.0))
140
+ self._span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_SEED,
141
+ self._kwargs.get("seed", ""))
142
+ self._span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_IS_STREAM,
143
+ True)
144
+ self._span.set_attribute(SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
145
+ prompt_tokens)
146
+ self._span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
147
+ completion_tokens)
148
+ self._span.set_attribute(SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
149
+ prompt_tokens + completion_tokens)
150
+ self._span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST,
151
+ cost)
152
+ if trace_content:
153
+ self._span.add_event(
154
+ name=SemanticConvetion.GEN_AI_CONTENT_PROMPT_EVENT,
155
+ attributes={
156
+ SemanticConvetion.GEN_AI_CONTENT_PROMPT: prompt,
157
+ },
158
+ )
159
+ self._span.add_event(
160
+ name=SemanticConvetion.GEN_AI_CONTENT_COMPLETION_EVENT,
161
+ attributes={
162
+ SemanticConvetion.GEN_AI_CONTENT_COMPLETION: self._llmresponse,
163
+ },
164
+ )
165
+
166
+ self._span.set_status(Status(StatusCode.OK))
167
+
168
+ if disable_metrics is False:
169
+ attributes = {
170
+ TELEMETRY_SDK_NAME:
171
+ "openlit",
172
+ SemanticConvetion.GEN_AI_APPLICATION_NAME:
173
+ application_name,
174
+ SemanticConvetion.GEN_AI_SYSTEM:
175
+ SemanticConvetion.GEN_AI_SYSTEM_PREMAI,
176
+ SemanticConvetion.GEN_AI_ENVIRONMENT:
177
+ environment,
178
+ SemanticConvetion.GEN_AI_TYPE:
179
+ SemanticConvetion.GEN_AI_TYPE_CHAT,
180
+ SemanticConvetion.GEN_AI_REQUEST_MODEL:
181
+ self._kwargs.get("model", "gpt-3.5-turbo")
182
+ }
183
+
184
+ metrics["genai_requests"].add(1, attributes)
185
+ metrics["genai_total_tokens"].add(
186
+ prompt_tokens + completion_tokens, attributes
187
+ )
188
+ metrics["genai_completion_tokens"].add(completion_tokens, attributes)
189
+ metrics["genai_prompt_tokens"].add(prompt_tokens, attributes)
190
+ metrics["genai_cost"].record(cost, attributes)
191
+
192
+ except Exception as e:
193
+ handle_exception(self._span, e)
194
+ logger.error("Error in trace creation: %s", e)
195
+ finally:
196
+ self._span.end()
197
+
198
+ def wrapper(wrapped, instance, args, kwargs):
199
+ """
200
+ Wraps the 'chat.completions' API call to add telemetry.
201
+
202
+ This collects metrics such as execution time, cost, and token usage, and handles errors
203
+ gracefully, adding details to the trace for observability.
204
+
205
+ Args:
206
+ wrapped: The original 'chat.completions' method to be wrapped.
207
+ instance: The instance of the class where the original method is defined.
208
+ args: Positional arguments for the 'chat.completions' method.
209
+ kwargs: Keyword arguments for the 'chat.completions' method.
210
+
211
+ Returns:
212
+ The response from the original 'chat.completions' method.
213
+ """
214
+
215
+ # Check if streaming is enabled for the API call
216
+ streaming = kwargs.get("stream", False)
217
+
218
+ # pylint: disable=no-else-return
219
+ if streaming:
220
+ # Special handling for streaming response to accommodate the nature of data flow
221
+ awaited_wrapped = wrapped(*args, **kwargs)
222
+ span = tracer.start_span(gen_ai_endpoint, kind=SpanKind.CLIENT)
223
+
224
+ return TracedSyncStream(awaited_wrapped, span, kwargs)
225
+
226
+ # Handling for non-streaming responses
227
+ else:
228
+ # pylint: disable=line-too-long
229
+ with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
230
+ response = wrapped(*args, **kwargs)
231
+
232
+ response_dict = response_as_dict(response)
233
+
234
+ try:
235
+ # Format 'messages' into a single string
236
+ message_prompt = kwargs.get("messages", "")
237
+ formatted_messages = []
238
+ for message in message_prompt:
239
+ role = message["role"]
240
+ content = message["content"]
241
+
242
+ if isinstance(content, list):
243
+ content_str = ", ".join(
244
+ # pylint: disable=line-too-long
245
+ f'{item["type"]}: {item["text"] if "text" in item else item["image_url"]}'
246
+ if "type" in item else f'text: {item["text"]}'
247
+ for item in content
248
+ )
249
+ formatted_messages.append(f"{role}: {content_str}")
250
+ else:
251
+ formatted_messages.append(f"{role}: {content}")
252
+ prompt = "\n".join(formatted_messages)
253
+
254
+ # Set base span attribues
255
+ span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
256
+ span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM,
257
+ SemanticConvetion.GEN_AI_SYSTEM_PREMAI)
258
+ span.set_attribute(SemanticConvetion.GEN_AI_TYPE,
259
+ SemanticConvetion.GEN_AI_TYPE_CHAT)
260
+ span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
261
+ gen_ai_endpoint)
262
+ span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_ID,
263
+ response_dict.additional_properties["id"])
264
+ span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
265
+ environment)
266
+ span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
267
+ application_name)
268
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL,
269
+ kwargs.get("model", "gpt-3.5-turbo"))
270
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TOP_P,
271
+ kwargs.get("top_p", 1.0))
272
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MAX_TOKENS,
273
+ kwargs.get("max_tokens", -1))
274
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_USER,
275
+ kwargs.get("user", ""))
276
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE,
277
+ kwargs.get("temperature", 1.0))
278
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_PRESENCE_PENALTY,
279
+ kwargs.get("presence_penalty", 0.0))
280
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_FREQUENCY_PENALTY,
281
+ kwargs.get("frequency_penalty", 0.0))
282
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_SEED,
283
+ kwargs.get("seed", ""))
284
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_IS_STREAM,
285
+ False)
286
+ if trace_content:
287
+ span.add_event(
288
+ name=SemanticConvetion.GEN_AI_CONTENT_PROMPT_EVENT,
289
+ attributes={
290
+ SemanticConvetion.GEN_AI_CONTENT_PROMPT: prompt,
291
+ },
292
+ )
293
+
294
+ # Set span attributes when tools is not passed to the function call
295
+ if "tools" not in kwargs:
296
+ # Calculate cost of the operation
297
+ cost = get_chat_model_cost(kwargs.get("model", "gpt-4o-mini"),
298
+ pricing_info, response_dict.usage.prompt_tokens,
299
+ response_dict.usage.completion_tokens)
300
+
301
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
302
+ response_dict.usage.prompt_tokens)
303
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
304
+ response_dict.usage.completion_tokens)
305
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
306
+ response_dict.usage.total_tokens)
307
+ span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_FINISH_REASON,
308
+ [response_dict.choices[0].finish_reason])
309
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST,
310
+ cost)
311
+
312
+ # Set span attributes for when n = 1 (default)
313
+ if "n" not in kwargs or kwargs["n"] == 1:
314
+ if trace_content:
315
+ span.add_event(
316
+ name=SemanticConvetion.GEN_AI_CONTENT_COMPLETION_EVENT,
317
+ attributes={
318
+ SemanticConvetion.GEN_AI_CONTENT_COMPLETION: response_dict.choices[0].message.content,
319
+ },
320
+ )
321
+
322
+ # Set span attributes for when n > 0
323
+ else:
324
+ i = 0
325
+ while i < kwargs["n"] and trace_content is True:
326
+ attribute_name = f"gen_ai.content.completion.{i}"
327
+ span.add_event(
328
+ name=attribute_name,
329
+ attributes={
330
+ SemanticConvetion.GEN_AI_CONTENT_COMPLETION: response_dict.choices[i].message.content,
331
+ },
332
+ )
333
+ i += 1
334
+
335
+ # Return original response
336
+ return response
337
+
338
+ # Set span attributes when tools is passed to the function call
339
+ elif "tools" in kwargs:
340
+ # Calculate cost of the operation
341
+ cost = get_chat_model_cost(kwargs.get("model", "gpt-3.5-turbo"),
342
+ pricing_info, response_dict.usage.prompt_tokens,
343
+ response_dict.usage.completion_tokens)
344
+ span.add_event(
345
+ name=SemanticConvetion.GEN_AI_CONTENT_COMPLETION_EVENT,
346
+ attributes={
347
+ SemanticConvetion.GEN_AI_CONTENT_COMPLETION: "Function called with tools",
348
+ },
349
+ )
350
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
351
+ response_dict.usage.prompt_tokens)
352
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
353
+ response_dict.usage.completion_tokens)
354
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
355
+ response_dict.usage.total_tokens)
356
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST,
357
+ cost)
358
+
359
+ span.set_status(Status(StatusCode.OK))
360
+
361
+ if disable_metrics is False:
362
+ attributes = {
363
+ TELEMETRY_SDK_NAME:
364
+ "openlit",
365
+ SemanticConvetion.GEN_AI_APPLICATION_NAME:
366
+ application_name,
367
+ SemanticConvetion.GEN_AI_SYSTEM:
368
+ SemanticConvetion.GEN_AI_SYSTEM_PREMAI,
369
+ SemanticConvetion.GEN_AI_ENVIRONMENT:
370
+ environment,
371
+ SemanticConvetion.GEN_AI_TYPE:
372
+ SemanticConvetion.GEN_AI_TYPE_CHAT,
373
+ SemanticConvetion.GEN_AI_REQUEST_MODEL:
374
+ kwargs.get("model", "gpt-3.5-turbo")
375
+ }
376
+
377
+ metrics["genai_requests"].add(1, attributes)
378
+ metrics["genai_total_tokens"].add(response_dict.usage.total_tokens, attributes)
379
+ metrics["genai_completion_tokens"].add(response_dict.usage.completion_tokens, attributes)
380
+ metrics["genai_prompt_tokens"].add(response_dict.usage.prompt_tokens, attributes)
381
+ metrics["genai_cost"].record(cost, attributes)
382
+
383
+ # Return original response
384
+ return response
385
+
386
+ except Exception as e:
387
+ handle_exception(span, e)
388
+ logger.error("Error in trace creation: %s", e)
389
+
390
+ # Return original response
391
+ return response
392
+
393
+ return wrapper
394
+
395
+ def embedding(gen_ai_endpoint, version, environment, application_name,
396
+ tracer, pricing_info, trace_content, metrics, disable_metrics):
397
+ """
398
+ Generates a telemetry wrapper for embeddings to collect metrics.
399
+
400
+ Args:
401
+ gen_ai_endpoint: Endpoint identifier for logging and tracing.
402
+ version: Version of the monitoring package.
403
+ environment: Deployment environment (e.g., production, staging).
404
+ application_name: Name of the application using the Prem AI API.
405
+ tracer: OpenTelemetry tracer for creating spans.
406
+ pricing_info: Information used for calculating the cost of Prem AI usage.
407
+ trace_content: Flag indicating whether to trace the actual content.
408
+
409
+ Returns:
410
+ A function that wraps the embeddings method to add telemetry.
411
+ """
412
+
413
+ def wrapper(wrapped, instance, args, kwargs):
414
+ """
415
+ Wraps the 'embeddings' API call to add telemetry.
416
+
417
+ This collects metrics such as execution time, cost, and token usage, and handles errors
418
+ gracefully, adding details to the trace for observability.
419
+
420
+ Args:
421
+ wrapped: The original 'embeddings' method to be wrapped.
422
+ instance: The instance of the class where the original method is defined.
423
+ args: Positional arguments for the 'embeddings' method.
424
+ kwargs: Keyword arguments for the 'embeddings' method.
425
+
426
+ Returns:
427
+ The response from the original 'embeddings' method.
428
+ """
429
+
430
+ with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
431
+ response = wrapped(*args, **kwargs)
432
+ response_dict = response_as_dict(response)
433
+ try:
434
+ # Calculate cost of the operation
435
+ cost = get_embed_model_cost(kwargs.get("model", "text-embedding-ada-002"),
436
+ pricing_info, response_dict.usage.prompt_tokens)
437
+
438
+ # Set Span attributes
439
+ span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
440
+ span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM,
441
+ SemanticConvetion.GEN_AI_SYSTEM_PREMAI)
442
+ span.set_attribute(SemanticConvetion.GEN_AI_TYPE,
443
+ SemanticConvetion.GEN_AI_TYPE_EMBEDDING)
444
+ span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
445
+ gen_ai_endpoint)
446
+ span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
447
+ environment)
448
+ span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
449
+ application_name)
450
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL,
451
+ kwargs.get("model", "text-embedding-3-large"))
452
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_EMBEDDING_FORMAT,
453
+ kwargs.get("encoding_format", "float"))
454
+ # span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_EMBEDDING_DIMENSION,
455
+ # kwargs.get("dimensions", "null"))
456
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_USER,
457
+ kwargs.get("user", ""))
458
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
459
+ response_dict.usage.prompt_tokens)
460
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
461
+ response_dict.usage.total_tokens)
462
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST,
463
+ cost)
464
+ if trace_content:
465
+ span.add_event(
466
+ name=SemanticConvetion.GEN_AI_CONTENT_PROMPT_EVENT,
467
+ attributes={
468
+ SemanticConvetion.GEN_AI_CONTENT_PROMPT: kwargs.get("input", ""),
469
+ },
470
+ )
471
+
472
+ span.set_status(Status(StatusCode.OK))
473
+
474
+ if disable_metrics is False:
475
+ attributes = {
476
+ TELEMETRY_SDK_NAME:
477
+ "openlit",
478
+ SemanticConvetion.GEN_AI_APPLICATION_NAME:
479
+ application_name,
480
+ SemanticConvetion.GEN_AI_SYSTEM:
481
+ SemanticConvetion.GEN_AI_SYSTEM_PREMAI,
482
+ SemanticConvetion.GEN_AI_ENVIRONMENT:
483
+ environment,
484
+ SemanticConvetion.GEN_AI_TYPE:
485
+ SemanticConvetion.GEN_AI_TYPE_EMBEDDING,
486
+ SemanticConvetion.GEN_AI_REQUEST_MODEL:
487
+ kwargs.get("model", "text-embedding-ada-002")
488
+ }
489
+
490
+ metrics["genai_requests"].add(1, attributes)
491
+ metrics["genai_total_tokens"].add(
492
+ response_dict.usage.total_tokens, attributes)
493
+ metrics["genai_prompt_tokens"].add(
494
+ response_dict.usageprompt_tokens, attributes)
495
+ metrics["genai_cost"].record(cost, attributes)
496
+
497
+ # Return original response
498
+ return response
499
+
500
+ except Exception as e:
501
+ handle_exception(span, e)
502
+ logger.error("Error in trace creation: %s", e)
503
+
504
+ # Return original response
505
+ return response
506
+
507
+ return wrapper
@@ -0,0 +1,54 @@
1
+ # pylint: disable=useless-return, bad-staticmethod-argument, disable=duplicate-code
2
+ """Initializer of Auto Instrumentation of Reka Functions"""
3
+
4
+ from typing import Collection
5
+ import importlib.metadata
6
+ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
7
+ from wrapt import wrap_function_wrapper
8
+
9
+ from openlit.instrumentation.reka.reka import (
10
+ chat
11
+ )
12
+ from openlit.instrumentation.reka.async_reka import (
13
+ async_chat
14
+ )
15
+
16
+ _instruments = ("reka-api >= 3.2.0",)
17
+
18
+ class RekaInstrumentor(BaseInstrumentor):
19
+ """
20
+ An instrumentor for Reka's client library.
21
+ """
22
+
23
+ def instrumentation_dependencies(self) -> Collection[str]:
24
+ return _instruments
25
+
26
+ def _instrument(self, **kwargs):
27
+ application_name = kwargs.get("application_name", "default_application")
28
+ environment = kwargs.get("environment", "default_environment")
29
+ tracer = kwargs.get("tracer")
30
+ metrics = kwargs.get("metrics_dict")
31
+ pricing_info = kwargs.get("pricing_info", {})
32
+ trace_content = kwargs.get("trace_content", False)
33
+ disable_metrics = kwargs.get("disable_metrics")
34
+ version = importlib.metadata.version("reka-api")
35
+
36
+ # sync chat
37
+ wrap_function_wrapper(
38
+ "reka.chat.client",
39
+ "ChatClient.create",
40
+ chat("reka.chat", version, environment, application_name,
41
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
42
+ )
43
+
44
+ # async chat
45
+ wrap_function_wrapper(
46
+ "reka.chat.client",
47
+ "AsyncChatClient.create",
48
+ async_chat("reka.chat", version, environment, application_name,
49
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
50
+ )
51
+
52
+ def _uninstrument(self, **kwargs):
53
+ # Proper uninstrumentation logic to revert patched methods
54
+ pass