openlit 1.34.2__py3-none-any.whl → 1.34.4__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 (11) hide show
  1. openlit/instrumentation/ai21/__init__.py +22 -23
  2. openlit/instrumentation/ai21/ai21.py +58 -51
  3. openlit/instrumentation/ai21/async_ai21.py +58 -51
  4. openlit/instrumentation/ai21/utils.py +134 -243
  5. openlit/instrumentation/openai/__init__.py +17 -2
  6. openlit/instrumentation/openai/async_openai.py +161 -0
  7. openlit/instrumentation/openai/openai.py +161 -0
  8. {openlit-1.34.2.dist-info → openlit-1.34.4.dist-info}/METADATA +1 -1
  9. {openlit-1.34.2.dist-info → openlit-1.34.4.dist-info}/RECORD +11 -11
  10. {openlit-1.34.2.dist-info → openlit-1.34.4.dist-info}/LICENSE +0 -0
  11. {openlit-1.34.2.dist-info → openlit-1.34.4.dist-info}/WHEEL +0 -0
openlit/instrumentation/ai21/__init__.py CHANGED
@@ -13,53 +13,52 @@ from openlit.instrumentation.ai21.async_ai21 import (
13
13
  async_chat, async_chat_rag
14
14
  )
15
15
 
16
- _instruments = ('ai21 >= 3.0.0',)
16
+ _instruments = ("ai21 >= 3.0.0",)
17
17
 
18
18
  class AI21Instrumentor(BaseInstrumentor):
19
19
  """
20
- An instrumentor for AI21's client library.
20
+ An instrumentor for AI21 client library.
21
21
  """
22
22
 
23
23
  def instrumentation_dependencies(self) -> Collection[str]:
24
24
  return _instruments
25
25
 
26
26
  def _instrument(self, **kwargs):
27
- application_name = kwargs.get('application_name', 'default')
28
- environment = kwargs.get('environment', 'default')
29
- tracer = kwargs.get('tracer')
30
- event_provider = kwargs.get('event_provider')
31
- metrics = kwargs.get('metrics_dict')
32
- pricing_info = kwargs.get('pricing_info', {})
33
- capture_message_content = kwargs.get('capture_message_content', False)
34
- disable_metrics = kwargs.get('disable_metrics')
35
- version = importlib.metadata.version('ai21')
27
+ application_name = kwargs.get("application_name", "default")
28
+ environment = kwargs.get("environment", "default")
29
+ tracer = kwargs.get("tracer")
30
+ metrics = kwargs.get("metrics_dict")
31
+ pricing_info = kwargs.get("pricing_info", {})
32
+ capture_message_content = kwargs.get("capture_message_content", False)
33
+ disable_metrics = kwargs.get("disable_metrics")
34
+ version = importlib.metadata.version("ai21")
36
35
 
37
36
  #sync
38
37
  wrap_function_wrapper(
39
- 'ai21.clients.studio.resources.chat.chat_completions',
40
- 'ChatCompletions.create',
38
+ "ai21.clients.studio.resources.chat.chat_completions",
39
+ "ChatCompletions.create",
41
40
  chat(version, environment, application_name,
42
- tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
41
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics),
43
42
  )
44
43
  wrap_function_wrapper(
45
- 'ai21.clients.studio.resources.studio_conversational_rag',
46
- 'StudioConversationalRag.create',
44
+ "ai21.clients.studio.resources.studio_conversational_rag",
45
+ "StudioConversationalRag.create",
47
46
  chat_rag(version, environment, application_name,
48
- tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
47
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics),
49
48
  )
50
49
 
51
50
  #Async
52
51
  wrap_function_wrapper(
53
- 'ai21.clients.studio.resources.chat.async_chat_completions',
54
- 'AsyncChatCompletions.create',
52
+ "ai21.clients.studio.resources.chat.async_chat_completions",
53
+ "AsyncChatCompletions.create",
55
54
  async_chat(version, environment, application_name,
56
- tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
55
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics),
57
56
  )
58
57
  wrap_function_wrapper(
59
- 'ai21.clients.studio.resources.studio_conversational_rag',
60
- 'AsyncStudioConversationalRag.create',
58
+ "ai21.clients.studio.resources.studio_conversational_rag",
59
+ "AsyncStudioConversationalRag.create",
61
60
  async_chat_rag(version, environment, application_name,
62
- tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics),
61
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics),
63
62
  )
64
63
 
65
64
  def _uninstrument(self, **kwargs):
openlit/instrumentation/ai21/ai21.py CHANGED
@@ -22,7 +22,7 @@ from openlit.semcov import SemanticConvention
22
22
  logger = logging.getLogger(__name__)
23
23
 
24
24
  def chat(version, environment, application_name,
25
- tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
25
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics):
26
26
  """
27
27
  Generates a telemetry wrapper for GenAI function call
28
28
  """
@@ -46,9 +46,9 @@ def chat(version, environment, application_name,
46
46
  self._span = span
47
47
  self._span_name = span_name
48
48
  # Placeholder for aggregating streaming response
49
- self._llmresponse = ''
50
- self._response_id = ''
51
- self._finish_reason = ''
49
+ self._llmresponse = ""
50
+ self._response_id = ""
51
+ self._finish_reason = ""
52
52
  self._input_tokens = 0
53
53
  self._output_tokens = 0
54
54
  self._choices = []
@@ -92,14 +92,13 @@ def chat(version, environment, application_name,
92
92
  environment=environment,
93
93
  application_name=application_name,
94
94
  metrics=metrics,
95
- event_provider=event_provider,
96
95
  capture_message_content=capture_message_content,
97
96
  disable_metrics=disable_metrics,
98
97
  version=version
99
98
  )
100
99
  except Exception as e:
101
100
  handle_exception(self._span, e)
102
- logger.error('Error in trace creation: %s', e)
101
+ logger.error("Error in trace creation: %s", e)
103
102
  raise
104
103
 
105
104
  def wrapper(wrapped, instance, args, kwargs):
@@ -108,12 +107,12 @@ def chat(version, environment, application_name,
108
107
  """
109
108
 
110
109
  # Check if streaming is enabled for the API call
111
- streaming = kwargs.get('stream', False)
110
+ streaming = kwargs.get("stream", False)
112
111
 
113
- server_address, server_port = set_server_address_and_port(instance, 'api.ai21.com', 443)
114
- request_model = kwargs.get('model', 'jamba-1.5-mini')
112
+ server_address, server_port = set_server_address_and_port(instance, "api.ai21.com", 443)
113
+ request_model = kwargs.get("model", "jamba-1.5-mini")
115
114
 
116
- span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}'
115
+ span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
117
116
 
118
117
  # pylint: disable=no-else-return
119
118
  if streaming:
@@ -127,30 +126,34 @@ def chat(version, environment, application_name,
127
126
  with tracer.start_as_current_span(span_name, kind= SpanKind.CLIENT) as span:
128
127
  start_time = time.time()
129
128
  response = wrapped(*args, **kwargs)
130
- response = process_chat_response(
131
- response=response,
132
- request_model=request_model,
133
- pricing_info=pricing_info,
134
- server_port=server_port,
135
- server_address=server_address,
136
- environment=environment,
137
- application_name=application_name,
138
- metrics=metrics,
139
- event_provider=event_provider,
140
- start_time=start_time,
141
- span=span,
142
- capture_message_content=capture_message_content,
143
- disable_metrics=disable_metrics,
144
- version=version,
145
- **kwargs
146
- )
147
129
 
148
- return response
130
+ try:
131
+ response = process_chat_response(
132
+ response=response,
133
+ request_model=request_model,
134
+ pricing_info=pricing_info,
135
+ server_port=server_port,
136
+ server_address=server_address,
137
+ environment=environment,
138
+ application_name=application_name,
139
+ metrics=metrics,
140
+ start_time=start_time,
141
+ span=span,
142
+ capture_message_content=capture_message_content,
143
+ disable_metrics=disable_metrics,
144
+ version=version,
145
+ **kwargs
146
+ )
147
+
148
+ except Exception as e:
149
+ handle_exception(span, e)
150
+
151
+ return response
149
152
 
150
153
  return wrapper
151
154
 
152
155
  def chat_rag(version, environment, application_name,
153
- tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
156
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics):
154
157
  """
155
158
  Generates a telemetry wrapper for GenAI function call
156
159
  """
@@ -160,32 +163,36 @@ def chat_rag(version, environment, application_name,
160
163
  Wraps the GenAI function call.
161
164
  """
162
165
 
163
- server_address, server_port = set_server_address_and_port(instance, 'api.ai21.com', 443)
164
- request_model = kwargs.get('model', 'jamba-1.5-mini')
166
+ server_address, server_port = set_server_address_and_port(instance, "api.ai21.com", 443)
167
+ request_model = kwargs.get("model", "jamba-1.5-mini")
165
168
 
166
- span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}'
169
+ span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
167
170
 
168
171
  with tracer.start_as_current_span(span_name, kind= SpanKind.CLIENT) as span:
169
172
  start_time = time.time()
170
173
  response = wrapped(*args, **kwargs)
171
- response = process_chat_rag_response(
172
- response=response,
173
- request_model=request_model,
174
- pricing_info=pricing_info,
175
- server_port=server_port,
176
- server_address=server_address,
177
- environment=environment,
178
- application_name=application_name,
179
- metrics=metrics,
180
- event_provider=event_provider,
181
- start_time=start_time,
182
- span=span,
183
- capture_message_content=capture_message_content,
184
- disable_metrics=disable_metrics,
185
- version=version,
186
- **kwargs
187
- )
188
-
189
- return response
174
+
175
+ try:
176
+ response = process_chat_rag_response(
177
+ response=response,
178
+ request_model=request_model,
179
+ pricing_info=pricing_info,
180
+ server_port=server_port,
181
+ server_address=server_address,
182
+ environment=environment,
183
+ application_name=application_name,
184
+ metrics=metrics,
185
+ start_time=start_time,
186
+ span=span,
187
+ capture_message_content=capture_message_content,
188
+ disable_metrics=disable_metrics,
189
+ version=version,
190
+ **kwargs
191
+ )
192
+
193
+ except Exception as e:
194
+ handle_exception(span, e)
195
+
196
+ return response
190
197
 
191
198
  return wrapper
openlit/instrumentation/ai21/async_ai21.py CHANGED
@@ -22,7 +22,7 @@ from openlit.semcov import SemanticConvention
22
22
  logger = logging.getLogger(__name__)
23
23
 
24
24
  def async_chat(version, environment, application_name,
25
- tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
25
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics):
26
26
  """
27
27
  Generates a telemetry wrapper for GenAI function call
28
28
  """
@@ -46,9 +46,9 @@ def async_chat(version, environment, application_name,
46
46
  self._span = span
47
47
  self._span_name = span_name
48
48
  # Placeholder for aggregating streaming response
49
- self._llmresponse = ''
50
- self._response_id = ''
51
- self._finish_reason = ''
49
+ self._llmresponse = ""
50
+ self._response_id = ""
51
+ self._finish_reason = ""
52
52
  self._input_tokens = 0
53
53
  self._output_tokens = 0
54
54
  self._choices = []
@@ -92,14 +92,13 @@ def async_chat(version, environment, application_name,
92
92
  environment=environment,
93
93
  application_name=application_name,
94
94
  metrics=metrics,
95
- event_provider=event_provider,
96
95
  capture_message_content=capture_message_content,
97
96
  disable_metrics=disable_metrics,
98
97
  version=version
99
98
  )
100
99
  except Exception as e:
101
100
  handle_exception(self._span, e)
102
- logger.error('Error in trace creation: %s', e)
101
+
103
102
  raise
104
103
 
105
104
  async def wrapper(wrapped, instance, args, kwargs):
@@ -108,12 +107,12 @@ def async_chat(version, environment, application_name,
108
107
  """
109
108
 
110
109
  # Check if streaming is enabled for the API call
111
- streaming = kwargs.get('stream', False)
110
+ streaming = kwargs.get("stream", False)
112
111
 
113
- server_address, server_port = set_server_address_and_port(instance, 'api.ai21.com', 443)
114
- request_model = kwargs.get('model', 'jamba-1.5-mini')
112
+ server_address, server_port = set_server_address_and_port(instance, "api.ai21.com", 443)
113
+ request_model = kwargs.get("model", "jamba-1.5-mini")
115
114
 
116
- span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}'
115
+ span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
117
116
 
118
117
  # pylint: disable=no-else-return
119
118
  if streaming:
@@ -127,30 +126,34 @@ def async_chat(version, environment, application_name,
127
126
  with tracer.start_as_current_span(span_name, kind= SpanKind.CLIENT) as span:
128
127
  start_time = time.time()
129
128
  response = await wrapped(*args, **kwargs)
130
- response = process_chat_response(
131
- response=response,
132
- request_model=request_model,
133
- pricing_info=pricing_info,
134
- server_port=server_port,
135
- server_address=server_address,
136
- environment=environment,
137
- application_name=application_name,
138
- metrics=metrics,
139
- event_provider=event_provider,
140
- start_time=start_time,
141
- span=span,
142
- capture_message_content=capture_message_content,
143
- disable_metrics=disable_metrics,
144
- version=version,
145
- **kwargs
146
- )
147
129
 
148
- return response
130
+ try:
131
+ response = process_chat_response(
132
+ response=response,
133
+ request_model=request_model,
134
+ pricing_info=pricing_info,
135
+ server_port=server_port,
136
+ server_address=server_address,
137
+ environment=environment,
138
+ application_name=application_name,
139
+ metrics=metrics,
140
+ start_time=start_time,
141
+ span=span,
142
+ capture_message_content=capture_message_content,
143
+ disable_metrics=disable_metrics,
144
+ version=version,
145
+ **kwargs
146
+ )
147
+
148
+ except Exception as e:
149
+ handle_exception(span, e)
150
+
151
+ return response
149
152
 
150
153
  return wrapper
151
154
 
152
155
  def async_chat_rag(version, environment, application_name,
153
- tracer, event_provider, pricing_info, capture_message_content, metrics, disable_metrics):
156
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics):
154
157
  """
155
158
  Generates a telemetry wrapper for GenAI function call
156
159
  """
@@ -160,32 +163,36 @@ def async_chat_rag(version, environment, application_name,
160
163
  Wraps the GenAI function call.
161
164
  """
162
165
 
163
- server_address, server_port = set_server_address_and_port(instance, 'api.ai21.com', 443)
164
- request_model = kwargs.get('model', 'jamba-1.5-mini')
166
+ server_address, server_port = set_server_address_and_port(instance, "api.ai21.com", 443)
167
+ request_model = kwargs.get("model", "jamba-1.5-mini")
165
168
 
166
- span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}'
169
+ span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
167
170
 
168
171
  with tracer.start_as_current_span(span_name, kind= SpanKind.CLIENT) as span:
169
172
  start_time = time.time()
170
173
  response = await wrapped(*args, **kwargs)
171
- response = process_chat_rag_response(
172
- response=response,
173
- request_model=request_model,
174
- pricing_info=pricing_info,
175
- server_port=server_port,
176
- server_address=server_address,
177
- environment=environment,
178
- application_name=application_name,
179
- metrics=metrics,
180
- event_provider=event_provider,
181
- start_time=start_time,
182
- span=span,
183
- capture_message_content=capture_message_content,
184
- disable_metrics=disable_metrics,
185
- version=version,
186
- **kwargs
187
- )
188
-
189
- return response
174
+
175
+ try:
176
+ response = process_chat_rag_response(
177
+ response=response,
178
+ request_model=request_model,
179
+ pricing_info=pricing_info,
180
+ server_port=server_port,
181
+ server_address=server_address,
182
+ environment=environment,
183
+ application_name=application_name,
184
+ metrics=metrics,
185
+ start_time=start_time,
186
+ span=span,
187
+ capture_message_content=capture_message_content,
188
+ disable_metrics=disable_metrics,
189
+ version=version,
190
+ **kwargs
191
+ )
192
+
193
+ except Exception as e:
194
+ handle_exception(span, e)
195
+
196
+ return response
190
197
 
191
198
  return wrapper