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

openlit/instrumentation/pinecone/__init__.py

@@ -13,7 +13,11 @@ _instruments = ("pinecone >= 7.3.0",)
 # Pinecone sync operations
 PINECONE_SYNC_OPERATIONS = [
     ("pinecone.pinecone", "Pinecone.create_index", "pinecone.create_collection"),
-    ("pinecone.pinecone", "Pinecone.create_index_for_model", "pinecone.create_collection"),
+    (
+        "pinecone.pinecone",
+        "Pinecone.create_index_for_model",
+        "pinecone.create_collection",
+    ),
     ("pinecone.db_data.index", "Index.upsert", "pinecone.upsert"),
     ("pinecone.db_data.index", "Index.upsert_records", "pinecone.upsert_records"),
     ("pinecone.db_data.index", "Index.query", "pinecone.query"),
@@ -26,18 +30,35 @@ PINECONE_SYNC_OPERATIONS = [
 
 # Pinecone async operations
 PINECONE_ASYNC_OPERATIONS = [
-    ("pinecone.pinecone_asyncio", "PineconeAsyncio.create_index", "pinecone.create_index"),
-    ("pinecone.pinecone_asyncio", "PineconeAsyncio.create_index_for_model", "pinecone.create_index"),
+    (
+        "pinecone.pinecone_asyncio",
+        "PineconeAsyncio.create_index",
+        "pinecone.create_index",
+    ),
+    (
+        "pinecone.pinecone_asyncio",
+        "PineconeAsyncio.create_index_for_model",
+        "pinecone.create_index",
+    ),
     ("pinecone.db_data.index_asyncio", "_IndexAsyncio.upsert", "pinecone.upsert"),
-    ("pinecone.db_data.index_asyncio", "_IndexAsyncio.upsert_records", "pinecone.upsert_records"),
+    (
+        "pinecone.db_data.index_asyncio",
+        "_IndexAsyncio.upsert_records",
+        "pinecone.upsert_records",
+    ),
     ("pinecone.db_data.index_asyncio", "_IndexAsyncio.query", "pinecone.query"),
     ("pinecone.db_data.index_asyncio", "_IndexAsyncio.search", "pinecone.search"),
     ("pinecone.db_data.index_asyncio", "_IndexAsyncio.fetch", "pinecone.fetch"),
-    ("pinecone.db_data.index_asyncio", "_IndexAsyncio.search_records", "pinecone.search_records"),
+    (
+        "pinecone.db_data.index_asyncio",
+        "_IndexAsyncio.search_records",
+        "pinecone.search_records",
+    ),
     ("pinecone.db_data.index_asyncio", "_IndexAsyncio.update", "pinecone.update"),
     ("pinecone.db_data.index_asyncio", "_IndexAsyncio.delete", "pinecone.delete"),
 ]
 
+
 class PineconeInstrumentor(BaseInstrumentor):
     """
     An instrumentor for Pinecone client library.
@@ -61,8 +82,17 @@ class PineconeInstrumentor(BaseInstrumentor):
             wrap_function_wrapper(
                 module,
                 class_method,
-                general_wrap(endpoint, version, environment, application_name,
-                    tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+                general_wrap(
+                    endpoint,
+                    version,
+                    environment,
+                    application_name,
+                    tracer,
+                    pricing_info,
+                    capture_message_content,
+                    metrics,
+                    disable_metrics,
+                ),
             )
 
         # Wrap async operations
@@ -70,8 +100,17 @@ class PineconeInstrumentor(BaseInstrumentor):
             wrap_function_wrapper(
                 module,
                 class_method,
-                async_general_wrap(endpoint, version, environment, application_name,
-                    tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+                async_general_wrap(
+                    endpoint,
+                    version,
+                    environment,
+                    application_name,
+                    tracer,
+                    pricing_info,
+                    capture_message_content,
+                    metrics,
+                    disable_metrics,
+                ),
             )
 
     def _uninstrument(self, **kwargs):

openlit/instrumentation/pinecone/async_pinecone.py

@@ -12,8 +12,18 @@ from openlit.instrumentation.pinecone.utils import (
     set_server_address_and_port,
 )
 
-def async_general_wrap(gen_ai_endpoint, version, environment, application_name,
-    tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+
+def async_general_wrap(
+    gen_ai_endpoint,
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for Pinecone async function calls.
     """
@@ -43,9 +53,21 @@ def async_general_wrap(gen_ai_endpoint, version, environment, application_name,
             try:
                 # Process response and generate telemetry
                 response = process_vectordb_response(
-                    response, db_operation, server_address, server_port,
-                    environment, application_name, metrics, start_time, span,
-                    capture_message_content, disable_metrics, version, instance, args, **kwargs
+                    response,
+                    db_operation,
+                    server_address,
+                    server_port,
+                    environment,
+                    application_name,
+                    metrics,
+                    start_time,
+                    span,
+                    capture_message_content,
+                    disable_metrics,
+                    version,
+                    instance,
+                    args,
+                    **kwargs,
                 )
 
             except Exception as e:

openlit/instrumentation/pinecone/pinecone.py

@@ -12,8 +12,18 @@ from openlit.instrumentation.pinecone.utils import (
     set_server_address_and_port,
 )
 
-def general_wrap(gen_ai_endpoint, version, environment, application_name,
-    tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+
+def general_wrap(
+    gen_ai_endpoint,
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for Pinecone function calls.
     """
@@ -43,9 +53,21 @@ def general_wrap(gen_ai_endpoint, version, environment, application_name,
             try:
                 # Process response and generate telemetry
                 response = process_vectordb_response(
-                    response, db_operation, server_address, server_port,
-                    environment, application_name, metrics, start_time, span,
-                    capture_message_content, disable_metrics, version, instance, args, **kwargs
+                    response,
+                    db_operation,
+                    server_address,
+                    server_port,
+                    environment,
+                    application_name,
+                    metrics,
+                    start_time,
+                    span,
+                    capture_message_content,
+                    disable_metrics,
+                    version,
+                    instance,
+                    args,
+                    **kwargs,
                 )
 
             except Exception as e:

openlit/instrumentation/pinecone/utils.py

@@ -1,6 +1,7 @@
 """
 Pinecone OpenTelemetry instrumentation utility functions
 """
+
 import time
 from urllib.parse import urlparse
 from opentelemetry.trace import Status, StatusCode
@@ -24,12 +25,14 @@ DB_OPERATION_MAP = {
     "pinecone.search_records": SemanticConvention.DB_OPERATION_QUERY,
 }
 
+
 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 Pinecone client instance.
@@ -70,8 +73,17 @@ 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):
+
+def common_vectordb_logic(
+    scope,
+    environment,
+    application_name,
+    metrics,
+    capture_message_content,
+    disable_metrics,
+    version,
+    instance=None,
+):
     """
     Process vector database request and generate telemetry.
     """
@@ -79,25 +91,47 @@ 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_PINECONE, scope._server_address, scope._server_port,
-        environment, application_name, version)
+    common_db_span_attributes(
+        scope,
+        SemanticConvention.DB_SYSTEM_PINECONE,
+        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 Create Index operation specific attributes
     if scope._db_operation == SemanticConvention.DB_OPERATION_CREATE_COLLECTION:
         # Standard database attributes
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, scope._kwargs.get("name", ""))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, scope._kwargs.get("name", "")
+        )
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_DIMENSION, scope._kwargs.get("dimension", -1))
-        scope._span.set_attribute(SemanticConvention.DB_SEARCH_SIMILARITY_METRIC, scope._kwargs.get("metric", "cosine"))
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_SPEC, str(scope._kwargs.get("spec", "")))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_DIMENSION,
+            scope._kwargs.get("dimension", -1),
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_SEARCH_SIMILARITY_METRIC,
+            scope._kwargs.get("metric", "cosine"),
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_SPEC, str(scope._kwargs.get("spec", ""))
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_SEARCH:
-        namespace = scope._kwargs.get("namespace", "default") or (scope._args[0] if scope._args else "unknown")
+        namespace = scope._kwargs.get("namespace", "default") or (
+            scope._args[0] if scope._args else "unknown"
+        )
         query = scope._kwargs.get("query", {})
 
         # Extract query text or vector from different possible locations
@@ -110,15 +144,21 @@ def common_vectordb_logic(scope, environment, application_name,
         scope._span.set_attribute(SemanticConvention.DB_NAMESPACE, namespace)
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_QUERY_TOP_K, query.get("top_k", -1))
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_QUERY_TOP_K, query.get("top_k", -1)
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {namespace} "
             f"top_k={query.get('top_k', -1)} "
             f"text={query_text} "
-            f"vector={query_vector}")
+            f"vector={query_vector}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_QUERY:
-        namespace = scope._kwargs.get("namespace", "default") or (scope._args[0] if scope._args else "unknown")
+        namespace = scope._kwargs.get("namespace", "default") or (
+            scope._args[0] if scope._args else "unknown"
+        )
         query = scope._kwargs.get("vector", [])
 
         # Standard database attributes
@@ -126,16 +166,24 @@ def common_vectordb_logic(scope, environment, application_name,
         scope._span.set_attribute(SemanticConvention.DB_NAMESPACE, namespace)
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_QUERY_TOP_K, scope._kwargs.get("top_k", ""))
-        scope._span.set_attribute(SemanticConvention.DB_FILTER, str(scope._kwargs.get("filter", "")))
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_QUERY_TOP_K, scope._kwargs.get("top_k", "")
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_FILTER, str(scope._kwargs.get("filter", ""))
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {namespace} "
             f"top_k={scope._kwargs.get('top_k', -1)} "
             f"filtered={scope._kwargs.get('filter', '')} "
-            f"vector={scope._kwargs.get('vector', '')}")
+            f"vector={scope._kwargs.get('vector', '')}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_FETCH:
-        namespace = scope._kwargs.get("namespace", "default") or (scope._args[0] if scope._args else "unknown")
+        namespace = scope._kwargs.get("namespace", "default") or (
+            scope._args[0] if scope._args else "unknown"
+        )
         query = scope._kwargs.get("ids", [])
 
         # Standard database attributes
@@ -143,13 +191,19 @@ def common_vectordb_logic(scope, environment, application_name,
         scope._span.set_attribute(SemanticConvention.DB_NAMESPACE, namespace)
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
-            f"{scope._db_operation} {namespace} "
-            f"ids={query}")
-        scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS, object_count(scope._response.vectors))
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
+            f"{scope._db_operation} {namespace} ids={query}",
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
+            object_count(scope._response.vectors),
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_UPDATE:
-        namespace = scope._kwargs.get("namespace") or (scope._args[0] if scope._args else "unknown")
+        namespace = scope._kwargs.get("namespace") or (
+            scope._args[0] if scope._args else "unknown"
+        )
         query = scope._kwargs.get("id", "")
 
         # Standard database attributes
@@ -157,54 +211,98 @@ def common_vectordb_logic(scope, environment, application_name,
         scope._span.set_attribute(SemanticConvention.DB_NAMESPACE, namespace)
 
         # 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} {namespace} "
             f"id={query} "
             f"values={scope._kwargs.get('values', [])} "
-            f"set_metadata={scope._kwargs.get('set_metadata', '')}")
+            f"set_metadata={scope._kwargs.get('set_metadata', '')}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_UPSERT:
-        namespace = scope._kwargs.get("namespace") or (scope._args[0] if scope._args else "unknown")
-        query = scope._kwargs.get("vectors") or (scope._args[1] if len(scope._args) > 1 else None)
+        namespace = scope._kwargs.get("namespace") or (
+            scope._args[0] if scope._args else "unknown"
+        )
+        query = scope._kwargs.get("vectors") or (
+            scope._args[1] if len(scope._args) > 1 else None
+        )
 
         # Standard database attributes
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(query))
         scope._span.set_attribute(SemanticConvention.DB_NAMESPACE, namespace)
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(query))
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
-            f"{scope._db_operation} {namespace} vectors_count={object_count(query)}")
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_COUNT, object_count(query)
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
+            f"{scope._db_operation} {namespace} vectors_count={object_count(query)}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_DELETE:
-        namespace = scope._kwargs.get("namespace") or (scope._args[0] if scope._args else "unknown")
-        query = scope._kwargs.get("ids") or (scope._args[1] if len(scope._args) > 1 else None)
+        namespace = scope._kwargs.get("namespace") or (
+            scope._args[0] if scope._args else "unknown"
+        )
+        query = scope._kwargs.get("ids") or (
+            scope._args[1] if len(scope._args) > 1 else None
+        )
 
         # Standard database attributes
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(query))
         scope._span.set_attribute(SemanticConvention.DB_NAMESPACE, namespace)
 
         # Vector database specific attributes (extensions)
-        scope._span.set_attribute(SemanticConvention.DB_ID_COUNT, object_count(scope._kwargs.get("ids")))
-        scope._span.set_attribute(SemanticConvention.DB_FILTER, str(scope._kwargs.get("filter", "")))
-        scope._span.set_attribute(SemanticConvention.DB_DELETE_ALL, scope._kwargs.get("delete_all", False))
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_ID_COUNT, object_count(scope._kwargs.get("ids"))
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_FILTER, str(scope._kwargs.get("filter", ""))
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_DELETE_ALL, scope._kwargs.get("delete_all", False)
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {namespace} "
             f"ids={query} "
             f"filter={scope._kwargs.get('filter', '')} "
-            f"delete_all={scope._kwargs.get('delete_all', False)}")
+            f"delete_all={scope._kwargs.get('delete_all', False)}",
+        )
 
     scope._span.set_status(Status(StatusCode.OK))
 
     # Record metrics using helper
     if not disable_metrics:
-        record_db_metrics(metrics, SemanticConvention.DB_SYSTEM_PINECONE, 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, **kwargs):
+        record_db_metrics(
+            metrics,
+            SemanticConvention.DB_SYSTEM_PINECONE,
+            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,
+    **kwargs,
+):
     """
     Process vector database response and generate telemetry following OpenTelemetry conventions.
     """
@@ -220,7 +318,15 @@ 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)
+    common_vectordb_logic(
+        scope,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        instance,
+    )
 
     return response

openlit/instrumentation/premai/__init__.py

@@ -5,12 +5,11 @@ import importlib.metadata
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from wrapt import wrap_function_wrapper
 
-from openlit.instrumentation.premai.premai import (
-    chat, embedding
-)
+from openlit.instrumentation.premai.premai import chat, embedding
 
 _instruments = ("premai >= 0.3.79",)
 
+
 class PremAIInstrumentor(BaseInstrumentor):
     """
     An instrumentor for Prem AI's client library.
@@ -33,16 +32,32 @@ class PremAIInstrumentor(BaseInstrumentor):
         wrap_function_wrapper(
             "premai.api",
             "ChatCompletionsModule.create",
-            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 embedding
         wrap_function_wrapper(
             "premai.api",
             "EmbeddingsModule.create",
-            embedding(version, environment, application_name,
-                  tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            embedding(
+                version,
+                environment,
+                application_name,
+                tracer,
+                pricing_info,
+                capture_message_content,
+                metrics,
+                disable_metrics,
+            ),
         )
 
     def _uninstrument(self, **kwargs):