openlit 1.34.26__py3-none-any.whl → 1.34.28__py3-none-any.whl

openlit/instrumentation/astra/__init__.py

@@ -1,178 +1,90 @@
-"""Initializer of Auto Instrumentation of AstraDB Functions"""
+"""
+OpenLIT AstraDB Instrumentation
+"""
+
 from typing import Collection
 import importlib.metadata
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
 from wrapt import wrap_function_wrapper
 
 from openlit.instrumentation.astra.astra import general_wrap
-# from openlit.instrumentation.astra.async_astra import asyc_general_wrap
+from openlit.instrumentation.astra.async_astra import async_general_wrap
+
+_instruments = ("astrapy >= 1.5.2",)
+
+# AstraDB sync operations
+ASTRA_SYNC_OPERATIONS = [
+    # Database operations
+    ("astrapy.database", "Database.create_collection", "astra.create_collection"),
+    ("astrapy.database", "Database.drop_collection", "astra.drop_collection"),
+
+    # Collection operations
+    ("astrapy.collection", "Collection.insert_one", "astra.insert_one"),
+    ("astrapy.collection", "Collection.insert_many", "astra.insert_many"),
+    ("astrapy.collection", "Collection.update_one", "astra.update_one"),
+    ("astrapy.collection", "Collection.update_many", "astra.update_many"),
+    ("astrapy.collection", "Collection.find_one_and_update", "astra.find_one_and_update"),
+    ("astrapy.collection", "Collection.find", "astra.find"),
+    ("astrapy.collection", "Collection.replace_one", "astra.replace_one"),
+    ("astrapy.collection", "Collection.find_one_and_delete", "astra.find_one_and_delete"),
+    ("astrapy.collection", "Collection.delete_one", "astra.delete_one"),
+    ("astrapy.collection", "Collection.delete_many", "astra.delete_many"),
+]
+
+# AstraDB async operations
+ASTRA_ASYNC_OPERATIONS = [
+    # Async Database operations
+    ("astrapy.database", "AsyncDatabase.create_collection", "astra.create_collection"),
+    ("astrapy.database", "AsyncDatabase.drop_collection", "astra.drop_collection"),
 
-_instruments = ('astrapy >= 1.5.2',)
+    # Async Collection operations
+    ("astrapy.collection", "AsyncCollection.insert_one", "astra.insert_one"),
+    ("astrapy.collection", "AsyncCollection.insert_many", "astra.insert_many"),
+    ("astrapy.collection", "AsyncCollection.update_one", "astra.update_one"),
+    ("astrapy.collection", "AsyncCollection.update_many", "astra.update_many"),
+    ("astrapy.collection", "AsyncCollection.find_one_and_update", "astra.find_one_and_update"),
+    ("astrapy.collection", "AsyncCollection.find", "astra.find"),
+    ("astrapy.collection", "AsyncCollection.replace_one", "astra.replace_one"),
+    ("astrapy.collection", "AsyncCollection.find_one_and_delete", "astra.find_one_and_delete"),
+    ("astrapy.collection", "AsyncCollection.delete_one", "astra.delete_one"),
+    ("astrapy.collection", "AsyncCollection.delete_many", "astra.delete_many"),
+]
 
 class AstraInstrumentor(BaseInstrumentor):
-    """An instrumentor for AstraDB's client library."""
+    """
+    An instrumentor for AstraDB's client library.
+    """
 
     def instrumentation_dependencies(self) -> Collection[str]:
         return _instruments
 
     def _instrument(self, **kwargs):
-        application_name = kwargs.get('application_name')
-        environment = kwargs.get('environment')
-        tracer = kwargs.get('tracer')
-        metrics = kwargs.get('metrics_dict')
-        pricing_info = kwargs.get('pricing_info')
-        capture_message_content = kwargs.get('capture_message_content')
-        disable_metrics = kwargs.get('disable_metrics')
-        version = importlib.metadata.version('astrapy')
+        version = importlib.metadata.version("astrapy")
+        environment = kwargs.get("environment", "default")
+        application_name = kwargs.get("application_name", "default")
+        tracer = kwargs.get("tracer")
+        pricing_info = kwargs.get("pricing_info", {})
+        capture_message_content = kwargs.get("capture_message_content", False)
+        metrics = kwargs.get("metrics_dict")
+        disable_metrics = kwargs.get("disable_metrics")
 
-        # Sync
-        wrap_function_wrapper(
-            'astrapy.database',
-            'Database.create_collection',
-            general_wrap('astra.create_collection', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.database',
-            'Database.drop_collection',
-            general_wrap('astra.drop_collection', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.insert_one',
-            general_wrap('astra.insert', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.insert_many',
-            general_wrap('astra.insert', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.update_one',
-            general_wrap('astra.update', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.update_many',
-            general_wrap('astra.update', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.find_one_and_update',
-            general_wrap('astra.find_one_and_update', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.find',
-            general_wrap('astra.find', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.replace_one',
-            general_wrap('astra.replace_one', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.find_one_and_delete',
-            general_wrap('astra.find_one_and_delete', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.delete_one',
-            general_wrap('astra.delete', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'Collection.delete_many',
-            general_wrap('astra.delete', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
+        # Wrap sync operations
+        for module, class_method, endpoint in ASTRA_SYNC_OPERATIONS:
+            wrap_function_wrapper(
+                module,
+                class_method,
+                general_wrap(endpoint, version, environment, application_name,
+                    tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            )
 
-        # ASync
-        wrap_function_wrapper(
-            'astrapy.database',
-            'AsyncDatabase.create_collection',
-            general_wrap('astra.create_collection', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.database',
-            'AsyncDatabase.drop_collection',
-            general_wrap('astra.drop_collection', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.insert_one',
-            general_wrap('astra.insert_one', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.insert_many',
-            general_wrap('astra.insert_many', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.update_one',
-            general_wrap('astra.update_one', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.update_many',
-            general_wrap('astra.update_many', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.find_one_and_update',
-            general_wrap('astra.find_one_and_update', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.find',
-            general_wrap('astra.find', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.replace_one',
-            general_wrap('astra.replace_one', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.find_one_and_delete',
-            general_wrap('astra.find_one_and_delete', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.delete_one',
-            general_wrap('astra.delete_one', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
-        wrap_function_wrapper(
-            'astrapy.collection',
-            'AsyncCollection.delete_many',
-            general_wrap('astra.delete_many', version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics),
-        )
+        # Wrap async operations
+        for module, class_method, endpoint in ASTRA_ASYNC_OPERATIONS:
+            wrap_function_wrapper(
+                module,
+                class_method,
+                async_general_wrap(endpoint, version, environment, application_name,
+                    tracer, pricing_info, capture_message_content, metrics, disable_metrics),
+            )
 
     def _uninstrument(self, **kwargs):
-        # Proper uninstrumentation logic to revert patched methods
         pass

openlit/instrumentation/astra/astra.py

@@ -1,44 +1,54 @@
 """
-Module for monitoring AstraDB.
+Module for monitoring AstraDB API calls.
 """
 
 import time
 from opentelemetry.trace import SpanKind
+from opentelemetry import context as context_api
+from openlit.__helpers import handle_exception
 from openlit.instrumentation.astra.utils import (
+    process_astra_response,
     DB_OPERATION_MAP,
-    process_db_operations
+    set_server_address_and_port,
 )
-from openlit.semcov import SemanticConvention
 
 def general_wrap(gen_ai_endpoint, version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+    tracer, pricing_info, capture_message_content, metrics, disable_metrics):
     """
-    Generates a telemetry wrapper for VectorDB function call
+    Generates a telemetry wrapper for AstraDB function calls.
     """
-
     def wrapper(wrapped, instance, args, kwargs):
-        """
-        Wraps the VectorDB function call.
-        """
+        # CRITICAL: Suppression check
+        if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
+            return wrapped(*args, **kwargs)
+
+        # Get server address and port using the standard helper
+        server_address, server_port = set_server_address_and_port(instance)
+
+        db_operation = DB_OPERATION_MAP.get(gen_ai_endpoint, "unknown")
 
-        db_operation = DB_OPERATION_MAP.get(gen_ai_endpoint, "UNKNOWN")
-        if db_operation == SemanticConvention.DB_OPERATION_REPLACE and kwargs.get('upsert'):
-            db_operation = SemanticConvention.DB_OPERATION_UPSERT
+        # Handle upsert detection for replace operations
+        if db_operation == "DB_OPERATION_REPLACE" and kwargs.get("upsert"):
+            db_operation = "DB_OPERATION_UPSERT"
 
-        span_name = f"{db_operation} {instance.name}"
+        # Get collection name for span naming
+        collection_name = getattr(instance, "name", "unknown")
+        span_name = f"{db_operation} {collection_name}"
 
         with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
             start_time = time.time()
             response = wrapped(*args, **kwargs)
-            server_address = getattr(getattr(instance, 'database', instance), 'api_endpoint', '')
-            server_port = 443
-            collection_name = instance.name
-            response = process_db_operations(
-                response, span, start_time, gen_ai_endpoint,
-                version, environment, application_name, capture_message_content,
-                metrics, disable_metrics, server_address, server_port,
-                collection_name, db_operation, kwargs, args
-            )
+
+            try:
+                # Process response and generate telemetry
+                response = process_astra_response(
+                    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:
+                handle_exception(span, e)
 
             return response
 

openlit/instrumentation/astra/async_astra.py

@@ -1,44 +1,54 @@
 """
-Module for monitoring AstraDB.
+Module for monitoring AstraDB async API calls.
 """
 
 import time
 from opentelemetry.trace import SpanKind
+from opentelemetry import context as context_api
+from openlit.__helpers import handle_exception
 from openlit.instrumentation.astra.utils import (
+    process_astra_response,
     DB_OPERATION_MAP,
-    process_db_operations
+    set_server_address_and_port,
 )
-from openlit.semcov import SemanticConvention
 
 def async_general_wrap(gen_ai_endpoint, version, environment, application_name,
-                 tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+    tracer, pricing_info, capture_message_content, metrics, disable_metrics):
     """
-    Generates a telemetry wrapper for VectorDB function call
+    Generates a telemetry wrapper for AstraDB async function calls.
     """
-
     async def wrapper(wrapped, instance, args, kwargs):
-        """
-        Wraps the VectorDB function call.
-        """
+        # CRITICAL: Suppression check
+        if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
+            return wrapped(*args, **kwargs)
+
+        # Get server address and port using the standard helper
+        server_address, server_port = set_server_address_and_port(instance)
+
+        db_operation = DB_OPERATION_MAP.get(gen_ai_endpoint, "unknown")
 
-        db_operation = DB_OPERATION_MAP.get(gen_ai_endpoint, "UNKNOWN")
-        if db_operation == SemanticConvention.DB_OPERATION_REPLACE and kwargs.get('upsert'):
-            db_operation = SemanticConvention.DB_OPERATION_UPSERT
+        # Handle upsert detection for replace operations
+        if db_operation == "DB_OPERATION_REPLACE" and kwargs.get("upsert"):
+            db_operation = "DB_OPERATION_UPSERT"
 
-        span_name = f"{db_operation} {instance.name}"
+        # Get collection name for span naming
+        collection_name = getattr(instance, "name", "unknown")
+        span_name = f"{db_operation} {collection_name}"
 
         with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
             start_time = time.time()
             response = await wrapped(*args, **kwargs)
-            server_address = getattr(getattr(instance, 'database', instance), 'api_endpoint', '')
-            server_port = 443
-            collection_name = instance.name
-            response = process_db_operations(
-                response, span, start_time, gen_ai_endpoint,
-                version, environment, application_name, capture_message_content,
-                metrics, disable_metrics, server_address, server_port,
-                collection_name, db_operation, kwargs, args
-            )
+
+            try:
+                # Process response and generate telemetry
+                response = process_astra_response(
+                    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:
+                handle_exception(span, e)
 
             return response