openlit 1.32.2__py3-none-any.whl → 1.32.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 (12) hide show
  1. openlit/__init__.py +6 -0
  2. openlit/instrumentation/ai21/__init__.py +66 -0
  3. openlit/instrumentation/ai21/ai21.py +604 -0
  4. openlit/instrumentation/ai21/async_ai21.py +603 -0
  5. openlit/instrumentation/astra/__init__.py +179 -0
  6. openlit/instrumentation/astra/astra.py +226 -0
  7. openlit/instrumentation/astra/async_astra.py +226 -0
  8. openlit/semcov/__init__.py +14 -0
  9. {openlit-1.32.2.dist-info → openlit-1.32.4.dist-info}/METADATA +15 -16
  10. {openlit-1.32.2.dist-info → openlit-1.32.4.dist-info}/RECORD +12 -6
  11. {openlit-1.32.2.dist-info → openlit-1.32.4.dist-info}/LICENSE +0 -0
  12. {openlit-1.32.2.dist-info → openlit-1.32.4.dist-info}/WHEEL +0 -0
openlit/instrumentation/astra/__init__.py ADDED
@@ -0,0 +1,179 @@
1
+ # pylint: disable=useless-return, bad-staticmethod-argument, disable=duplicate-code
2
+ """Initializer of Auto Instrumentation of AstraDB 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.astra.astra import general_wrap
9
+ # from openlit.instrumentation.astra.async_astra import asyc_general_wrap
10
+
11
+ _instruments = ("astrapy >= 1.5.2",)
12
+
13
+ class AstraInstrumentor(BaseInstrumentor):
14
+ """An instrumentor for AstraDB's client library."""
15
+
16
+ def instrumentation_dependencies(self) -> Collection[str]:
17
+ return _instruments
18
+
19
+ def _instrument(self, **kwargs):
20
+ application_name = kwargs.get("application_name")
21
+ environment = kwargs.get("environment")
22
+ tracer = kwargs.get("tracer")
23
+ metrics = kwargs.get("metrics_dict")
24
+ pricing_info = kwargs.get("pricing_info")
25
+ trace_content = kwargs.get("trace_content")
26
+ disable_metrics = kwargs.get("disable_metrics")
27
+ version = importlib.metadata.version("astrapy")
28
+
29
+ # Sync
30
+ wrap_function_wrapper(
31
+ "astrapy.database",
32
+ "Database.create_collection",
33
+ general_wrap("astra.create_collection", version, environment, application_name,
34
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
35
+ )
36
+ wrap_function_wrapper(
37
+ "astrapy.database",
38
+ "Database.drop_collection",
39
+ general_wrap("astra.drop_collection", version, environment, application_name,
40
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
41
+ )
42
+ wrap_function_wrapper(
43
+ "astrapy.collection",
44
+ "Collection.insert_one",
45
+ general_wrap("astra.insert_one", version, environment, application_name,
46
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
47
+ )
48
+ wrap_function_wrapper(
49
+ "astrapy.collection",
50
+ "Collection.insert_many",
51
+ general_wrap("astra.insert_many", version, environment, application_name,
52
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
53
+ )
54
+ wrap_function_wrapper(
55
+ "astrapy.collection",
56
+ "Collection.update_one",
57
+ general_wrap("astra.update_one", version, environment, application_name,
58
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
59
+ )
60
+ wrap_function_wrapper(
61
+ "astrapy.collection",
62
+ "Collection.update_many",
63
+ general_wrap("astra.update_many", version, environment, application_name,
64
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
65
+ )
66
+ wrap_function_wrapper(
67
+ "astrapy.collection",
68
+ "Collection.find_one_and_update",
69
+ general_wrap("astra.find_one_and_update", version, environment, application_name,
70
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
71
+ )
72
+ wrap_function_wrapper(
73
+ "astrapy.collection",
74
+ "Collection.find",
75
+ general_wrap("astra.find", version, environment, application_name,
76
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
77
+ )
78
+ wrap_function_wrapper(
79
+ "astrapy.collection",
80
+ "Collection.replace_one",
81
+ general_wrap("astra.replace_one", version, environment, application_name,
82
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
83
+ )
84
+ wrap_function_wrapper(
85
+ "astrapy.collection",
86
+ "Collection.find_one_and_delete",
87
+ general_wrap("astra.find_one_and_delete", version, environment, application_name,
88
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
89
+ )
90
+ wrap_function_wrapper(
91
+ "astrapy.collection",
92
+ "Collection.delete_one",
93
+ general_wrap("astra.delete_one", version, environment, application_name,
94
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
95
+ )
96
+ wrap_function_wrapper(
97
+ "astrapy.collection",
98
+ "Collection.delete_many",
99
+ general_wrap("astra.delete_many", version, environment, application_name,
100
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
101
+ )
102
+
103
+ # ASync
104
+ wrap_function_wrapper(
105
+ "astrapy.database",
106
+ "AsyncDatabase.create_collection",
107
+ general_wrap("astra.create_collection", version, environment, application_name,
108
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
109
+ )
110
+ wrap_function_wrapper(
111
+ "astrapy.database",
112
+ "AsyncDatabase.drop_collection",
113
+ general_wrap("astra.drop_collection", version, environment, application_name,
114
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
115
+ )
116
+ wrap_function_wrapper(
117
+ "astrapy.collection",
118
+ "AsyncCollection.insert_one",
119
+ general_wrap("astra.insert_one", version, environment, application_name,
120
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
121
+ )
122
+ wrap_function_wrapper(
123
+ "astrapy.collection",
124
+ "AsyncCollection.insert_many",
125
+ general_wrap("astra.insert_many", version, environment, application_name,
126
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
127
+ )
128
+ wrap_function_wrapper(
129
+ "astrapy.collection",
130
+ "AsyncCollection.update_one",
131
+ general_wrap("astra.update_one", version, environment, application_name,
132
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
133
+ )
134
+ wrap_function_wrapper(
135
+ "astrapy.collection",
136
+ "AsyncCollection.update_many",
137
+ general_wrap("astra.update_many", version, environment, application_name,
138
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
139
+ )
140
+ wrap_function_wrapper(
141
+ "astrapy.collection",
142
+ "AsyncCollection.find_one_and_update",
143
+ general_wrap("astra.find_one_and_update", version, environment, application_name,
144
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
145
+ )
146
+ wrap_function_wrapper(
147
+ "astrapy.collection",
148
+ "AsyncCollection.find",
149
+ general_wrap("astra.find", version, environment, application_name,
150
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
151
+ )
152
+ wrap_function_wrapper(
153
+ "astrapy.collection",
154
+ "AsyncCollection.replace_one",
155
+ general_wrap("astra.replace_one", version, environment, application_name,
156
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
157
+ )
158
+ wrap_function_wrapper(
159
+ "astrapy.collection",
160
+ "AsyncCollection.find_one_and_delete",
161
+ general_wrap("astra.find_one_and_delete", version, environment, application_name,
162
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
163
+ )
164
+ wrap_function_wrapper(
165
+ "astrapy.collection",
166
+ "AsyncCollection.delete_one",
167
+ general_wrap("astra.delete_one", version, environment, application_name,
168
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
169
+ )
170
+ wrap_function_wrapper(
171
+ "astrapy.collection",
172
+ "AsyncCollection.delete_many",
173
+ general_wrap("astra.delete_many", version, environment, application_name,
174
+ tracer, pricing_info, trace_content, metrics, disable_metrics),
175
+ )
176
+
177
+ @staticmethod
178
+ def _uninstrument(self, **kwargs):
179
+ pass
openlit/instrumentation/astra/astra.py ADDED
@@ -0,0 +1,226 @@
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 AstraDB.
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
+ return len(obj) if obj else None
21
+
22
+ def general_wrap(gen_ai_endpoint, version, environment, application_name,
23
+ tracer, pricing_info, trace_content, metrics, disable_metrics):
24
+ """
25
+ Wraps a AstraDB operation to trace and log its execution metrics.
26
+
27
+ This function is intended to wrap around AstraDB operations in order to
28
+ measure their execution time, log relevant information, and trace the execution
29
+ using OpenTelemetry. This helps in monitoring and debugging operations within
30
+ the AstraDB space.
31
+
32
+ Parameters:
33
+ - operation (str): The specific AstraDB operation being monitored.
34
+ Examples include 'create_index', 'query', 'upsert', etc.
35
+ - version (str): The version of the application interfacing with AstraDB.
36
+ - environment (str): The deployment environment, such as 'production' or 'development'.
37
+ - application_name (str): The name of the application performing the AstraDB operation.
38
+ - tracer (opentelemetry.trace.Tracer): An object used for OpenTelemetry tracing.
39
+ - pricing_info (dict): Information about pricing, not used in current implementation.
40
+ - trace_content (bool): A flag indicating whether the content of responses should be traced.
41
+
42
+ Returns:
43
+ - function: A decorator function that, when applied, wraps the target function with
44
+ additional functionality for tracing and logging AstraDB operations.
45
+ """
46
+
47
+ def wrapper(wrapped, instance, args, kwargs):
48
+ """
49
+ Executes the wrapped AstraDB operation, adding tracing and logging.
50
+
51
+ This inner wrapper function captures the execution of AstraDB operations,
52
+ annotating the operation with relevant metrics and tracing information, and
53
+ ensuring any exceptions are caught and logged appropriately.
54
+
55
+ Parameters:
56
+ - wrapped (Callable): The AstraDB operation to be wrapped and executed.
57
+ - instance (object): The instance on which the operation is called (for class methods).
58
+ - args (tuple): Positional arguments for the AstraDB operation.
59
+ - kwargs (dict): Keyword arguments for the AstraDB operation.
60
+
61
+ Returns:
62
+ - Any: The result of executing the wrapped AstraDB operation.
63
+ """
64
+
65
+ with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
66
+ response = wrapped(*args, **kwargs)
67
+
68
+ try:
69
+ span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
70
+ span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
71
+ gen_ai_endpoint)
72
+ span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
73
+ environment)
74
+ span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
75
+ application_name)
76
+ span.set_attribute(SemanticConvetion.GEN_AI_TYPE,
77
+ SemanticConvetion.GEN_AI_TYPE_VECTORDB)
78
+ span.set_attribute(SemanticConvetion.DB_SYSTEM,
79
+ SemanticConvetion.DB_SYSTEM_ASTRA)
80
+
81
+ if gen_ai_endpoint == "astra.create_collection":
82
+ db_operation = SemanticConvetion.DB_OPERATION_CREATE_COLLECTION
83
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
84
+ db_operation)
85
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
86
+ response.name)
87
+ span.set_attribute(SemanticConvetion.DB_INDEX_DIMENSION,
88
+ kwargs.get("dimension", ""))
89
+ span.set_attribute(SemanticConvetion.DB_INDEX_METRIC,
90
+ kwargs.get("metric", ""))
91
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
92
+ instance.api_endpoint)
93
+
94
+ elif gen_ai_endpoint == "astra.drop_collection":
95
+ db_operation = SemanticConvetion.DB_OPERATION_DELETE_COLLECTION
96
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
97
+ db_operation)
98
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
99
+ kwargs.get("name_or_collection", ""))
100
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
101
+ instance.api_endpoint)
102
+
103
+ elif gen_ai_endpoint == "astra.insert_one":
104
+ db_operation = SemanticConvetion.DB_OPERATION_INSERT
105
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
106
+ db_operation)
107
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
108
+ instance.name)
109
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
110
+ 1)
111
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
112
+ instance.database.api_endpoint)
113
+ span.set_attribute(SemanticConvetion.DB_OPERATION_ID,
114
+ response.inserted_id)
115
+
116
+ elif gen_ai_endpoint == "astra.insert_many":
117
+ db_operation = SemanticConvetion.DB_OPERATION_INSERT
118
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
119
+ db_operation)
120
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
121
+ object_count(args[0]))
122
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
123
+ instance.name)
124
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
125
+ instance.database.api_endpoint)
126
+
127
+ elif gen_ai_endpoint in ["astra.update_one", "astra.update_many"]:
128
+ db_operation = SemanticConvetion.DB_OPERATION_UPDATE
129
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
130
+ db_operation)
131
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
132
+ instance.name)
133
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
134
+ instance.database.api_endpoint)
135
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
136
+ response.update_info.get("nModified", 0))
137
+
138
+ elif gen_ai_endpoint == "astra.find_one_and_update":
139
+ db_operation = SemanticConvetion.DB_OPERATION_UPDATE
140
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
141
+ db_operation)
142
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
143
+ instance.name)
144
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
145
+ instance.database.api_endpoint)
146
+
147
+ elif gen_ai_endpoint == "astra.find":
148
+ db_operation = SemanticConvetion.DB_OPERATION_QUERY
149
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
150
+ db_operation)
151
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
152
+ instance.name)
153
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
154
+ instance.database.api_endpoint)
155
+ span.set_attribute(SemanticConvetion.DB_STATEMENT,
156
+ str(args))
157
+
158
+ elif gen_ai_endpoint == "astra.replace_one":
159
+ if kwargs.get("upsert") is True:
160
+ db_operation = SemanticConvetion.DB_OPERATION_UPSERT
161
+ else:
162
+ db_operation = SemanticConvetion.DB_OPERATION_REPLACE
163
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
164
+ db_operation)
165
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
166
+ instance.name)
167
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
168
+ instance.database.api_endpoint)
169
+ span.set_attribute(SemanticConvetion.DB_STATEMENT,
170
+ str(args))
171
+
172
+ elif gen_ai_endpoint in ["astra.delete_one", "astra.delete_many"]:
173
+ db_operation = SemanticConvetion.DB_OPERATION_DELETE
174
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
175
+ db_operation)
176
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
177
+ instance.name)
178
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
179
+ instance.database.api_endpoint)
180
+ span.set_attribute(SemanticConvetion.DB_STATEMENT,
181
+ str(args))
182
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
183
+ response.deleted_count)
184
+
185
+ elif gen_ai_endpoint == "astra.find_one_and_delete":
186
+ db_operation = SemanticConvetion.DB_OPERATION_DELETE
187
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
188
+ db_operation)
189
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
190
+ instance.name)
191
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
192
+ instance.database.api_endpoint)
193
+ span.set_attribute(SemanticConvetion.DB_STATEMENT,
194
+ str(args))
195
+
196
+ span.set_status(Status(StatusCode.OK))
197
+
198
+ if disable_metrics is False:
199
+ attributes = {
200
+ TELEMETRY_SDK_NAME:
201
+ "openlit",
202
+ SemanticConvetion.GEN_AI_APPLICATION_NAME:
203
+ application_name,
204
+ SemanticConvetion.DB_SYSTEM:
205
+ SemanticConvetion.DB_SYSTEM_ASTRA,
206
+ SemanticConvetion.GEN_AI_ENVIRONMENT:
207
+ environment,
208
+ SemanticConvetion.GEN_AI_TYPE:
209
+ SemanticConvetion.GEN_AI_TYPE_VECTORDB,
210
+ SemanticConvetion.DB_OPERATION:
211
+ db_operation
212
+ }
213
+
214
+ metrics["db_requests"].add(1, attributes)
215
+
216
+ # Return original response
217
+ return response
218
+
219
+ except Exception as e:
220
+ handle_exception(span, e)
221
+ logger.error("Error in trace creation: %s", e)
222
+
223
+ # Return original response
224
+ return response
225
+
226
+ return wrapper
openlit/instrumentation/astra/async_astra.py ADDED
@@ -0,0 +1,226 @@
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 AstraDB.
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
+ return len(obj) if obj else None
21
+
22
+ def general_wrap(gen_ai_endpoint, version, environment, application_name,
23
+ tracer, pricing_info, trace_content, metrics, disable_metrics):
24
+ """
25
+ Wraps a AstraDB operation to trace and log its execution metrics.
26
+
27
+ This function is intended to wrap around AstraDB operations in order to
28
+ measure their execution time, log relevant information, and trace the execution
29
+ using OpenTelemetry. This helps in monitoring and debugging operations within
30
+ the AstraDB space.
31
+
32
+ Parameters:
33
+ - operation (str): The specific AstraDB operation being monitored.
34
+ Examples include 'create_index', 'query', 'upsert', etc.
35
+ - version (str): The version of the application interfacing with AstraDB.
36
+ - environment (str): The deployment environment, such as 'production' or 'development'.
37
+ - application_name (str): The name of the application performing the AstraDB operation.
38
+ - tracer (opentelemetry.trace.Tracer): An object used for OpenTelemetry tracing.
39
+ - pricing_info (dict): Information about pricing, not used in current implementation.
40
+ - trace_content (bool): A flag indicating whether the content of responses should be traced.
41
+
42
+ Returns:
43
+ - function: A decorator function that, when applied, wraps the target function with
44
+ additional functionality for tracing and logging AstraDB operations.
45
+ """
46
+
47
+ async def wrapper(wrapped, instance, args, kwargs):
48
+ """
49
+ Executes the wrapped AstraDB operation, adding tracing and logging.
50
+
51
+ This inner wrapper function captures the execution of AstraDB operations,
52
+ annotating the operation with relevant metrics and tracing information, and
53
+ ensuring any exceptions are caught and logged appropriately.
54
+
55
+ Parameters:
56
+ - wrapped (Callable): The AstraDB operation to be wrapped and executed.
57
+ - instance (object): The instance on which the operation is called (for class methods).
58
+ - args (tuple): Positional arguments for the AstraDB operation.
59
+ - kwargs (dict): Keyword arguments for the AstraDB operation.
60
+
61
+ Returns:
62
+ - Any: The result of executing the wrapped AstraDB operation.
63
+ """
64
+
65
+ with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
66
+ response = await wrapped(*args, **kwargs)
67
+
68
+ try:
69
+ span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
70
+ span.set_attribute(SemanticConvetion.GEN_AI_ENDPOINT,
71
+ gen_ai_endpoint)
72
+ span.set_attribute(SemanticConvetion.GEN_AI_ENVIRONMENT,
73
+ environment)
74
+ span.set_attribute(SemanticConvetion.GEN_AI_APPLICATION_NAME,
75
+ application_name)
76
+ span.set_attribute(SemanticConvetion.GEN_AI_TYPE,
77
+ SemanticConvetion.GEN_AI_TYPE_VECTORDB)
78
+ span.set_attribute(SemanticConvetion.DB_SYSTEM,
79
+ SemanticConvetion.DB_SYSTEM_ASTRA)
80
+
81
+ if gen_ai_endpoint == "astra.create_collection":
82
+ db_operation = SemanticConvetion.DB_OPERATION_CREATE_COLLECTION
83
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
84
+ db_operation)
85
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
86
+ response.name)
87
+ span.set_attribute(SemanticConvetion.DB_INDEX_DIMENSION,
88
+ kwargs.get("dimension", ""))
89
+ span.set_attribute(SemanticConvetion.DB_INDEX_METRIC,
90
+ kwargs.get("metric", ""))
91
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
92
+ instance.api_endpoint)
93
+
94
+ elif gen_ai_endpoint == "astra.drop_collection":
95
+ db_operation = SemanticConvetion.DB_OPERATION_DELETE_COLLECTION
96
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
97
+ db_operation)
98
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
99
+ kwargs.get("name_or_collection", ""))
100
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
101
+ instance.api_endpoint)
102
+
103
+ elif gen_ai_endpoint == "astra.insert_one":
104
+ db_operation = SemanticConvetion.DB_OPERATION_INSERT
105
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
106
+ db_operation)
107
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
108
+ instance.name)
109
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
110
+ 1)
111
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
112
+ instance.database.api_endpoint)
113
+ span.set_attribute(SemanticConvetion.DB_OPERATION_ID,
114
+ response.inserted_id)
115
+
116
+ elif gen_ai_endpoint == "astra.insert_many":
117
+ db_operation = SemanticConvetion.DB_OPERATION_INSERT
118
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
119
+ db_operation)
120
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
121
+ object_count(args[0]))
122
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
123
+ instance.name)
124
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
125
+ instance.database.api_endpoint)
126
+
127
+ elif gen_ai_endpoint in ["astra.update_one", "astra.update_many"]:
128
+ db_operation = SemanticConvetion.DB_OPERATION_UPDATE
129
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
130
+ db_operation)
131
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
132
+ instance.name)
133
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
134
+ instance.database.api_endpoint)
135
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
136
+ response.update_info.get("nModified", 0))
137
+
138
+ elif gen_ai_endpoint == "astra.find_one_and_update":
139
+ db_operation = SemanticConvetion.DB_OPERATION_UPDATE
140
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
141
+ db_operation)
142
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
143
+ instance.name)
144
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
145
+ instance.database.api_endpoint)
146
+
147
+ elif gen_ai_endpoint == "astra.find":
148
+ db_operation = SemanticConvetion.DB_OPERATION_QUERY
149
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
150
+ db_operation)
151
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
152
+ instance.name)
153
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
154
+ instance.database.api_endpoint)
155
+ span.set_attribute(SemanticConvetion.DB_STATEMENT,
156
+ str(args))
157
+
158
+ elif gen_ai_endpoint == "astra.replace_one":
159
+ if kwargs.get("upsert") is True:
160
+ db_operation = SemanticConvetion.DB_OPERATION_UPSERT
161
+ else:
162
+ db_operation = SemanticConvetion.DB_OPERATION_REPLACE
163
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
164
+ db_operation)
165
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
166
+ instance.name)
167
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
168
+ instance.database.api_endpoint)
169
+ span.set_attribute(SemanticConvetion.DB_STATEMENT,
170
+ str(args))
171
+
172
+ elif gen_ai_endpoint in ["astra.delete_one", "astra.delete_many"]:
173
+ db_operation = SemanticConvetion.DB_OPERATION_DELETE
174
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
175
+ db_operation)
176
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
177
+ instance.name)
178
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
179
+ instance.database.api_endpoint)
180
+ span.set_attribute(SemanticConvetion.DB_STATEMENT,
181
+ str(args))
182
+ span.set_attribute(SemanticConvetion.DB_DOCUMENTS_COUNT,
183
+ response.deleted_count)
184
+
185
+ elif gen_ai_endpoint == "astra.find_one_and_delete":
186
+ db_operation = SemanticConvetion.DB_OPERATION_DELETE
187
+ span.set_attribute(SemanticConvetion.DB_OPERATION,
188
+ db_operation)
189
+ span.set_attribute(SemanticConvetion.DB_COLLECTION_NAME,
190
+ instance.name)
191
+ span.set_attribute(SemanticConvetion.DB_OPERATION_API_ENDPOINT,
192
+ instance.database.api_endpoint)
193
+ span.set_attribute(SemanticConvetion.DB_STATEMENT,
194
+ str(args))
195
+
196
+ span.set_status(Status(StatusCode.OK))
197
+
198
+ if disable_metrics is False:
199
+ attributes = {
200
+ TELEMETRY_SDK_NAME:
201
+ "openlit",
202
+ SemanticConvetion.GEN_AI_APPLICATION_NAME:
203
+ application_name,
204
+ SemanticConvetion.DB_SYSTEM:
205
+ SemanticConvetion.DB_SYSTEM_ASTRA,
206
+ SemanticConvetion.GEN_AI_ENVIRONMENT:
207
+ environment,
208
+ SemanticConvetion.GEN_AI_TYPE:
209
+ SemanticConvetion.GEN_AI_TYPE_VECTORDB,
210
+ SemanticConvetion.DB_OPERATION:
211
+ db_operation
212
+ }
213
+
214
+ metrics["db_requests"].add(1, attributes)
215
+
216
+ # Return original response
217
+ return response
218
+
219
+ except Exception as e:
220
+ handle_exception(span, e)
221
+ logger.error("Error in trace creation: %s", e)
222
+
223
+ # Return original response
224
+ return response
225
+
226
+ return wrapper