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

openlit/instrumentation/mem0/mem0.py

@@ -5,15 +5,27 @@ Module for monitoring mem0 applications.
 
 import logging
 from opentelemetry.trace import SpanKind, Status, StatusCode
-from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
+from opentelemetry.sdk.resources import (
+    SERVICE_NAME,
+    TELEMETRY_SDK_NAME,
+    DEPLOYMENT_ENVIRONMENT,
+)
 from openlit.__helpers import handle_exception
 from openlit.semcov import SemanticConvention
 
 # Initialize logger for logging potential issues and operations
 logger = logging.getLogger(__name__)
 
-def mem0_wrap(gen_ai_endpoint, version, environment, application_name,
-                 tracer, pricing_info, capture_message_content):
+
+def mem0_wrap(
+    gen_ai_endpoint,
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+):
     """
     Creates a wrapper around a function call to trace and log its execution metrics.
 
@@ -54,51 +66,60 @@ def mem0_wrap(gen_ai_endpoint, version, environment, application_name,
         on the span based on the function's execution and response, and ensures
         errors are handled and logged appropriately.
         """
-        with tracer.start_as_current_span(gen_ai_endpoint, kind= SpanKind.CLIENT) as span:
+        with tracer.start_as_current_span(
+            gen_ai_endpoint, kind=SpanKind.CLIENT
+        ) as span:
             response = wrapped(*args, **kwargs)
 
             try:
                 span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
-                span.set_attribute(SemanticConvention.GEN_AI_SYSTEM,
-                                    SemanticConvention.GEN_AI_SYSTEM_MEM0)
-                span.set_attribute(SemanticConvention.GEN_AI_ENDPOINT,
-                                    gen_ai_endpoint)
-                span.set_attribute(DEPLOYMENT_ENVIRONMENT,
-                                    environment)
-                span.set_attribute(SemanticConvention.GEN_AI_OPERATION,
-                                    SemanticConvention.GEN_AI_OPERATION_TYPE_FRAMEWORK)
-                span.set_attribute(SERVICE_NAME,
-                                    application_name)
+                span.set_attribute(
+                    SemanticConvention.GEN_AI_SYSTEM,
+                    SemanticConvention.GEN_AI_SYSTEM_MEM0,
+                )
+                span.set_attribute(SemanticConvention.GEN_AI_ENDPOINT, gen_ai_endpoint)
+                span.set_attribute(DEPLOYMENT_ENVIRONMENT, environment)
+                span.set_attribute(
+                    SemanticConvention.GEN_AI_OPERATION,
+                    SemanticConvention.GEN_AI_OPERATION_TYPE_FRAMEWORK,
+                )
+                span.set_attribute(SERVICE_NAME, application_name)
 
                 if gen_ai_endpoint == "mem0.memory_add":
-                    span.set_attribute(SemanticConvention.DB_METADATA,
-                                        str(kwargs.get('metadata', '')))
+                    span.set_attribute(
+                        SemanticConvention.DB_METADATA, str(kwargs.get("metadata", ""))
+                    )
                     if response:
-                        span.set_attribute(SemanticConvention.GEN_AI_DATA_SOURCES,
-                                        len(response))
+                        span.set_attribute(
+                            SemanticConvention.GEN_AI_DATA_SOURCES, len(response)
+                        )
 
                 elif gen_ai_endpoint == "mem0.memory_get_all":
                     if response:
-                        span.set_attribute(SemanticConvention.GEN_AI_DATA_SOURCES,
-                                        len(response))
+                        span.set_attribute(
+                            SemanticConvention.GEN_AI_DATA_SOURCES, len(response)
+                        )
 
                 elif gen_ai_endpoint == "mem0.memory_get":
                     if response:
-                        span.set_attribute(SemanticConvention.GEN_AI_DATA_SOURCES,
-                                        len(response))
+                        span.set_attribute(
+                            SemanticConvention.GEN_AI_DATA_SOURCES, len(response)
+                        )
 
                 elif gen_ai_endpoint == "mem0.memory_search":
-                    span.set_attribute(SemanticConvention.DB_STATEMENT,
-                                    kwargs.get("query", ""))
+                    span.set_attribute(
+                        SemanticConvention.DB_STATEMENT, kwargs.get("query", "")
+                    )
 
                 elif gen_ai_endpoint == "mem0.memory_update":
-                    span.set_attribute(SemanticConvention.DB_UPDATE_ID,
-                                    kwargs.get("memory_id", ""))
+                    span.set_attribute(
+                        SemanticConvention.DB_UPDATE_ID, kwargs.get("memory_id", "")
+                    )
 
                 elif gen_ai_endpoint == "mem0.memory_delete":
-                    span.set_attribute(SemanticConvention.DB_DELETE_ID,
-                                    kwargs.get("memory_id", ""))
-
+                    span.set_attribute(
+                        SemanticConvention.DB_DELETE_ID, kwargs.get("memory_id", "")
+                    )
 
                 span.set_status(Status(StatusCode.OK))
 

openlit/instrumentation/milvus/__init__.py

@@ -23,6 +23,7 @@ MILVUS_OPERATIONS = [
     ("delete", "milvus.delete"),
 ]
 
+
 class MilvusInstrumentor(BaseInstrumentor):
     """
     An instrumentor for Milvus's client library.
@@ -47,8 +48,15 @@ class MilvusInstrumentor(BaseInstrumentor):
                 "pymilvus",
                 f"MilvusClient.{method_name}",
                 general_wrap(
-                    endpoint, version, environment, application_name, tracer,
-                    pricing_info, capture_message_content, metrics, disable_metrics
+                    endpoint,
+                    version,
+                    environment,
+                    application_name,
+                    tracer,
+                    pricing_info,
+                    capture_message_content,
+                    metrics,
+                    disable_metrics,
                 ),
             )
 

openlit/instrumentation/milvus/milvus.py

@@ -12,11 +12,22 @@ from openlit.instrumentation.milvus.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 Milvus function calls.
     """
+
     def wrapper(wrapped, instance, args, kwargs):
         # CRITICAL: Suppression check
         if context_api.get_value(context_api._SUPPRESS_INSTRUMENTATION_KEY):
@@ -27,7 +38,9 @@ def general_wrap(gen_ai_endpoint, version, environment, application_name,
 
         db_operation = DB_OPERATION_MAP.get(gen_ai_endpoint, "unknown")
         if db_operation == "create_collection":
-            collection_name = kwargs.get("collection_name") or (args[0] if args else "unknown")
+            collection_name = kwargs.get("collection_name") or (
+                args[0] if args else "unknown"
+            )
         else:
             collection_name = kwargs.get("collection_name", "unknown")
         span_name = f"{db_operation} {collection_name}"
@@ -39,9 +52,21 @@ def general_wrap(gen_ai_endpoint, version, environment, application_name,
             try:
                 # Process response and generate telemetry
                 response = process_milvus_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/milvus/utils.py

@@ -23,12 +23,14 @@ DB_OPERATION_MAP = {
     "milvus.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 Milvus client instance.
@@ -68,9 +70,18 @@ def set_server_address_and_port(instance):
 
     return server_address, server_port
 
-def common_milvus_logic(scope, environment, application_name,
-    metrics, capture_message_content, disable_metrics, version,
-    instance=None, endpoint=None):
+
+def common_milvus_logic(
+    scope,
+    environment,
+    application_name,
+    metrics,
+    capture_message_content,
+    disable_metrics,
+    version,
+    instance=None,
+    endpoint=None,
+):
     """
     Process Milvus database request and generate telemetry.
 
@@ -88,153 +99,233 @@ def common_milvus_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_MILVUS,
-        scope._server_address, scope._server_port, environment, application_name, version)
+    common_db_span_attributes(
+        scope,
+        SemanticConvention.DB_SYSTEM_MILVUS,
+        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,
+    )
 
     if scope._db_operation == SemanticConvention.DB_OPERATION_CREATE_COLLECTION:
         collection_name = scope._kwargs.get("collection_name", "unknown")
 
         # Standard database attributes
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_DIMENSION,
-            scope._kwargs.get("dimension", -1))
-
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_DIMENSION,
+            scope._kwargs.get("dimension", -1),
+        )
+
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {collection_name} "
-            f"dimension={scope._kwargs.get('dimension', 'None')}")
+            f"dimension={scope._kwargs.get('dimension', 'None')}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_DELETE_COLLECTION:
         collection_name = scope._kwargs.get("collection_name", "unknown")
 
         # Standard database attributes
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
 
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
-            f"{scope._db_operation} {collection_name}")
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
+            f"{scope._db_operation} {collection_name}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_INSERT:
         collection_name = scope._kwargs.get("collection_name", "unknown")
         data = scope._kwargs.get("data", [])
 
         # Standard database attributes
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(data))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_COUNT, object_count(data)
+        )
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(data))
 
         # Response metrics
         if scope._response and scope._response.get("insert_count"):
-            scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
-                scope._response["insert_count"])
+            scope._span.set_attribute(
+                SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
+                scope._response["insert_count"],
+            )
 
         if scope._response and scope._response.get("cost"):
-            scope._span.set_attribute(SemanticConvention.DB_OPERATION_COST,
-                scope._response["cost"])
+            scope._span.set_attribute(
+                SemanticConvention.DB_OPERATION_COST, scope._response["cost"]
+            )
 
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
-            f"{scope._db_operation} {collection_name} "
-            f"data_count={object_count(data)}")
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
+            f"{scope._db_operation} {collection_name} data_count={object_count(data)}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_UPSERT:
         collection_name = scope._kwargs.get("collection_name", "unknown")
         data = scope._kwargs.get("data", [])
 
         # Standard database attributes
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_COUNT, object_count(data))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_COUNT, object_count(data)
+        )
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(data))
 
         # Response metrics
         if scope._response and scope._response.get("upsert_count"):
-            scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
-                scope._response["upsert_count"])
+            scope._span.set_attribute(
+                SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
+                scope._response["upsert_count"],
+            )
 
         if scope._response and scope._response.get("cost"):
-            scope._span.set_attribute(SemanticConvention.DB_OPERATION_COST,
-                scope._response["cost"])
+            scope._span.set_attribute(
+                SemanticConvention.DB_OPERATION_COST, scope._response["cost"]
+            )
 
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
-            f"{scope._db_operation} {collection_name} "
-            f"data_count={object_count(data)}")
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
+            f"{scope._db_operation} {collection_name} data_count={object_count(data)}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_SEARCH:
         collection_name = scope._kwargs.get("collection_name", "unknown")
         data = scope._kwargs.get("data", [])
 
         # Standard database attributes
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(data))
-        scope._span.set_attribute(SemanticConvention.DB_VECTOR_QUERY_TOP_K,
-            scope._kwargs.get("limit", -1))
+        scope._span.set_attribute(
+            SemanticConvention.DB_VECTOR_QUERY_TOP_K, scope._kwargs.get("limit", -1)
+        )
 
         # Search specific attributes
-        scope._span.set_attribute(SemanticConvention.DB_FILTER,
-            str(scope._kwargs.get("filter", "")))
+        scope._span.set_attribute(
+            SemanticConvention.DB_FILTER, str(scope._kwargs.get("filter", ""))
+        )
 
         # Response metrics
         if scope._response and isinstance(scope._response, list):
-            scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
-                len(scope._response))
+            scope._span.set_attribute(
+                SemanticConvention.DB_RESPONSE_RETURNED_ROWS, len(scope._response)
+            )
 
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {collection_name} "
             f"data={str(data)[:100]}... "
-            f"limit={scope._kwargs.get('limit', 'None')}")
+            f"limit={scope._kwargs.get('limit', 'None')}",
+        )
 
-    elif scope._db_operation in [SemanticConvention.DB_OPERATION_QUERY, SemanticConvention.DB_OPERATION_GET]:
+    elif scope._db_operation in [
+        SemanticConvention.DB_OPERATION_QUERY,
+        SemanticConvention.DB_OPERATION_GET,
+    ]:
         collection_name = scope._kwargs.get("collection_name", "unknown")
         output_fields = scope._kwargs.get("output_fields", [])
 
         # Standard database attributes
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
         scope._span.set_attribute(SemanticConvention.DB_QUERY_TEXT, str(output_fields))
 
         # Query specific attributes
-        scope._span.set_attribute(SemanticConvention.DB_FILTER,
-            str(scope._kwargs.get("filter", "")))
+        scope._span.set_attribute(
+            SemanticConvention.DB_FILTER, str(scope._kwargs.get("filter", ""))
+        )
 
         # Response metrics
         if scope._response and isinstance(scope._response, list):
-            scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
-                len(scope._response))
+            scope._span.set_attribute(
+                SemanticConvention.DB_RESPONSE_RETURNED_ROWS, len(scope._response)
+            )
 
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {collection_name} "
             f"output_fields={output_fields} "
-            f"filter={scope._kwargs.get('filter', 'None')}")
+            f"filter={scope._kwargs.get('filter', 'None')}",
+        )
 
     elif scope._db_operation == SemanticConvention.DB_OPERATION_DELETE:
         collection_name = scope._kwargs.get("collection_name", "unknown")
 
         # Standard database attributes
-        scope._span.set_attribute(SemanticConvention.DB_COLLECTION_NAME, collection_name)
-        scope._span.set_attribute(SemanticConvention.DB_FILTER,
-            str(scope._kwargs.get("filter", "")))
+        scope._span.set_attribute(
+            SemanticConvention.DB_COLLECTION_NAME, collection_name
+        )
+        scope._span.set_attribute(
+            SemanticConvention.DB_FILTER, str(scope._kwargs.get("filter", ""))
+        )
 
         # Response metrics
         if scope._response and scope._response.get("delete_count"):
-            scope._span.set_attribute(SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
-                scope._response["delete_count"])
+            scope._span.set_attribute(
+                SemanticConvention.DB_RESPONSE_RETURNED_ROWS,
+                scope._response["delete_count"],
+            )
 
-        scope._span.set_attribute(SemanticConvention.DB_QUERY_SUMMARY,
+        scope._span.set_attribute(
+            SemanticConvention.DB_QUERY_SUMMARY,
             f"{scope._db_operation} {collection_name} "
-            f"filter={scope._kwargs.get('filter', 'None')}")
+            f"filter={scope._kwargs.get('filter', 'None')}",
+        )
 
     scope._span.set_status(Status(StatusCode.OK))
 
     # Record metrics using helper
     if not disable_metrics:
-        record_db_metrics(metrics, SemanticConvention.DB_SYSTEM_MILVUS,
-            scope._server_address, scope._server_port, environment, application_name,
-            scope._start_time, scope._end_time)
-
-def process_milvus_response(response, db_operation, server_address, server_port,
-    environment, application_name, metrics, start_time, span,
-    capture_message_content, disable_metrics, version, instance, args, **kwargs):
+        record_db_metrics(
+            metrics,
+            SemanticConvention.DB_SYSTEM_MILVUS,
+            scope._server_address,
+            scope._server_port,
+            environment,
+            application_name,
+            scope._start_time,
+            scope._end_time,
+        )
+
+
+def process_milvus_response(
+    response,
+    db_operation,
+    server_address,
+    server_port,
+    environment,
+    application_name,
+    metrics,
+    start_time,
+    span,
+    capture_message_content,
+    disable_metrics,
+    version,
+    instance,
+    args,
+    **kwargs,
+):
     """
     Process Milvus response and generate telemetry.
 
@@ -270,7 +361,15 @@ def process_milvus_response(response, db_operation, server_address, server_port,
     scope._kwargs = kwargs
 
     # Process the response using common logic
-    common_milvus_logic(scope, environment, application_name,
-        metrics, capture_message_content, disable_metrics, version, instance)
+    common_milvus_logic(
+        scope,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        instance,
+    )
 
     return response