openlit 0.0.1__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,291 @@
1
+ # pylint: disable=duplicate-code, broad-exception-caught, too-many-statements, unused-argument
2
+ """
3
+ Module for monitoring Anthropic 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 get_chat_model_cost, handle_exception
10
+ from openlit.semcov import SemanticConvetion
11
+
12
+ # Initialize logger for logging potential issues and operations
13
+ logger = logging.getLogger(__name__)
14
+
15
+ def async_messages(gen_ai_endpoint, version, environment, application_name,
16
+ tracer, pricing_info, trace_content, metrics, disable_metrics):
17
+ """
18
+ Generates a telemetry wrapper for messages to collect metrics.
19
+
20
+ Args:
21
+ gen_ai_endpoint: Endpoint identifier for logging and tracing.
22
+ version: Version of the monitoring package.
23
+ environment: Deployment environment (e.g., production, staging).
24
+ application_name: Name of the application using the OpenAI API.
25
+ tracer: OpenTelemetry tracer for creating spans.
26
+ pricing_info: Information used for calculating the cost of OpenAI usage.
27
+ trace_content: Flag indicating whether to trace the actual content.
28
+
29
+ Returns:
30
+ A function that wraps the chat method to add telemetry.
31
+ """
32
+
33
+ async def wrapper(wrapped, instance, args, kwargs):
34
+ """
35
+ Wraps the 'messages' API call to add telemetry.
36
+
37
+ This collects metrics such as execution time, cost, and token usage, and handles errors
38
+ gracefully, adding details to the trace for observability.
39
+
40
+ Args:
41
+ wrapped: The original 'messages' method to be wrapped.
42
+ instance: The instance of the class where the original method is defined.
43
+ args: Positional arguments for the 'messages' method.
44
+ kwargs: Keyword arguments for the 'messages' method.
45
+
46
+ Returns:
47
+ The response from the original 'messages' method.
48
+ """
49
+
50
+ # Check if streaming is enabled for the API call
51
+ streaming = kwargs.get("stream", False)
52
+
53
+ # pylint: disable=no-else-return
54
+ if streaming:
55
+ # Special handling for streaming response to accommodate the nature of data flow
56
+ async def stream_generator():
57
+ with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
58
+ # Placeholder for aggregating streaming response
59
+ llmresponse = ""
60
+
61
+ # Loop through streaming events capturing relevant details
62
+ async for event in await wrapped(*args, **kwargs):
63
+
64
+ # Collect message IDs and input token from events
65
+ if event.type == "message_start":
66
+ response_id = event.message.id
67
+ prompt_tokens = event.message.usage.input_tokens
68
+
69
+ # Aggregate response content
70
+ if event.type == "content_block_delta":
71
+ llmresponse += event.delta.text
72
+
73
+ # Collect output tokens and stop reason from events
74
+ if event.type == "message_delta":
75
+ completion_tokens = event.usage.output_tokens
76
+ finish_reason = event.delta.stop_reason
77
+ yield event
78
+
79
+ # Handling exception ensure observability without disrupting operation
80
+ try:
81
+ # Format 'messages' into a single string
82
+ message_prompt = kwargs.get("messages", "")
83
+ formatted_messages = []
84
+ for message in message_prompt:
85
+ role = message["role"]
86
+ content = message["content"]
87
+
88
+ if isinstance(content, list):
89
+ content_str = ", ".join(
90
+ # pylint: disable=line-too-long
91
+ f'{item["type"]}: {item["text"] if "text" in item else item["image_url"]}'
92
+ if "type" in item else f'text: {item["text"]}'
93
+ for item in content
94
+ )
95
+ formatted_messages.append(f"{role}: {content_str}")
96
+ else:
97
+ formatted_messages.append(f"{role}: {content}")
98
+ prompt = "\n".join(formatted_messages)
99
+
100
+ # Calculate cost of the operation
101
+ cost = get_chat_model_cost(
102
+ kwargs.get("model", "claude-3-sonnet-20240229"),
103
+ pricing_info, prompt_tokens, completion_tokens
104
+ )
105
+
106
+ # Set Span attributes
107
+ span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
108
+ span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM,
109
+ SemanticConvetion.GEN_AI_SYSTEM_ANTHROPIC)
110
+ span.set_attribute(SemanticConvetion.GEN_AI_TYPE,
111
+ SemanticConvetion.GEN_AI_TYPE_CHAT)
112
+ span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
113
+ gen_ai_endpoint)
114
+ span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_ID,
115
+ response_id)
116
+ span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
117
+ environment)
118
+ span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
119
+ application_name)
120
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL,
121
+ kwargs.get("model", "claude-3-sonnet-20240229"))
122
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MAX_TOKENS,
123
+ kwargs.get("max_tokens", ""))
124
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_IS_STREAM,
125
+ True)
126
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE,
127
+ kwargs.get("temperature", 1.0))
128
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TOP_P,
129
+ kwargs.get("top_p", ""))
130
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TOP_K,
131
+ kwargs.get("top_k", ""))
132
+ span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_FINISH_REASON,
133
+ finish_reason)
134
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
135
+ prompt_tokens)
136
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
137
+ completion_tokens)
138
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
139
+ prompt_tokens + completion_tokens)
140
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST,
141
+ cost)
142
+ if trace_content:
143
+ span.set_attribute(SemanticConvetion.GEN_AI_CONTENT_PROMPT,
144
+ prompt)
145
+ span.set_attribute(SemanticConvetion.GEN_AI_CONTENT_COMPLETION,
146
+ llmresponse)
147
+
148
+ span.set_status(Status(StatusCode.OK))
149
+
150
+ if disable_metrics is False:
151
+ attributes = {
152
+ TELEMETRY_SDK_NAME:
153
+ "openlit",
154
+ SemanticConvetion.GEN_AI_APPLICATION_NAME:
155
+ application_name,
156
+ SemanticConvetion.GEN_AI_SYSTEM:
157
+ SemanticConvetion.GEN_AI_SYSTEM_ANTHROPIC,
158
+ SemanticConvetion.GEN_AI_ENVIRONMENT:
159
+ environment,
160
+ SemanticConvetion.GEN_AI_TYPE:
161
+ SemanticConvetion.GEN_AI_TYPE_CHAT,
162
+ SemanticConvetion.GEN_AI_REQUEST_MODEL:
163
+ kwargs.get("model", "claude-3-sonnet-20240229")
164
+ }
165
+
166
+ metrics["genai_requests"].add(1, attributes)
167
+ metrics["genai_total_tokens"].add(
168
+ prompt_tokens + completion_tokens, attributes
169
+ )
170
+ metrics["genai_completion_tokens"].add(completion_tokens, attributes)
171
+ metrics["genai_prompt_tokens"].add(prompt_tokens, attributes)
172
+ metrics["genai_cost"].record(cost, attributes)
173
+
174
+ except Exception as e:
175
+ handle_exception(span, e)
176
+ logger.error("Error in trace creation: %s", e)
177
+
178
+ return stream_generator()
179
+
180
+ # Handling for non-streaming responses
181
+ else:
182
+ with tracer.start_as_current_span(gen_ai_endpoint, kind=SpanKind.CLIENT) as span:
183
+ response = await wrapped(*args, **kwargs)
184
+
185
+ try:
186
+ # Format 'messages' into a single string
187
+ message_prompt = kwargs.get("messages", "")
188
+ formatted_messages = []
189
+ for message in message_prompt:
190
+ role = message["role"]
191
+ content = message["content"]
192
+
193
+ if isinstance(content, list):
194
+ content_str = ", ".join(
195
+ # pylint: disable=line-too-long
196
+ f'{item["type"]}: {item["text"] if "text" in item else item["image_url"]}'
197
+ if "type" in item else f'text: {item["text"]}'
198
+ for item in content
199
+ )
200
+ formatted_messages.append(f"{role}: {content_str}")
201
+ else:
202
+ formatted_messages.append(f"{role}: {content}")
203
+ prompt = "\n".join(formatted_messages)
204
+
205
+ # Calculate cost of the operation
206
+ cost = get_chat_model_cost(kwargs.get("model", "claude-3-sonnet-20240229"),
207
+ pricing_info, response.usage.input_tokens,
208
+ response.usage.output_tokens)
209
+
210
+ # Set Span attribues
211
+ span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
212
+ span.set_attribute(SemanticConvetion.GEN_AI_SYSTEM,
213
+ SemanticConvetion.GEN_AI_SYSTEM_ANTHROPIC)
214
+ span.set_attribute(SemanticConvetion.GEN_AI_TYPE,
215
+ SemanticConvetion.GEN_AI_TYPE_CHAT)
216
+ span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
217
+ gen_ai_endpoint)
218
+ span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_ID,
219
+ response.id)
220
+ span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
221
+ environment)
222
+ span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
223
+ application_name)
224
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MODEL,
225
+ kwargs.get("model", "claude-3-sonnet-20240229"))
226
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_MAX_TOKENS,
227
+ kwargs.get("max_tokens", ""))
228
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_IS_STREAM,
229
+ False)
230
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TEMPERATURE,
231
+ kwargs.get("temperature", 1.0))
232
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TOP_P,
233
+ kwargs.get("top_p", ""))
234
+ span.set_attribute(SemanticConvetion.GEN_AI_REQUEST_TOP_K,
235
+ kwargs.get("top_k", ""))
236
+ span.set_attribute(SemanticConvetion.GEN_AI_RESPONSE_FINISH_REASON,
237
+ response.stop_reason)
238
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_PROMPT_TOKENS,
239
+ response.usage.input_tokens)
240
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COMPLETION_TOKENS,
241
+ response.usage.output_tokens)
242
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_TOTAL_TOKENS,
243
+ response.usage.input_tokens +
244
+ response.usage.output_tokens)
245
+ span.set_attribute(SemanticConvetion.GEN_AI_USAGE_COST,
246
+ cost)
247
+ if trace_content:
248
+ span.set_attribute(SemanticConvetion.GEN_AI_CONTENT_PROMPT,
249
+ prompt)
250
+ # pylint: disable=line-too-long
251
+ span.set_attribute(SemanticConvetion.GEN_AI_CONTENT_COMPLETION, response.content[0].text if response.content else "")
252
+
253
+ span.set_status(Status(StatusCode.OK))
254
+
255
+ if disable_metrics is False:
256
+ attributes = {
257
+ TELEMETRY_SDK_NAME:
258
+ "openlit",
259
+ SemanticConvetion.GEN_AI_APPLICATION_NAME:
260
+ application_name,
261
+ SemanticConvetion.GEN_AI_SYSTEM:
262
+ SemanticConvetion.GEN_AI_SYSTEM_ANTHROPIC,
263
+ SemanticConvetion.GEN_AI_ENVIRONMENT:
264
+ environment,
265
+ SemanticConvetion.GEN_AI_TYPE:
266
+ SemanticConvetion.GEN_AI_TYPE_CHAT,
267
+ SemanticConvetion.GEN_AI_REQUEST_MODEL:
268
+ kwargs.get("model", "claude-3-sonnet-20240229")
269
+ }
270
+
271
+ metrics["genai_requests"].add(1, attributes)
272
+ metrics["genai_total_tokens"].add(
273
+ response.usage.input_tokens +
274
+ response.usage.output_tokens, attributes)
275
+ metrics["genai_completion_tokens"].add(
276
+ response.usage.output_tokens, attributes)
277
+ metrics["genai_prompt_tokens"].add(
278
+ response.usage.input_tokens, attributes)
279
+ metrics["genai_cost"].record(cost, attributes)
280
+
281
+ # Return original response
282
+ return response
283
+
284
+ except Exception as e:
285
+ handle_exception(span, e)
286
+ logger.error("Error in trace creation: %s", e)
287
+
288
+ # Return original response
289
+ return response
290
+
291
+ return wrapper
@@ -0,0 +1,86 @@
1
+ # pylint: disable=useless-return, bad-staticmethod-argument, disable=duplicate-code
2
+ """Initializer of Auto Instrumentation of Chromadb Functions"""
3
+ from typing import Collection
4
+ import importlib.metadata
5
+ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
6
+ from wrapt import wrap_function_wrapper
7
+
8
+ from openlit.instrumentation.chroma.chroma import general_wrap
9
+
10
+ _instruments = ("chromadb >= 0.4.0",)
11
+
12
+ class ChromaInstrumentor(BaseInstrumentor):
13
+ """An instrumentor for Chromadb's client library."""
14
+
15
+ def instrumentation_dependencies(self) -> Collection[str]:
16
+ return _instruments
17
+
18
+ def _instrument(self, **kwargs):
19
+ application_name = kwargs.get("application_name")
20
+ environment = kwargs.get("environment")
21
+ tracer = kwargs.get("tracer")
22
+ metrics = kwargs.get("metrics_dict")
23
+ pricing_info = kwargs.get("pricing_info")
24
+ trace_content = kwargs.get("trace_content")
25
+ disable_metrics = kwargs.get("disable_metrics")
26
+ version = importlib.metadata.version("chromadb")
27
+
28
+ wrap_function_wrapper(
29
+ "chromadb.db",
30
+ "DB.create_collection",
31
+ general_wrap("chroma.create_collection", version, environment, application_name,
32
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
33
+ )
34
+
35
+ wrap_function_wrapper(
36
+ "chromadb",
37
+ "Collection.add",
38
+ general_wrap("chroma.add", version, environment, application_name,
39
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
40
+ )
41
+
42
+ wrap_function_wrapper(
43
+ "chromadb",
44
+ "Collection.get",
45
+ general_wrap("chroma.get", version, environment, application_name,
46
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
47
+ )
48
+
49
+ wrap_function_wrapper(
50
+ "chromadb",
51
+ "Collection.peek",
52
+ general_wrap("chroma.peek", version, environment, application_name,
53
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
54
+ )
55
+
56
+ wrap_function_wrapper(
57
+ "chromadb",
58
+ "Collection.query",
59
+ general_wrap("chroma.query", version, environment, application_name,
60
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
61
+ )
62
+
63
+ wrap_function_wrapper(
64
+ "chromadb",
65
+ "Collection.update",
66
+ general_wrap("chroma.update", version, environment, application_name,
67
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
68
+ )
69
+
70
+ wrap_function_wrapper(
71
+ "chromadb",
72
+ "Collection.upsert",
73
+ general_wrap("chroma.upsert", version, environment, application_name,
74
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
75
+ )
76
+ wrap_function_wrapper(
77
+ "chromadb",
78
+ "Collection.delete",
79
+ general_wrap("chroma.delete", version, environment, application_name,
80
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
81
+ )
82
+
83
+
84
+ @staticmethod
85
+ def _uninstrument(self, **kwargs):
86
+ pass
@@ -0,0 +1,197 @@
1
+ # pylint: disable=duplicate-code, broad-exception-caught, too-many-statements, unused-argument
2
+ """
3
+ Module for monitoring ChromaDB.
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 handle_exception
10
+ from openlit.semcov import SemanticConvetion
11
+
12
+ # Initialize logger for logging potential issues and operations
13
+ logger = logging.getLogger(__name__)
14
+
15
+ def object_count(obj):
16
+ """
17
+ Counts Length of object if it exists, Else returns None
18
+ """
19
+
20
+ if obj:
21
+ return len(obj)
22
+
23
+ return None
24
+
25
+ def general_wrap(gen_ai_endpoint, version, environment, application_name,
26
+ tracer, pricing_info, trace_content, metrics, disable_metrics):
27
+ """
28
+ Creates a wrapper around a function call to trace and log its execution metrics.
29
+
30
+ This function wraps any given function to measure its execution time,
31
+ log its operation, and trace its execution using OpenTelemetry.
32
+
33
+ Parameters:
34
+ - gen_ai_endpoint (str): A descriptor or name for the endpoint being traced.
35
+ - version (str): The version of the Langchain application.
36
+ - environment (str): The deployment environment (e.g., 'production', 'development').
37
+ - application_name (str): Name of the Langchain application.
38
+ - tracer (opentelemetry.trace.Tracer): The tracer object used for OpenTelemetry tracing.
39
+ - pricing_info (dict): Information about the pricing for internal metrics (currently not used).
40
+ - trace_content (bool): Flag indicating whether to trace the content of the response.
41
+
42
+ Returns:
43
+ - function: A higher-order function that takes a function 'wrapped' and returns
44
+ a new function that wraps 'wrapped' with additional tracing and logging.
45
+ """
46
+
47
+ def wrapper(wrapped, instance, args, kwargs):
48
+ """
49
+ An inner wrapper function that executes the wrapped function, measures execution
50
+ time, and records trace data using OpenTelemetry.
51
+
52
+ Parameters:
53
+ - wrapped (Callable): The original function that this wrapper will execute.
54
+ - instance (object): The instance to which the wrapped function belongs. This
55
+ is used for instance methods. For static and classmethods,
56
+ this may be None.
57
+ - args (tuple): Positional arguments passed to the wrapped function.
58
+ - kwargs (dict): Keyword arguments passed to the wrapped function.
59
+
60
+ Returns:
61
+ - The result of the wrapped function call.
62
+
63
+ The wrapper initiates a span with the provided tracer, sets various attributes
64
+ on the span based on the function's execution and response, and ensures
65
+ errors are handled and logged appropriately.
66
+ """
67
+ with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
68
+ response = wrapped(*args, **kwargs)
69
+
70
+ try:
71
+ span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
72
+ span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
73
+ gen_ai_endpoint)
74
+ span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
75
+ environment)
76
+ span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
77
+ application_name)
78
+ span.set_attribute(SemanticConvetion.GEN_AI_TYPE,
79
+ SemanticConvetion.GEN_AI_TYPE_VECTORDB)
80
+ span.set_attribute(SemanticConvetion.DB_SYSTEM,
81
+ SemanticConvetion.DB_SYSTEM_CHROMA)
82
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
83
+ instance.name)
84
+
85
+ if gen_ai_endpoint == "chroma.add":
86
+ db_operation = SemanticConvetion.DB_OPERATION_ADD
87
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
88
+ SemanticConvetion.DB_OPERATION_GET)
89
+ span.set_attribute(SemanticConvetion.DB_ID_COUNT,
90
+ object_count(kwargs.get("ids")))
91
+ span.set_attribute(SemanticConvetion.DB_VECTOR_COUNT,
92
+ object_count(kwargs.get("embeddings")))
93
+ span.set_attribute(SemanticConvetion.DB_VECTOR_COUNT,
94
+ object_count(kwargs.get("metadatas")))
95
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
96
+ object_count(kwargs.get("documents")))
97
+
98
+ elif gen_ai_endpoint == "chroma.get":
99
+ db_operation = SemanticConvetion.DB_OPERATION_GET
100
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
101
+ SemanticConvetion.DB_OPERATION_GET)
102
+ span.set_attribute(SemanticConvetion.DB_ID_COUNT,
103
+ object_count(kwargs.get("ids")))
104
+ span.set_attribute(SemanticConvetion.DB_QUERY_LIMIT,
105
+ kwargs.get("limit"))
106
+ span.set_attribute(SemanticConvetion.DB_OFFSET,
107
+ kwargs.get("offset"))
108
+ span.set_attribute(SemanticConvetion.DB_WHERE_DOCUMENT,
109
+ str(kwargs.get("where_document", "")))
110
+
111
+ elif gen_ai_endpoint == "chroma.query":
112
+ db_operation = SemanticConvetion.DB_OPERATION_QUERY
113
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
114
+ SemanticConvetion.DB_OPERATION_QUERY)
115
+ span.set_attribute(SemanticConvetion.DB_STATEMENT,
116
+ str(kwargs.get("query_texts")))
117
+ span.set_attribute(SemanticConvetion.DB_N_RESULTS,
118
+ kwargs.get("n_results", ""))
119
+ span.set_attribute(SemanticConvetion.DB_FILTER,
120
+ str(kwargs.get("where", "")))
121
+ span.set_attribute(SemanticConvetion.DB_WHERE_DOCUMENT,
122
+ str(kwargs.get("where_document", "")))
123
+
124
+ elif gen_ai_endpoint == "chroma.update":
125
+ db_operation = SemanticConvetion.DB_OPERATION_UPDATE
126
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
127
+ SemanticConvetion.DB_OPERATION_UPDATE)
128
+ span.set_attribute(SemanticConvetion.DB_VECTOR_COUNT,
129
+ object_count(kwargs.get("embeddings")))
130
+ span.set_attribute(SemanticConvetion.DB_VECTOR_COUNT,
131
+ object_count(kwargs.get("metadatas")))
132
+ span.set_attribute(SemanticConvetion.DB_ID_COUNT,
133
+ object_count(kwargs.get("ids")))
134
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
135
+ object_count(kwargs.get("documents")))
136
+
137
+ elif gen_ai_endpoint == "chroma.upsert":
138
+ db_operation = SemanticConvetion.DB_OPERATION_UPSERT
139
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
140
+ SemanticConvetion.DB_OPERATION_UPSERT)
141
+ span.set_attribute(SemanticConvetion.DB_VECTOR_COUNT,
142
+ object_count(kwargs.get("embeddings")))
143
+ span.set_attribute(SemanticConvetion.DB_VECTOR_COUNT,
144
+ object_count(kwargs.get("metadatas")))
145
+ span.set_attribute(SemanticConvetion.DB_ID_COUNT,
146
+ object_count(kwargs.get("ids")))
147
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
148
+ object_count(kwargs.get("documents")))
149
+
150
+ elif gen_ai_endpoint == "chroma.delete":
151
+ db_operation = SemanticConvetion.DB_OPERATION_DELETE
152
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
153
+ SemanticConvetion.DB_OPERATION_DELETE)
154
+ span.set_attribute(SemanticConvetion.DB_ID_COUNT,
155
+ object_count(kwargs.get("ids")))
156
+ span.set_attribute(SemanticConvetion.DB_FILTER,
157
+ str(kwargs.get("where", "")))
158
+ span.set_attribute(SemanticConvetion.DB_DELETE_ALL,
159
+ kwargs.get("delete_all", False))
160
+ span.set_attribute(SemanticConvetion.DB_WHERE_DOCUMENT,
161
+ str(kwargs.get("where_document", "")))
162
+
163
+ elif gen_ai_endpoint == "chroma.peek":
164
+ db_operation = SemanticConvetion.DB_OPERATION_PEEK
165
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
166
+ SemanticConvetion.DB_OPERATION_PEEK)
167
+
168
+ span.set_status(Status(StatusCode.OK))
169
+
170
+ if disable_metrics is False:
171
+ attributes = {
172
+ TELEMETRY_SDK_NAME:
173
+ "openlit",
174
+ SemanticConvetion.GEN_AI_APPLICATION_NAME:
175
+ application_name,
176
+ SemanticConvetion.DB_SYSTEM:
177
+ SemanticConvetion.DB_SYSTEM_CHROMA,
178
+ SemanticConvetion.GEN_AI_ENVIRONMENT:
179
+ environment,
180
+ SemanticConvetion.GEN_AI_TYPE:
181
+ SemanticConvetion.GEN_AI_TYPE_VECTORDB,
182
+ SemanticConvetion.DB_OPERATION:
183
+ db_operation
184
+ }
185
+
186
+ metrics["db_requests"].add(1, attributes)
187
+
188
+ return response
189
+
190
+ except Exception as e:
191
+ handle_exception(span, e)
192
+ logger.error("Error in trace creation: %s", e)
193
+
194
+ # Return original response
195
+ return response
196
+
197
+ return wrapper
@@ -0,0 +1,51 @@
1
+ # pylint: disable=useless-return, bad-staticmethod-argument, disable=duplicate-code
2
+ """Initializer of Auto Instrumentation of Cohere Functions"""
3
+ from typing import Collection
4
+ import importlib.metadata
5
+ from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
6
+ from wrapt import wrap_function_wrapper
7
+
8
+ from openlit.instrumentation.cohere.cohere import chat, chat_stream, embed
9
+
10
+ _instruments = ("cohere >= 5.0.0",)
11
+
12
+ class CohereInstrumentor(BaseInstrumentor):
13
+ """An instrumentor for Cohere's client library."""
14
+
15
+ def instrumentation_dependencies(self) -> Collection[str]:
16
+ return _instruments
17
+
18
+ def _instrument(self, **kwargs):
19
+ application_name = kwargs.get("application_name")
20
+ environment = kwargs.get("environment")
21
+ tracer = kwargs.get("tracer")
22
+ metrics = kwargs.get("metrics_dict")
23
+ pricing_info = kwargs.get("pricing_info")
24
+ trace_content = kwargs.get("trace_content")
25
+ disable_metrics = kwargs.get("disable_metrics")
26
+ version = importlib.metadata.version("cohere")
27
+
28
+ wrap_function_wrapper(
29
+ "cohere.client",
30
+ "Client.chat",
31
+ chat("cohere.chat", version, environment, application_name,
32
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
33
+ )
34
+
35
+ wrap_function_wrapper(
36
+ "cohere.client",
37
+ "Client.chat_stream",
38
+ chat_stream("cohere.chat", version, environment, application_name,
39
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
40
+ )
41
+
42
+ wrap_function_wrapper(
43
+ "cohere.client",
44
+ "Client.embed",
45
+ embed("cohere.embed", version, environment, application_name,
46
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
47
+ )
48
+
49
+ @staticmethod
50
+ def _uninstrument(self, **kwargs):
51
+ pass