openlit 1.34.26__py3-none-any.whl → 1.34.28__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 (37) hide show
  1. openlit/__helpers.py +38 -0
  2. openlit/__init__.py +22 -155
  3. openlit/_instrumentors.py +144 -0
  4. openlit/guard/all.py +3 -3
  5. openlit/instrumentation/astra/__init__.py +71 -159
  6. openlit/instrumentation/astra/astra.py +32 -22
  7. openlit/instrumentation/astra/async_astra.py +32 -22
  8. openlit/instrumentation/astra/utils.py +263 -88
  9. openlit/instrumentation/chroma/utils.py +2 -2
  10. openlit/instrumentation/controlflow/controlflow.py +2 -2
  11. openlit/instrumentation/embedchain/embedchain.py +4 -4
  12. openlit/instrumentation/groq/__init__.py +4 -4
  13. openlit/instrumentation/haystack/__init__.py +57 -28
  14. openlit/instrumentation/haystack/async_haystack.py +54 -0
  15. openlit/instrumentation/haystack/haystack.py +35 -65
  16. openlit/instrumentation/haystack/utils.py +377 -0
  17. openlit/instrumentation/julep/async_julep.py +2 -2
  18. openlit/instrumentation/julep/julep.py +2 -2
  19. openlit/instrumentation/langchain_community/utils.py +2 -2
  20. openlit/instrumentation/llamaindex/__init__.py +165 -37
  21. openlit/instrumentation/llamaindex/async_llamaindex.py +53 -0
  22. openlit/instrumentation/llamaindex/llamaindex.py +32 -64
  23. openlit/instrumentation/llamaindex/utils.py +412 -0
  24. openlit/instrumentation/mem0/mem0.py +2 -2
  25. openlit/instrumentation/milvus/__init__.py +30 -68
  26. openlit/instrumentation/milvus/milvus.py +34 -161
  27. openlit/instrumentation/milvus/utils.py +276 -0
  28. openlit/instrumentation/openai/__init__.py +24 -24
  29. openlit/instrumentation/openai/utils.py +10 -4
  30. openlit/instrumentation/pinecone/utils.py +2 -2
  31. openlit/instrumentation/qdrant/utils.py +2 -2
  32. openlit/instrumentation/together/__init__.py +8 -8
  33. openlit/semcov/__init__.py +79 -0
  34. {openlit-1.34.26.dist-info → openlit-1.34.28.dist-info}/METADATA +1 -1
  35. {openlit-1.34.26.dist-info → openlit-1.34.28.dist-info}/RECORD +37 -31
  36. {openlit-1.34.26.dist-info → openlit-1.34.28.dist-info}/LICENSE +0 -0
  37. {openlit-1.34.26.dist-info → openlit-1.34.28.dist-info}/WHEEL +0 -0
@@ -1,102 +1,277 @@
1
1
  """
2
- Astra OpenTelemetry instrumentation utility functions
2
+ Utility functions for AstraDB instrumentation.
3
3
  """
4
4
 
5
5
  import time
6
- import logging
6
+ from urllib.parse import urlparse
7
7
  from opentelemetry.trace import Status, StatusCode
8
- from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
9
- from openlit.__helpers import handle_exception
8
+ from openlit.__helpers import (
9
+ common_db_span_attributes,
10
+ record_db_metrics,
11
+ )
10
12
  from openlit.semcov import SemanticConvention
11
13
 
12
- # Initialize logger for logging potential issues and operations
13
- logger = logging.getLogger(__name__)
14
+ # Operation mapping for simple span naming
15
+ DB_OPERATION_MAP = {
16
+ "astra.create_collection": SemanticConvention.DB_OPERATION_CREATE_COLLECTION,
17
+ "astra.drop_collection": SemanticConvention.DB_OPERATION_DELETE_COLLECTION,
18
+ "astra.insert": SemanticConvention.DB_OPERATION_INSERT,
19
+ "astra.insert_one": SemanticConvention.DB_OPERATION_INSERT,
20
+ "astra.insert_many": SemanticConvention.DB_OPERATION_INSERT,
21
+ "astra.update": SemanticConvention.DB_OPERATION_UPDATE,
22
+ "astra.update_one": SemanticConvention.DB_OPERATION_UPDATE,
23
+ "astra.update_many": SemanticConvention.DB_OPERATION_UPDATE,
24
+ "astra.find": SemanticConvention.DB_OPERATION_SELECT,
25
+ "astra.find_one_and_update": SemanticConvention.DB_OPERATION_REPLACE,
26
+ "astra.replace_one": SemanticConvention.DB_OPERATION_REPLACE,
27
+ "astra.find_one_and_delete": SemanticConvention.DB_OPERATION_FIND_AND_DELETE,
28
+ "astra.delete": SemanticConvention.DB_OPERATION_DELETE,
29
+ "astra.delete_one": SemanticConvention.DB_OPERATION_DELETE,
30
+ "astra.delete_many": SemanticConvention.DB_OPERATION_DELETE,
31
+ }
14
32
 
15
33
  def object_count(obj):
16
- """Counts Length of object if it exists, Else returns None."""
17
- return len(obj) if isinstance(obj, list) else 1
34
+ """
35
+ Counts length of object if it exists, else returns 0.
36
+ """
37
+ if isinstance(obj, list):
38
+ return len(obj)
39
+ elif obj is not None:
40
+ return 1
41
+ return 0
18
42
 
19
- DB_OPERATION_MAP = {
20
- 'astra.create_collection': SemanticConvention.DB_OPERATION_CREATE_COLLECTION,
21
- 'astra.drop_collection': SemanticConvention.DB_OPERATION_DELETE_COLLECTION,
22
- 'astra.insert': SemanticConvention.DB_OPERATION_INSERT,
23
- 'astra.update': SemanticConvention.DB_OPERATION_UPDATE,
24
- 'astra.find': SemanticConvention.DB_OPERATION_SELECT,
25
- 'astra.find_one_and_update': SemanticConvention.DB_OPERATION_REPLACE,
26
- 'astra.replace_one': SemanticConvention.DB_OPERATION_REPLACE,
27
- 'astra.delete': SemanticConvention.DB_OPERATION_DELETE,
28
- 'astra.find_one_and_delete': SemanticConvention.DB_OPERATION_FIND_AND_DELETE
29
- }
43
+ def set_server_address_and_port(instance):
44
+ """
45
+ Extracts server address and port from AstraDB client instance.
46
+
47
+ Args:
48
+ instance: AstraDB client instance
49
+
50
+ Returns:
51
+ tuple: (server_address, server_port)
52
+ """
53
+ server_address = "astra.datastax.com"
54
+ server_port = 443
55
+
56
+ # Try getting api_endpoint from instance or its database
57
+ api_endpoint = getattr(instance, "api_endpoint", None)
58
+ if not api_endpoint:
59
+ # Try getting from database attribute
60
+ database = getattr(instance, "database", None)
61
+ if database:
62
+ api_endpoint = getattr(database, "api_endpoint", None)
63
+
64
+ if api_endpoint and isinstance(api_endpoint, str):
65
+ if api_endpoint.startswith(("http://", "https://")):
66
+ try:
67
+ parsed = urlparse(api_endpoint)
68
+ server_address = parsed.hostname or server_address
69
+ server_port = parsed.port or server_port
70
+ except Exception:
71
+ pass
72
+ else:
73
+ # Handle cases like "hostname:port" or just "hostname"
74
+ if ":" in api_endpoint:
75
+ parts = api_endpoint.split(":")
76
+ server_address = parts[0]
77
+ try:
78
+ server_port = int(parts[1])
79
+ except (ValueError, IndexError):
80
+ pass
81
+ else:
82
+ server_address = api_endpoint
83
+
84
+ return server_address, server_port
30
85
 
31
- def process_db_operations(response, span, start_time, gen_ai_endpoint,
32
- version, environment, application_name,
33
- capture_message_content, metrics, disable_metrics, server_address,
34
- server_port, collection_name, db_operation, kwargs, args):
86
+ def common_astra_logic(scope, environment, application_name,
87
+ metrics, capture_message_content, disable_metrics, version,
88
+ instance=None, endpoint=None):
35
89
  """
36
- Process DB operation and generate Telemetry
90
+ Process AstraDB request and generate telemetry.
91
+
92
+ Args:
93
+ scope: Scope object containing span, response, and operation details
94
+ environment: Deployment environment
95
+ application_name: Name of the application
96
+ metrics: Metrics dictionary for recording telemetry
97
+ capture_message_content: Flag to capture message content
98
+ disable_metrics: Flag to disable metrics collection
99
+ version: Version of the AstraDB client
100
+ instance: AstraDB client instance
101
+ endpoint: Operation endpoint for differentiation
37
102
  """
103
+ scope._end_time = time.time()
104
+
105
+ # Set common database span attributes using helper
106
+ common_db_span_attributes(scope, SemanticConvention.DB_SYSTEM_ASTRA,
107
+ scope._server_address, scope._server_port, environment, application_name, version)
108
+
109
+ # Set DB operation specific attributes
110
+ scope._span.set_attribute(SemanticConvention.DB_OPERATION_NAME, scope._db_operation)
111
+ scope._span.set_attribute(SemanticConvention.DB_CLIENT_OPERATION_DURATION,
112
+ scope._end_time - scope._start_time)
113
+
114
+ # Get collection name from instance
115
+ collection_name = getattr(instance, "name", "unknown")
116
+ scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
117
+
118
+ if scope._db_operation == SemanticConvention.DB_OPERATION_CREATE_COLLECTION:
119
+ # Handle create_collection operation
120
+ scope._span.set_attribute(SemanticConvention.DB_COLLECTION_DIMENSION,
121
+ scope._kwargs.get("dimension", -1))
122
+ scope._span.set_attribute(SemanticConvention.DB_INDEX_METRIC,
123
+ str(scope._kwargs.get("metric", "")))
124
+
125
+ # Set namespace if available in response
126
+ if scope._response and hasattr(scope._response, "keyspace"):
127
+ scope._span.set_attribute(SemanticConvention.DB_NAMESPACE, scope._response.keyspace)
128
+
129
+ if scope._response and hasattr(scope._response, "name"):
130
+ scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, scope._response.name)
131
+
132
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
133
+ f"{scope._db_operation} {collection_name} "
134
+ f"dimension={scope._kwargs.get('dimension', 'None')} "
135
+ f"metric={scope._kwargs.get('metric', 'None')}")
136
+
137
+ elif scope._db_operation == SemanticConvention.DB_OPERATION_DELETE_COLLECTION:
138
+ # Handle drop_collection operation
139
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
140
+ f"{scope._db_operation} {collection_name}")
141
+
142
+ elif scope._db_operation == SemanticConvention.DB_OPERATION_INSERT:
143
+ # Handle insert operations (insert_one, insert_many, regular insert)
144
+ documents = scope._args[0] if scope._args else scope._kwargs.get("documents", [])
145
+
146
+ scope._span.set_attribute(SemanticConvention.DB_DOCUMENTS_COUNT, object_count(documents))
147
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(documents))
148
+
149
+ # Response metrics
150
+ if scope._response and hasattr(scope._response, "inserted_ids"):
151
+ scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
152
+ len(scope._response.inserted_ids))
153
+
154
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
155
+ f"{scope._db_operation} {collection_name} "
156
+ f"documents_count={object_count(documents)}")
157
+
158
+ elif scope._db_operation == SemanticConvention.DB_OPERATION_UPDATE:
159
+ # Handle update operations (update_one, update_many, regular update)
160
+ update_query = scope._args[1] if len(scope._args) > 1 else scope._kwargs.get("update", {})
161
+ filter_query = scope._args[0] if scope._args else scope._kwargs.get("filter", {})
162
+
163
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(update_query))
164
+ scope._span.set_attribute(SemanticConvention.DB_FILTER, str(filter_query))
165
+
166
+ # Response metrics
167
+ if scope._response and hasattr(scope._response, "update_info"):
168
+ scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
169
+ scope._response.update_info.get("nModified", 0))
170
+
171
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
172
+ f"{scope._db_operation} {collection_name} "
173
+ f"filter={str(filter_query)[:100]}... "
174
+ f"update={str(update_query)[:100]}...")
175
+
176
+ elif scope._db_operation == SemanticConvention.DB_OPERATION_REPLACE:
177
+ # Handle replace operations (find_one_and_update, replace_one)
178
+ filter_query = scope._args[0] if scope._args else scope._kwargs.get("filter", {})
179
+
180
+ # Check if it's an upsert operation
181
+ if scope._kwargs.get("upsert"):
182
+ scope._db_operation = SemanticConvention.DB_OPERATION_UPSERT
183
+ scope._span.set_attribute(SemanticConvention.DB_OPERATION_NAME,
184
+ SemanticConvention.DB_OPERATION_UPSERT)
185
+
186
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(filter_query))
187
+ scope._span.set_attribute(SemanticConvention.DB_FILTER, str(filter_query))
188
+
189
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
190
+ f"{scope._db_operation} {collection_name} "
191
+ f"filter={str(filter_query)[:100]}... "
192
+ f"upsert={scope._kwargs.get('upsert', False)}")
193
+
194
+ elif scope._db_operation == SemanticConvention.DB_OPERATION_SELECT:
195
+ # Handle find operations
196
+ filter_query = scope._args[0] if scope._args else scope._kwargs.get("filter", {})
197
+
198
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(filter_query))
199
+ scope._span.set_attribute(SemanticConvention.DB_FILTER, str(filter_query))
200
+
201
+ # Response metrics
202
+ if scope._response and hasattr(scope._response, "__len__"):
203
+ scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
204
+ len(scope._response))
205
+
206
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
207
+ f"{scope._db_operation} {collection_name} "
208
+ f"filter={str(filter_query)[:100]}...")
209
+
210
+ elif scope._db_operation in [SemanticConvention.DB_OPERATION_DELETE,
211
+ SemanticConvention.DB_OPERATION_FIND_AND_DELETE]:
212
+ # Handle delete operations (delete_one, delete_many, find_one_and_delete)
213
+ filter_query = scope._args[0] if scope._args else scope._kwargs.get("filter", {})
214
+
215
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(filter_query))
216
+ scope._span.set_attribute(SemanticConvention.DB_FILTER, str(filter_query))
217
+
218
+ # Response metrics
219
+ if scope._response and hasattr(scope._response, "deleted_count"):
220
+ scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
221
+ scope._response.deleted_count)
222
+
223
+ scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
224
+ f"{scope._db_operation} {collection_name} "
225
+ f"filter={str(filter_query)[:100]}...")
226
+
227
+ scope._span.set_status(Status(StatusCode.OK))
228
+
229
+ # Record metrics using helper
230
+ if not disable_metrics:
231
+ record_db_metrics(metrics, SemanticConvention.DB_SYSTEM_ASTRA,
232
+ scope._server_address, scope._server_port, environment, application_name,
233
+ scope._start_time, scope._end_time)
234
+
235
+ def process_astra_response(response, db_operation, server_address, server_port,
236
+ environment, application_name, metrics, start_time, span,
237
+ capture_message_content, disable_metrics, version, instance, args, **kwargs):
238
+ """
239
+ Process AstraDB response and generate telemetry.
240
+
241
+ Args:
242
+ response: Response from AstraDB operation
243
+ db_operation: Database operation type
244
+ server_address: Server address
245
+ server_port: Server port
246
+ environment: Deployment environment
247
+ application_name: Application name
248
+ metrics: Metrics dictionary
249
+ start_time: Start time of the operation
250
+ span: OpenTelemetry span
251
+ capture_message_content: Flag to capture message content
252
+ disable_metrics: Flag to disable metrics
253
+ version: AstraDB client version
254
+ instance: AstraDB client instance
255
+ args: Positional arguments
256
+ **kwargs: Keyword arguments
257
+
258
+ Returns:
259
+ Original response
260
+ """
261
+
262
+ # Create a scope object to hold all the context
263
+ scope = type("GenericScope", (), {})()
264
+ scope._response = response
265
+ scope._db_operation = db_operation
266
+ scope._server_address = server_address
267
+ scope._server_port = server_port
268
+ scope._start_time = start_time
269
+ scope._span = span
270
+ scope._kwargs = kwargs
271
+ scope._args = args
272
+
273
+ # Process the response using common logic
274
+ common_astra_logic(scope, environment, application_name,
275
+ metrics, capture_message_content, disable_metrics, version, instance)
38
276
 
39
- end_time = time.time()
40
-
41
- try:
42
- span.set_attribute(TELEMETRY_SDK_NAME, 'openlit')
43
- span.set_attribute(SemanticConvention.GEN_AI_OPERATION, SemanticConvention.GEN_AI_OPERATION_TYPE_VECTORDB)
44
- span.set_attribute(SemanticConvention.DB_SYSTEM_NAME, SemanticConvention.DB_SYSTEM_ASTRA)
45
- span.set_attribute(SemanticConvention.DB_CLIENT_OPERATION_DURATION, end_time - start_time)
46
- span.set_attribute(SemanticConvention.SERVER_ADDRESS, server_address)
47
- span.set_attribute(SemanticConvention.SERVER_PORT, server_port)
48
- span.set_attribute(DEPLOYMENT_ENVIRONMENT, environment)
49
- span.set_attribute(SERVICE_NAME, application_name)
50
- span.set_attribute(SemanticConvention.DB_OPERATION_NAME, db_operation)
51
- span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
52
- span.set_attribute(SemanticConvention.DB_SDK_VERSION, version)
53
-
54
- if db_operation == SemanticConvention.DB_OPERATION_CREATE_COLLECTION:
55
- span.set_attribute(SemanticConvention.DB_NAMESPACE, response.keyspace)
56
- span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, response.name)
57
- span.set_attribute(SemanticConvention.DB_INDEX_DIMENSION, kwargs.get('dimension', ''))
58
- span.set_attribute(SemanticConvention.DB_INDEX_METRIC, str(kwargs.get('metric', '')))
59
-
60
- if db_operation == SemanticConvention.DB_OPERATION_INSERT:
61
- span.set_attribute(SemanticConvention.DB_DOCUMENTS_COUNT, object_count(args[0]))
62
- span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(args[0] or kwargs.get('documents', {})))
63
-
64
- elif db_operation == SemanticConvention.DB_OPERATION_UPDATE:
65
- span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS, response.update_info.get('nModified', 0))
66
- span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(args[1] or kwargs.get('update', {})))
67
-
68
- elif db_operation == SemanticConvention.DB_OPERATION_DELETE:
69
- span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS, response.deleted_count)
70
- span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(args[0] or kwargs.get('filter', {})))
71
-
72
- elif db_operation in [
73
- SemanticConvention.DB_OPERATION_SELECT,
74
- SemanticConvention.DB_OPERATION_FIND_AND_DELETE,
75
- SemanticConvention.DB_OPERATION_REPLACE
76
- ]:
77
- span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(args or kwargs.get('filter', {})))
78
-
79
- span.set_status(Status(StatusCode.OK))
80
-
81
- if not disable_metrics:
82
- attributes = {
83
- TELEMETRY_SDK_NAME: 'openlit',
84
- SERVICE_NAME: application_name,
85
- SemanticConvention.DB_SYSTEM_NAME: SemanticConvention.DB_SYSTEM_ASTRA,
86
- DEPLOYMENT_ENVIRONMENT: environment,
87
- SemanticConvention.GEN_AI_OPERATION: SemanticConvention.GEN_AI_OPERATION_TYPE_VECTORDB,
88
- SemanticConvention.DB_OPERATION_NAME: db_operation
89
- }
90
-
91
- metrics['db_requests'].add(1, attributes)
92
- metrics['db_client_operation_duration'].record(end_time - start_time, attributes)
93
-
94
- # Return original response
95
- return response
96
-
97
- except Exception as e:
98
- handle_exception(span, e)
99
- logger.error('Error in trace creation: %s', e)
100
-
101
- # Return original response
102
- return response
277
+ return response
@@ -32,10 +32,10 @@ def object_count(obj):
32
32
  def set_server_address_and_port(instance):
33
33
  """
34
34
  Extracts server address and port from ChromaDB client instance.
35
-
35
+
36
36
  Args:
37
37
  instance: ChromaDB client instance
38
-
38
+
39
39
  Returns:
40
40
  tuple: (server_address, server_port)
41
41
  """
@@ -19,7 +19,7 @@ def wrap_controlflow(gen_ai_endpoint, version, environment, application_name,
19
19
 
20
20
  This function wraps any given function to measure its execution time,
21
21
  log its operation, and trace its execution using OpenTelemetry.
22
-
22
+
23
23
  Parameters:
24
24
  - gen_ai_endpoint (str): A descriptor or name for the endpoint being traced.
25
25
  - version (str): The version of the Langchain application.
@@ -49,7 +49,7 @@ def wrap_controlflow(gen_ai_endpoint, version, environment, application_name,
49
49
 
50
50
  Returns:
51
51
  - The result of the wrapped function call.
52
-
52
+
53
53
  The wrapper initiates a span with the provided tracer, sets various attributes
54
54
  on the span based on the function's execution and response, and ensures
55
55
  errors are handled and logged appropriately.
@@ -19,7 +19,7 @@ def evaluate(gen_ai_endpoint, version, environment, application_name,
19
19
 
20
20
  This function wraps any given function to measure its execution time,
21
21
  log its operation, and trace its execution using OpenTelemetry.
22
-
22
+
23
23
  Parameters:
24
24
  - gen_ai_endpoint (str): A descriptor or name for the endpoint being traced.
25
25
  - version (str): The version of the EmbedChain application.
@@ -49,7 +49,7 @@ def evaluate(gen_ai_endpoint, version, environment, application_name,
49
49
 
50
50
  Returns:
51
51
  - The result of the wrapped function call.
52
-
52
+
53
53
  The wrapper initiates a span with the provided tracer, sets various attributes
54
54
  on the span based on the function's execution and response, and ensures
55
55
  errors are handled and logged appropriately.
@@ -97,7 +97,7 @@ def get_data_sources(gen_ai_endpoint, version, environment, application_name,
97
97
 
98
98
  This function wraps any given function to measure its execution time,
99
99
  log its operation, and trace its execution using OpenTelemetry.
100
-
100
+
101
101
  Parameters:
102
102
  - gen_ai_endpoint (str): A descriptor or name for the endpoint being traced.
103
103
  - version (str): The version of the EmbedChain application.
@@ -127,7 +127,7 @@ def get_data_sources(gen_ai_endpoint, version, environment, application_name,
127
127
 
128
128
  Returns:
129
129
  - The result of the wrapped function call.
130
-
130
+
131
131
  The wrapper initiates a span with the provided tracer, sets various attributes
132
132
  on the span based on the function's execution and response, and ensures
133
133
  errors are handled and logged appropriately.
@@ -30,16 +30,16 @@ class GroqInstrumentor(BaseInstrumentor):
30
30
 
31
31
  # Chat completions
32
32
  wrap_function_wrapper(
33
- "groq.resources.chat.completions",
34
- "Completions.create",
33
+ "groq.resources.chat.completions",
34
+ "Completions.create",
35
35
  chat(version, environment, application_name,
36
36
  tracer, pricing_info, capture_message_content, metrics, disable_metrics),
37
37
  )
38
38
 
39
39
  # Chat completions
40
40
  wrap_function_wrapper(
41
- "groq.resources.chat.completions",
42
- "AsyncCompletions.create",
41
+ "groq.resources.chat.completions",
42
+ "AsyncCompletions.create",
43
43
  async_chat(version, environment, application_name,
44
44
  tracer, pricing_info, capture_message_content, metrics, disable_metrics),
45
45
  )
@@ -1,49 +1,78 @@
1
- # pylint: disable=useless-return, bad-staticmethod-argument, disable=duplicate-code
2
- """Initializer of Auto Instrumentation of Haystack Functions"""
1
+ """
2
+ OpenLIT Haystack Instrumentation
3
+ """
4
+
3
5
  from typing import Collection
4
6
  import importlib.metadata
5
7
  from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
6
8
  from wrapt import wrap_function_wrapper
7
9
 
8
- from openlit.instrumentation.haystack.haystack import join_data
10
+ from openlit.instrumentation.haystack.haystack import general_wrap
11
+ from openlit.instrumentation.haystack.async_haystack import async_general_wrap
9
12
 
10
13
  _instruments = ("haystack-ai >= 2.0.0",)
11
14
 
12
- WRAPPED_METHODS = [
13
- {
14
- "package": "haystack.components.joiners.document_joiner",
15
- "object": "DocumentJoiner",
16
- "endpoint": "haystack.join_data",
17
- "wrapper": join_data,
18
- }
19
- ]
20
-
21
15
  class HaystackInstrumentor(BaseInstrumentor):
22
- """An instrumentor for Cohere's client library."""
16
+ """Optimized instrumentor for Haystack with minimal overhead"""
23
17
 
24
18
  def instrumentation_dependencies(self) -> Collection[str]:
25
19
  return _instruments
26
20
 
27
21
  def _instrument(self, **kwargs):
28
- application_name = kwargs.get("application_name")
29
- environment = kwargs.get("environment")
30
- tracer = kwargs.get("tracer")
31
- pricing_info = kwargs.get("pricing_info")
32
- capture_message_content = kwargs.get("capture_message_content")
33
22
  version = importlib.metadata.version("haystack-ai")
23
+ environment = kwargs.get("environment", "default")
24
+ application_name = kwargs.get("application_name", "default")
25
+ tracer = kwargs.get("tracer")
26
+ pricing_info = kwargs.get("pricing_info", {})
27
+ capture_message_content = kwargs.get("capture_message_content", False)
28
+ metrics = kwargs.get("metrics_dict")
29
+ disable_metrics = kwargs.get("disable_metrics")
30
+ detailed_tracing = kwargs.get("detailed_tracing", False)
34
31
 
35
- for wrapped_method in WRAPPED_METHODS:
36
- wrap_package = wrapped_method.get("package")
37
- wrap_object = wrapped_method.get("object")
38
- gen_ai_endpoint = wrapped_method.get("endpoint")
39
- wrapper = wrapped_method.get("wrapper")
32
+ # Pipeline operations (always enabled)
33
+ try:
40
34
  wrap_function_wrapper(
41
- wrap_package,
42
- wrap_object,
43
- wrapper(gen_ai_endpoint, version, environment, application_name,
44
- tracer, pricing_info, capture_message_content),
35
+ "haystack", "Pipeline.run",
36
+ general_wrap("pipeline", version, environment, application_name,
37
+ tracer, pricing_info, capture_message_content,
38
+ metrics, disable_metrics)
45
39
  )
40
+ wrap_function_wrapper(
41
+ "haystack", "AsyncPipeline.run_async",
42
+ async_general_wrap("pipeline", version, environment,
43
+ application_name, tracer, pricing_info,
44
+ capture_message_content, metrics, disable_metrics)
45
+ )
46
+ except Exception:
47
+ pass
48
+
49
+ # Component operations (only if detailed_tracing enabled)
50
+ if detailed_tracing:
51
+ components = [
52
+ ("haystack.components.retrievers.in_memory",
53
+ "InMemoryBM25Retriever.run", "bm25_retriever"),
54
+ ("haystack.components.builders.prompt_builder",
55
+ "PromptBuilder.run", "prompt_builder"),
56
+ ("haystack.components.generators.openai",
57
+ "OpenAIGenerator.run", "openai_generator"),
58
+ ("haystack.components.generators.chat.openai",
59
+ "OpenAIChatGenerator.run", "openai_chat_generator"),
60
+ ("haystack.components.embedders.openai_text_embedder",
61
+ "OpenAITextEmbedder.run", "text_embedder"),
62
+ ("haystack.components.embedders.openai_document_embedder",
63
+ "OpenAIDocumentEmbedder.run", "document_embedder"),
64
+ ]
65
+
66
+ for module, method, component_type in components:
67
+ try:
68
+ wrap_function_wrapper(
69
+ module, method,
70
+ general_wrap(component_type, version, environment,
71
+ application_name, tracer, pricing_info,
72
+ capture_message_content, metrics, disable_metrics)
73
+ )
74
+ except Exception:
75
+ pass
46
76
 
47
- @staticmethod
48
77
  def _uninstrument(self, **kwargs):
49
78
  pass
@@ -0,0 +1,54 @@
1
+ """
2
+ Haystack async wrapper
3
+ """
4
+
5
+ import time
6
+ from opentelemetry.trace import SpanKind
7
+ from opentelemetry import context as context_api
8
+ from openlit.__helpers import handle_exception
9
+ from openlit.instrumentation.haystack.utils import (
10
+ process_haystack_response,
11
+ OPERATION_MAP,
12
+ set_server_address_and_port,
13
+ )
14
+
15
+ def async_general_wrap(endpoint, version, environment, application_name,
16
+ tracer, pricing_info, capture_message_content, metrics, disable_metrics):
17
+ """Optimized async wrapper for Haystack operations"""
18
+
19
+ async def wrapper(wrapped, instance, args, kwargs):
20
+ """Fast async wrapper with minimal overhead"""
21
+
22
+ # CRITICAL: Suppression check
23
+ if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
24
+ return await wrapped(*args, **kwargs)
25
+
26
+ # Fast operation mapping
27
+ operation_type = OPERATION_MAP.get(endpoint, "framework")
28
+
29
+ # Optimized span naming
30
+ if endpoint == "pipeline":
31
+ span_name = f"{operation_type} pipeline"
32
+ else:
33
+ span_name = f"{operation_type} {endpoint}"
34
+
35
+ # Fast server address
36
+ server_address, server_port = set_server_address_and_port(instance)
37
+
38
+ with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
39
+ start_time = time.time()
40
+ response = await wrapped(*args, **kwargs)
41
+
42
+ try:
43
+ response = process_haystack_response(
44
+ response, operation_type, server_address, server_port,
45
+ environment, application_name, metrics, start_time, span,
46
+ capture_message_content, disable_metrics, version,
47
+ instance, args, endpoint=endpoint, **kwargs
48
+ )
49
+ except Exception as e:
50
+ handle_exception(span, e)
51
+
52
+ return response
53
+
54
+ return wrapper