openlit 1.34.30__py3-none-any.whl → 1.34.31__py3-none-any.whl

openlit/instrumentation/chroma/utils.py

@@ -1,6 +1,7 @@
 """
 ChromaDB OpenTelemetry instrumentation utility functions
 """
+
 import time
 from urllib.parse import urlparse
 from opentelemetry.trace import Status, StatusCode
@@ -23,12 +24,14 @@ DB_OPERATION_MAP = {
     "chroma.delete": SemanticConvention.DB_OPERATION_DELETE,
 }
 
+
 def object_count(obj):
     """
     Counts length of object if it exists, else returns 0.
     """
     return len(obj) if obj else 0
 
+
 def set_server_address_and_port(instance):
     """
     Extracts server address and port from ChromaDB client instance.
@@ -74,8 +77,18 @@ def set_server_address_and_port(instance):
 
     return server_address, server_port
 
-def common_vectordb_logic(scope, environment, application_name,
-    metrics, capture_message_content, disable_metrics, version, instance=None, endpoint=None):
+
+def common_vectordb_logic(
+    scope,
+    environment,
+    application_name,
+    metrics,
+    capture_message_content,
+    disable_metrics,
+    version,
+    instance=None,
+    endpoint=None,
+):
     """
     Process vector database request and generate telemetry.
     """
@@ -83,12 +96,22 @@ def common_vectordb_logic(scope, environment, application_name,
     scope._end_time = time.time()
 
     # Set common database span attributes using helper
-    common_db_span_attributes(scope, SemanticConvention.DB_SYSTEM_CHROMA, scope._server_address, scope._server_port,
-        environment, application_name, version)
+    common_db_span_attributes(
+        scope,
+        SemanticConvention.DB_SYSTEM_CHROMA,
+        scope._server_address,
+        scope._server_port,
+        environment,
+        application_name,
+        version,
+    )
 
     # Set DB operation specific attributes
     scope._span.set_attribute(SemanticConvention.DB_OPERATION_NAME, scope._db_operation)
-    scope._span.set_attribute(SemanticConvention.DB_CLIENT_OPERATION_DURATION, scope._end_time - scope._start_time)
+    scope._span.set_attribute(
+        SemanticConvention.DB_CLIENT_OPERATION_DURATION,
+        scope._end_time - scope._start_time,
+    )
 
     # Set collection name from instance
     if hasattr(instance, "name"):
@@ -97,13 +120,19 @@ def common_vectordb_logic(scope, environment, application_name,
     # Set Create Collection operation specific attributes
     if scope._db_operation == SemanticConvention.DB_OPERATION_CREATE_COLLECTION:
         # Standard database attributes
-        collection_name = scope._kwargs.get("name") or (scope._args[0] if scope._args else "unknown")
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
+        collection_name = scope._kwargs.get("name") or (
+            scope._args[0] if scope._args else "unknown"
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
 
         # Vector database specific attributes (extensions)
         metadata = scope._kwargs.get("metadata", {})
         if metadata:
-            scope._span.set_attribute(SemanticConvention.DB_COLLECTION_METADATA, str(metadata))
+            scope._span.set_attribute(
+                SemanticConvention.DB_COLLECTION_METADATA, str(metadata)
+            )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_INSERT:
         collection_name = getattr(instance, "name", "unknown")
@@ -111,14 +140,20 @@ def common_vectordb_logic(scope, environment, application_name,
 
         # Standard database attributes
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(query))
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(query))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_COUNT, object_count(query)
+        )
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {collection_name} "
             f"ids={query} "
-            f"documents={scope._kwargs.get('documents', [])}")
+            f"documents={scope._kwargs.get('documents', [])}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_GET:
         collection_name = getattr(instance, "name", "unknown")
@@ -130,17 +165,25 @@ def common_vectordb_logic(scope, environment, application_name,
 
             # Standard database attributes
             scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(query))
-            scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-            scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(query))
+            scope._span.set_attribute(
+                SemanticConvention.DB_COLLECTION_NAME, collection_name
+            )
+            scope._span.set_attribute(
+                SemanticConvention.DB_VECTOR_COUNT, object_count(query)
+            )
 
             # Vector database specific attributes (extensions)
-            scope._span.set_attribute(SemanticConvention.DB_FILTER, str(scope._kwargs.get("where", "")))
+            scope._span.set_attribute(
+                SemanticConvention.DB_FILTER, str(scope._kwargs.get("where", ""))
+            )
 
-            scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+            scope._span.set_attribute(
+                SemanticConvention.DB_QUERY_SUMMARY,
                 f"{scope._db_operation} {collection_name} "
                 f"ids={query} "
                 f"limit={scope._kwargs.get('limit', 'None')} "
-                f"offset={scope._kwargs.get('offset', 'None')}")
+                f"offset={scope._kwargs.get('offset', 'None')}",
+            )
 
         elif endpoint == "chroma.query":
             query_texts = scope._kwargs.get("query_texts", [])
@@ -156,24 +199,39 @@ def common_vectordb_logic(scope, environment, application_name,
 
             # Standard database attributes
             scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, query_content)
-            scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
+            scope._span.set_attribute(
+                SemanticConvention.DB_COLLECTION_NAME, collection_name
+            )
 
             # Vector database specific attributes (extensions)
-            scope._span.set_attribute(SemanticConvention.DB_VECTOR_QUERY_TOP_K, scope._kwargs.get("n_results", 10))
-            scope._span.set_attribute(SemanticConvention.DB_FILTER, str(scope._kwargs.get("where", "")))
+            scope._span.set_attribute(
+                SemanticConvention.DB_VECTOR_QUERY_TOP_K,
+                scope._kwargs.get("n_results", 10),
+            )
+            scope._span.set_attribute(
+                SemanticConvention.DB_FILTER, str(scope._kwargs.get("where", ""))
+            )
 
             # Extract response metrics if available
             if scope._response:
                 # Get number of results returned
                 if hasattr(scope._response, "get") and scope._response.get("ids"):
-                    returned_rows = object_count(scope._response["ids"][0]) if scope._response["ids"] else 0
-                    scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS, returned_rows)
-
-            scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+                    returned_rows = (
+                        object_count(scope._response["ids"][0])
+                        if scope._response["ids"]
+                        else 0
+                    )
+                    scope._span.set_attribute(
+                        SemanticConvention.DB_RESPONSE_RETURNED_ROWS, returned_rows
+                    )
+
+            scope._span.set_attribute(
+                SemanticConvention.DB_QUERY_SUMMARY,
                 f"{scope._db_operation} {collection_name} "
                 f"n_results={scope._kwargs.get('n_results', 10)} "
                 f"{query_content} "
-                f"filter={scope._kwargs.get('where', 'None')}")
+                f"filter={scope._kwargs.get('where', 'None')}",
+            )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_UPDATE:
         collection_name = getattr(instance, "name", "unknown")
@@ -181,16 +239,22 @@ def common_vectordb_logic(scope, environment, application_name,
 
         # Standard database attributes
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(query))
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(query))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_COUNT, object_count(query)
+        )
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {collection_name} "
             f"ids={query} "
             f"embeddings={scope._kwargs.get('embeddings', 'None')} "
             f"metadatas={scope._kwargs.get('metadatas', 'None')} "
-            f"documents={scope._kwargs.get('documents', 'None')}")
+            f"documents={scope._kwargs.get('documents', 'None')}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_UPSERT:
         collection_name = getattr(instance, "name", "unknown")
@@ -198,16 +262,22 @@ def common_vectordb_logic(scope, environment, application_name,
 
         # Standard database attributes
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(query))
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(query))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_COUNT, object_count(query)
+        )
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {collection_name} "
             f"ids={query} "
             f"embeddings={scope._kwargs.get('embeddings', 'None')} "
             f"metadatas={scope._kwargs.get('metadatas', 'None')} "
-            f"documents={scope._kwargs.get('documents', 'None')}")
+            f"documents={scope._kwargs.get('documents', 'None')}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_DELETE:
         collection_name = getattr(instance, "name", "unknown")
@@ -215,16 +285,24 @@ def common_vectordb_logic(scope, environment, application_name,
 
         # Standard database attributes
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(query))
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(query))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_COUNT, object_count(query)
+        )
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_FILTER, str(scope._kwargs.get("where", "")))
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_FILTER, str(scope._kwargs.get("where", ""))
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {collection_name} "
             f"ids={query} "
             f"filter={scope._kwargs.get('where', 'None')} "
-            f"delete_all={scope._kwargs.get('delete_all', False)}")
+            f"delete_all={scope._kwargs.get('delete_all', False)}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_PEEK:
         collection_name = getattr(instance, "name", "unknown")
@@ -232,25 +310,54 @@ def common_vectordb_logic(scope, environment, application_name,
 
         # Standard database attributes
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, query)
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, scope._kwargs.get("limit", ""))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_COUNT, scope._kwargs.get("limit", "")
+        )
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {collection_name} "
-            f"limit={scope._kwargs.get('limit', 'None')}")
+            f"limit={scope._kwargs.get('limit', 'None')}",
+        )
 
     scope._span.set_status(Status(StatusCode.OK))
 
     # Record metrics using helper
     if not disable_metrics:
-        record_db_metrics(metrics, SemanticConvention.DB_SYSTEM_CHROMA, scope._server_address, scope._server_port,
-            environment, application_name, scope._start_time, scope._end_time)
-
-def process_vectordb_response(response, db_operation, server_address, server_port,
-    environment, application_name, metrics, start_time, span,
-    capture_message_content=False, disable_metrics=False,
-    version="1.0.0", instance=None, args=None, endpoint=None, **kwargs):
+        record_db_metrics(
+            metrics,
+            SemanticConvention.DB_SYSTEM_CHROMA,
+            scope._server_address,
+            scope._server_port,
+            environment,
+            application_name,
+            scope._start_time,
+            scope._end_time,
+        )
+
+
+def process_vectordb_response(
+    response,
+    db_operation,
+    server_address,
+    server_port,
+    environment,
+    application_name,
+    metrics,
+    start_time,
+    span,
+    capture_message_content=False,
+    disable_metrics=False,
+    version="1.0.0",
+    instance=None,
+    args=None,
+    endpoint=None,
+    **kwargs,
+):
     """
     Process vector database response and generate telemetry following OpenTelemetry conventions.
     """
@@ -266,7 +373,16 @@ def process_vectordb_response(response, db_operation, server_address, server_por
     scope._server_address = server_address
     scope._server_port = server_port
 
-    common_vectordb_logic(scope, environment, application_name,
-        metrics, capture_message_content, disable_metrics, version, instance, endpoint)
+    common_vectordb_logic(
+        scope,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        instance,
+        endpoint,
+    )
 
     return response

openlit/instrumentation/cohere/__init__.py

@@ -5,19 +5,16 @@ import importlib.metadata
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from wrapt import wrap_function_wrapper
 
-from openlit.instrumentation.cohere.cohere import (
-    chat,
-    chat_stream,
-    embed
-)
+from openlit.instrumentation.cohere.cohere import chat, chat_stream, embed
 from openlit.instrumentation.cohere.async_cohere import (
     async_chat,
     async_chat_stream,
-    async_embed
+    async_embed,
 )
 
 _instruments = ("cohere >= 5.14.0",)
 
+
 class CohereInstrumentor(BaseInstrumentor):
     """
     An instrumentor for Cohere client library.
@@ -40,48 +37,96 @@ class CohereInstrumentor(BaseInstrumentor):
         wrap_function_wrapper(
             "cohere.client_v2",
             "ClientV2.chat",
-            chat(version, environment, application_name,
-                  tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            chat(
+                version,
+                environment,
+                application_name,
+                tracer,
+                pricing_info,
+                capture_message_content,
+                metrics,
+                disable_metrics,
+            ),
         )
 
         # sync chat streaming
         wrap_function_wrapper(
             "cohere.client_v2",
             "ClientV2.chat_stream",
-            chat_stream(version, environment, application_name,
-                  tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            chat_stream(
+                version,
+                environment,
+                application_name,
+                tracer,
+                pricing_info,
+                capture_message_content,
+                metrics,
+                disable_metrics,
+            ),
         )
 
         # sync embeddings
         wrap_function_wrapper(
             "cohere.client_v2",
             "ClientV2.embed",
-            embed(version, environment, application_name,
-                  tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            embed(
+                version,
+                environment,
+                application_name,
+                tracer,
+                pricing_info,
+                capture_message_content,
+                metrics,
+                disable_metrics,
+            ),
         )
 
         # async chat completions
         wrap_function_wrapper(
             "cohere.client_v2",
             "AsyncClientV2.chat",
-            async_chat(version, environment, application_name,
-                  tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            async_chat(
+                version,
+                environment,
+                application_name,
+                tracer,
+                pricing_info,
+                capture_message_content,
+                metrics,
+                disable_metrics,
+            ),
         )
 
         # async chat streaming
         wrap_function_wrapper(
             "cohere.client_v2",
             "AsyncClientV2.chat_stream",
-            async_chat_stream(version, environment, application_name,
-                  tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            async_chat_stream(
+                version,
+                environment,
+                application_name,
+                tracer,
+                pricing_info,
+                capture_message_content,
+                metrics,
+                disable_metrics,
+            ),
         )
 
         # async embeddings
         wrap_function_wrapper(
             "cohere.client_v2",
             "AsyncClientV2.embed",
-            async_embed(version, environment, application_name,
-                  tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            async_embed(
+                version,
+                environment,
+                application_name,
+                tracer,
+                pricing_info,
+                capture_message_content,
+                metrics,
+                disable_metrics,
+            ),
         )
 
     def _uninstrument(self, **kwargs):

openlit/instrumentation/cohere/async_cohere.py

@@ -16,8 +16,17 @@ from openlit.instrumentation.cohere.utils import (
 )
 from openlit.semcov import SemanticConvention
 
-def async_chat(version, environment, application_name,
-              tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+
+def async_chat(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for GenAI chat function call
     """
@@ -27,7 +36,9 @@ def async_chat(version, environment, application_name,
         Wraps the GenAI chat function call.
         """
 
-        server_address, server_port = set_server_address_and_port(instance, "api.cohere.com", 443)
+        server_address, server_port = set_server_address_and_port(
+            instance, "api.cohere.com", 443
+        )
         request_model = kwargs.get("model", "command-r-plus-08-2024")
 
         span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
@@ -49,15 +60,24 @@ def async_chat(version, environment, application_name,
                 capture_message_content=capture_message_content,
                 disable_metrics=disable_metrics,
                 version=version,
-                **kwargs
+                **kwargs,
             )
 
         return response
 
     return wrapper
 
-def async_chat_stream(version, environment, application_name,
-                     tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+
+def async_chat_stream(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for GenAI chat_stream function call
     """
@@ -68,15 +88,15 @@ def async_chat_stream(version, environment, application_name,
         """
 
         def __init__(
-                self,
-                wrapped,
-                span,
-                span_name,
-                kwargs,
-                server_address,
-                server_port,
-                **args,
-            ):
+            self,
+            wrapped,
+            span,
+            span_name,
+            kwargs,
+            server_address,
+            server_port,
+            **args,
+        ):
             self.__wrapped__ = wrapped
             self._span = span
             self._span_name = span_name
@@ -120,7 +140,9 @@ def async_chat_stream(version, environment, application_name,
                 return chunk
             except StopAsyncIteration:
                 try:
-                    with tracer.start_as_current_span(self._span_name, kind= SpanKind.CLIENT) as self._span:
+                    with tracer.start_as_current_span(
+                        self._span_name, kind=SpanKind.CLIENT
+                    ) as self._span:
                         process_streaming_chat_response(
                             self,
                             pricing_info=pricing_info,
@@ -129,7 +151,7 @@ def async_chat_stream(version, environment, application_name,
                             metrics=metrics,
                             capture_message_content=capture_message_content,
                             disable_metrics=disable_metrics,
-                            version=version
+                            version=version,
                         )
 
                 except Exception as e:
@@ -142,7 +164,9 @@ def async_chat_stream(version, environment, application_name,
         Wraps the GenAI chat_stream function call.
         """
 
-        server_address, server_port = set_server_address_and_port(instance, "api.cohere.com", 443)
+        server_address, server_port = set_server_address_and_port(
+            instance, "api.cohere.com", 443
+        )
         request_model = kwargs.get("model", "command-r-plus-08-2024")
 
         span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT} {request_model}"
@@ -151,12 +175,23 @@ def async_chat_stream(version, environment, application_name,
         awaited_wrapped = await wrapped(*args, **kwargs)
         span = tracer.start_span(span_name, kind=SpanKind.CLIENT)
 
-        return TracedAsyncStream(awaited_wrapped, span, span_name, kwargs, server_address, server_port)
+        return TracedAsyncStream(
+            awaited_wrapped, span, span_name, kwargs, server_address, server_port
+        )
 
     return wrapper
 
-def async_embed(version, environment, application_name,
-               tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+
+def async_embed(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for GenAI embedding function call
     """
@@ -166,10 +201,14 @@ def async_embed(version, environment, application_name,
         Wraps the GenAI embedding function call.
         """
 
-        server_address, server_port = set_server_address_and_port(instance, "api.cohere.com", 443)
+        server_address, server_port = set_server_address_and_port(
+            instance, "api.cohere.com", 443
+        )
         request_model = kwargs.get("model", "embed-english-v3.0")
 
-        span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_EMBEDDING} {request_model}"
+        span_name = (
+            f"{SemanticConvention.GEN_AI_OPERATION_TYPE_EMBEDDING} {request_model}"
+        )
 
         with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
             start_time = time.time()
@@ -190,7 +229,7 @@ def async_embed(version, environment, application_name,
                     capture_message_content=capture_message_content,
                     disable_metrics=disable_metrics,
                     version=version,
-                    **kwargs
+                    **kwargs,
                 )
 
             except Exception as e: