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

openlit/instrumentation/pydantic_ai/pydantic_ai.py

@@ -4,48 +4,101 @@ Module for monitoring Pydantic AI API calls.
 
 from openlit.instrumentation.pydantic_ai.utils import (
     common_agent_run,
-    common_agent_create
+    common_agent_create,
 )
 
-def agent_create(version, environment, application_name,
-    tracer, pricing_info, capture_message_content, metrics, disable_metrics):
 
+def agent_create(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for GenAI function call
     """
 
     def wrapper(wrapped, instance, args, kwargs):
         response = wrapped(*args, **kwargs)
-        return common_agent_create(wrapped, instance, args, kwargs, tracer,
-                                 version, environment, application_name,
-                                 capture_message_content, response=response)
+        return common_agent_create(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            tracer,
+            version,
+            environment,
+            application_name,
+            capture_message_content,
+            response=response,
+        )
 
     return wrapper
 
-def agent_run(version, environment, application_name,
-              tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+
+def agent_run(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for GenAI function call
     """
 
     def wrapper(wrapped, instance, args, kwargs):
         response = wrapped(*args, **kwargs)
-        return common_agent_run(wrapped, instance, args, kwargs, tracer,
-                                    version, environment, application_name,
-                                    capture_message_content, response=response)
+        return common_agent_run(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            tracer,
+            version,
+            environment,
+            application_name,
+            capture_message_content,
+            response=response,
+        )
 
     return wrapper
 
-def async_agent_run(version, environment, application_name,
-                    tracer, pricing_info, capture_message_content, metrics, disable_metrics):
+
+def async_agent_run(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
     """
     Generates a telemetry wrapper for GenAI function call
     """
 
     async def wrapper(wrapped, instance, args, kwargs):
         response = await wrapped(*args, **kwargs)
-        return common_agent_run(wrapped, instance, args, kwargs, tracer,
-                                          version, environment, application_name,
-                                          capture_message_content, response=response)
+        return common_agent_run(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            tracer,
+            version,
+            environment,
+            application_name,
+            capture_message_content,
+            response=response,
+        )
 
     return wrapper

openlit/instrumentation/pydantic_ai/utils.py

@@ -1,27 +1,42 @@
 """
 Pydantic AI OpenTelemetry instrumentation utility functions
 """
+
 import logging
-from opentelemetry.sdk.resources import SERVICE_NAME, TELEMETRY_SDK_NAME, DEPLOYMENT_ENVIRONMENT
-from opentelemetry.trace import Status, StatusCode, SpanKind
-from openlit.__helpers import (
-    handle_exception
+from opentelemetry.sdk.resources import (
+    SERVICE_NAME,
+    TELEMETRY_SDK_NAME,
+    DEPLOYMENT_ENVIRONMENT,
 )
+from opentelemetry.trace import Status, StatusCode, SpanKind
+from openlit.__helpers import handle_exception
 from openlit.semcov import SemanticConvention
 
 # Initialize logger for logging potential issues and operations
 logger = logging.getLogger(__name__)
 
-def set_span_attributes(span, version, operation_name, environment,
-        application_name, server_address, server_port, request_model, agent_name):
+
+def set_span_attributes(
+    span,
+    version,
+    operation_name,
+    environment,
+    application_name,
+    server_address,
+    server_port,
+    request_model,
+    agent_name,
+):
     """
     Set common OpenTelemetry span attributes for Pydantic AI operations.
     """
 
     # Set Span attributes (OTel Semconv)
-    span.set_attribute(TELEMETRY_SDK_NAME, 'openlit')
+    span.set_attribute(TELEMETRY_SDK_NAME, "openlit")
     span.set_attribute(SemanticConvention.GEN_AI_OPERATION, operation_name)
-    span.set_attribute(SemanticConvention.GEN_AI_SYSTEM, SemanticConvention.GEN_AI_SYSTEM_PYDANTIC_AI)
+    span.set_attribute(
+        SemanticConvention.GEN_AI_SYSTEM, SemanticConvention.GEN_AI_SYSTEM_PYDANTIC_AI
+    )
     span.set_attribute(SemanticConvention.GEN_AI_AGENT_NAME, agent_name)
     span.set_attribute(SemanticConvention.SERVER_ADDRESS, server_address)
     span.set_attribute(SemanticConvention.SERVER_PORT, server_port)
@@ -32,8 +47,19 @@ def set_span_attributes(span, version, operation_name, environment,
     span.set_attribute(SERVICE_NAME, application_name)
     span.set_attribute(SemanticConvention.GEN_AI_SDK_VERSION, version)
 
-def common_agent_run(wrapped, instance, args, kwargs, tracer, version, environment, application_name,
-                         capture_message_content, response):
+
+def common_agent_run(
+    wrapped,
+    instance,
+    args,
+    kwargs,
+    tracer,
+    version,
+    environment,
+    application_name,
+    capture_message_content,
+    response,
+):
     """
     Handle telemetry for Pydantic AI agent run operations.
     """
@@ -41,13 +67,27 @@ def common_agent_run(wrapped, instance, args, kwargs, tracer, version, environme
     server_address, server_port = instance.model.base_url, 443
     agent_name = instance.name or "pydantic_agent"
     request_model = str(instance.model.model_name)
-    span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK} {agent_name}'
+    span_name = (
+        f"{SemanticConvention.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK} {agent_name}"
+    )
 
     with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
         try:
-            set_span_attributes(span, version, SemanticConvention.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK,
-                                environment, application_name, server_address, server_port, request_model, agent_name)
-            span.set_attribute(SemanticConvention.GEN_AI_AGENT_DESCRIPTION, str(instance._system_prompts))
+            set_span_attributes(
+                span,
+                version,
+                SemanticConvention.GEN_AI_OPERATION_TYPE_EXECUTE_AGENT_TASK,
+                environment,
+                application_name,
+                server_address,
+                server_port,
+                request_model,
+                agent_name,
+            )
+            span.set_attribute(
+                SemanticConvention.GEN_AI_AGENT_DESCRIPTION,
+                str(instance._system_prompts),
+            )
             span.set_attribute(SemanticConvention.GEN_AI_RESPONSE_MODEL, request_model)
 
             if capture_message_content:
@@ -64,25 +104,50 @@ def common_agent_run(wrapped, instance, args, kwargs, tracer, version, environme
 
         except Exception as e:
             handle_exception(span, e)
-            logger.error('Error in trace creation: %s', e)
+            logger.error("Error in trace creation: %s", e)
             return response
 
-def common_agent_create(wrapped, instance, args, kwargs, tracer, version, environment, application_name,
-                       capture_message_content, response):
+
+def common_agent_create(
+    wrapped,
+    instance,
+    args,
+    kwargs,
+    tracer,
+    version,
+    environment,
+    application_name,
+    capture_message_content,
+    response,
+):
     """
     Handle telemetry for Pydantic AI agent creation operations.
     """
 
-    server_address, server_port = '127.0.0.1', 80
+    server_address, server_port = "127.0.0.1", 80
     agent_name = kwargs.get("name", "pydantic_agent")
-    span_name = f'{SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT} {agent_name}'
+    span_name = f"{SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT} {agent_name}"
 
     with tracer.start_as_current_span(span_name, kind=SpanKind.CLIENT) as span:
         try:
-            request_model = args[0] or kwargs.get("model", "google-gla:gemini-1.5-flash")
-            set_span_attributes(span, version, SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT,
-                environment, application_name, server_address, server_port, request_model, agent_name)
-            span.set_attribute(SemanticConvention.GEN_AI_AGENT_DESCRIPTION, str(kwargs.get("system_prompt", "")))
+            request_model = args[0] or kwargs.get(
+                "model", "google-gla:gemini-1.5-flash"
+            )
+            set_span_attributes(
+                span,
+                version,
+                SemanticConvention.GEN_AI_OPERATION_TYPE_CREATE_AGENT,
+                environment,
+                application_name,
+                server_address,
+                server_port,
+                request_model,
+                agent_name,
+            )
+            span.set_attribute(
+                SemanticConvention.GEN_AI_AGENT_DESCRIPTION,
+                str(kwargs.get("system_prompt", "")),
+            )
             span.set_attribute(SemanticConvention.GEN_AI_RESPONSE_MODEL, request_model)
 
             span.set_status(Status(StatusCode.OK))
@@ -91,5 +156,5 @@ def common_agent_create(wrapped, instance, args, kwargs, tracer, version, enviro
 
         except Exception as e:
             handle_exception(span, e)
-            logger.error('Error in trace creation: %s', e)
+            logger.error("Error in trace creation: %s", e)
             return response

openlit/instrumentation/qdrant/__init__.py

@@ -36,6 +36,7 @@ QDRANT_OPERATIONS = [
     ("query_points", "qdrant.query_points"),
 ]
 
+
 class QdrantInstrumentor(BaseInstrumentor):
     """
     An instrumentor for Qdrant client library.
@@ -60,8 +61,15 @@ class QdrantInstrumentor(BaseInstrumentor):
                 "qdrant_client",
                 f"QdrantClient.{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,
                 ),
             )
 
@@ -71,8 +79,15 @@ class QdrantInstrumentor(BaseInstrumentor):
                 "qdrant_client",
                 f"AsyncQdrantClient.{method_name}",
                 async_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/qdrant/async_qdrant.py

@@ -7,10 +7,24 @@ import time
 from opentelemetry.trace import SpanKind
 from opentelemetry import context as context_api
 from openlit.__helpers import handle_exception
-from openlit.instrumentation.qdrant.utils import DB_OPERATION_MAP, process_qdrant_response, set_server_address_and_port
+from openlit.instrumentation.qdrant.utils import (
+    DB_OPERATION_MAP,
+    process_qdrant_response,
+    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 function calls.
     """
@@ -41,10 +55,22 @@ def async_general_wrap(gen_ai_endpoint, version, environment, application_name,
             try:
                 # Process response with endpoint information
                 response = process_qdrant_response(
-                    response, db_operation, server_address, server_port,
-                    environment, application_name, metrics, start_time, span,
-                    capture_message_content, disable_metrics, version,
-                    instance=instance, args=args, endpoint=gen_ai_endpoint, **kwargs
+                    response,
+                    db_operation,
+                    server_address,
+                    server_port,
+                    environment,
+                    application_name,
+                    metrics,
+                    start_time,
+                    span,
+                    capture_message_content,
+                    disable_metrics,
+                    version,
+                    instance=instance,
+                    args=args,
+                    endpoint=gen_ai_endpoint,
+                    **kwargs,
                 )
 
             except Exception as e:

openlit/instrumentation/qdrant/qdrant.py

@@ -7,10 +7,24 @@ import time
 from opentelemetry.trace import SpanKind
 from opentelemetry import context as context_api
 from openlit.__helpers import handle_exception
-from openlit.instrumentation.qdrant.utils import DB_OPERATION_MAP, process_qdrant_response, set_server_address_and_port
+from openlit.instrumentation.qdrant.utils import (
+    DB_OPERATION_MAP,
+    process_qdrant_response,
+    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.
     """
@@ -41,10 +55,22 @@ def general_wrap(gen_ai_endpoint, version, environment, application_name, tracer
             try:
                 # Process response with endpoint information
                 response = process_qdrant_response(
-                    response, db_operation, server_address, server_port,
-                    environment, application_name, metrics, start_time, span,
-                    capture_message_content, disable_metrics, version,
-                    instance=instance, args=args, endpoint=gen_ai_endpoint, **kwargs
+                    response,
+                    db_operation,
+                    server_address,
+                    server_port,
+                    environment,
+                    application_name,
+                    metrics,
+                    start_time,
+                    span,
+                    capture_message_content,
+                    disable_metrics,
+                    version,
+                    instance=instance,
+                    args=args,
+                    endpoint=gen_ai_endpoint,
+                    **kwargs,
                 )
 
             except Exception as e: