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

openlit/instrumentation/ag2/utils.py

@@ -147,7 +147,7 @@ def common_agent_logic(
     scope._span.set_status(Status(StatusCode.OK))
 
     # Metrics
-    if not disable_metrics and hasattr(scope, "_input_tokens"):
+    if not disable_metrics and metrics is not None and hasattr(scope, "_input_tokens"):
         record_completion_metrics(
             metrics,
             operation_type,
@@ -196,7 +196,7 @@ def process_agent_creation(
     scope._end_time = time.time()
     scope._span = span
     scope._agent_name = agent_name
-    scope._request_model = llm_config.get("model", "gpt-4o")
+    scope._request_model = llm_config.get("model", "unknown")
     scope._response_model = scope._request_model
     scope._system_message = system_message
     scope._server_address, scope._server_port = server_address, server_port
@@ -276,3 +276,344 @@ def process_agent_run(
     )
 
     return response
+
+
+def process_agent_generate_reply(
+    response,
+    agent_name,
+    request_model,
+    messages,
+    sender,
+    pricing_info,
+    server_port,
+    server_address,
+    environment,
+    application_name,
+    metrics,
+    start_time,
+    span,
+    capture_message_content=False,
+    disable_metrics=False,
+    version="1.0.0",
+    **kwargs,
+):
+    """
+    Process agent generate_reply and generate Telemetry
+    """
+
+    # Create scope object
+    scope = type("GenericScope", (), {})()
+
+    scope._start_time = start_time
+    scope._end_time = time.time()
+    scope._span = span
+    scope._agent_name = agent_name
+    scope._request_model = request_model
+    scope._response_model = request_model
+    scope._server_address, scope._server_port = server_address, server_port
+    scope._messages = messages
+    scope._sender_name = getattr(sender, "name", "Unknown") if sender else "Unknown"
+
+    # Set agent-specific attributes
+    span.set_attribute(SemanticConvention.GEN_AI_AGENT_MESSAGE_TYPE, "generate_reply")
+    span.set_attribute(SemanticConvention.GEN_AI_AGENT_SENDER, scope._sender_name)
+
+    # Process response content
+    if response and isinstance(response, str):
+        scope._response_content = response
+    elif response and hasattr(response, "content"):
+        scope._response_content = response.content
+    else:
+        scope._response_content = str(response) if response else ""
+
+    # Try to extract token information if available
+    try:
+        # Mock token calculation for generate_reply
+        scope._input_tokens = len(str(messages)) // 4 if messages else 0
+        scope._output_tokens = len(scope._response_content) // 4
+        scope._cost = get_chat_model_cost(
+            request_model, pricing_info, scope._input_tokens, scope._output_tokens
+        )
+    except Exception:
+        scope._input_tokens = 0
+        scope._output_tokens = 0
+        scope._cost = 0.0
+
+    common_agent_logic(
+        scope,
+        pricing_info,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT,
+    )
+
+    return response
+
+
+def process_agent_receive(
+    message,
+    agent_name,
+    sender_name,
+    agent_instance,
+    pricing_info,
+    server_port,
+    server_address,
+    environment,
+    application_name,
+    metrics,
+    start_time,
+    span,
+    capture_message_content=False,
+    disable_metrics=False,
+    version="1.0.0",
+    **kwargs,
+):
+    """
+    Process agent receive and generate Telemetry
+    """
+
+    # Create scope object
+    scope = type("GenericScope", (), {})()
+
+    scope._start_time = start_time
+    scope._end_time = time.time()
+    scope._span = span
+    scope._agent_name = agent_name
+    scope._sender_name = sender_name
+    scope._server_address, scope._server_port = server_address, server_port
+    scope._message = message
+
+    # Extract model from agent instance
+    if hasattr(agent_instance, "llm_config") and isinstance(
+        agent_instance.llm_config, dict
+    ):
+        scope._request_model = agent_instance.llm_config.get("model", "unknown")
+    else:
+        scope._request_model = "unknown"
+    scope._response_model = scope._request_model
+
+    # Set agent-specific attributes
+    span.set_attribute(SemanticConvention.GEN_AI_AGENT_MESSAGE_TYPE, "receive")
+    span.set_attribute(SemanticConvention.GEN_AI_AGENT_SENDER, sender_name)
+
+    # Content capture for received message
+    if capture_message_content:
+        span.set_attribute(SemanticConvention.GEN_AI_CONTENT_PROMPT, str(message))
+
+    common_agent_logic(
+        scope,
+        pricing_info,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT,
+    )
+
+
+def process_agent_send(
+    message,
+    agent_name,
+    recipient_name,
+    agent_instance,
+    pricing_info,
+    server_port,
+    server_address,
+    environment,
+    application_name,
+    metrics,
+    start_time,
+    span,
+    capture_message_content=False,
+    disable_metrics=False,
+    version="1.0.0",
+    **kwargs,
+):
+    """
+    Process agent send and generate Telemetry
+    """
+
+    # Create scope object
+    scope = type("GenericScope", (), {})()
+
+    scope._start_time = start_time
+    scope._end_time = time.time()
+    scope._span = span
+    scope._agent_name = agent_name
+    scope._recipient_name = recipient_name
+    scope._server_address, scope._server_port = server_address, server_port
+    scope._message = message
+
+    # Extract model from agent instance
+    if hasattr(agent_instance, "llm_config") and isinstance(
+        agent_instance.llm_config, dict
+    ):
+        scope._request_model = agent_instance.llm_config.get("model", "unknown")
+    else:
+        scope._request_model = "unknown"
+    scope._response_model = scope._request_model
+
+    # Set agent-specific attributes
+    span.set_attribute(SemanticConvention.GEN_AI_AGENT_MESSAGE_TYPE, "send")
+    span.set_attribute(SemanticConvention.GEN_AI_AGENT_RECIPIENT, recipient_name)
+
+    # Content capture for sent message
+    if capture_message_content:
+        span.set_attribute(SemanticConvention.GEN_AI_CONTENT_COMPLETION, str(message))
+
+    common_agent_logic(
+        scope,
+        pricing_info,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        SemanticConvention.GEN_AI_OPERATION_TYPE_CHAT,
+    )
+
+
+def process_groupchat_operation(
+    group_name,
+    participants,
+    messages,
+    sender,
+    max_turns,
+    request_model,
+    pricing_info,
+    server_port,
+    server_address,
+    environment,
+    application_name,
+    metrics,
+    start_time,
+    span,
+    capture_message_content=False,
+    disable_metrics=False,
+    version="1.0.0",
+    **kwargs,
+):
+    """
+    Process GroupChat operation and generate Telemetry
+    """
+
+    # Create scope object
+    scope = type("GenericScope", (), {})()
+
+    scope._start_time = start_time
+    scope._end_time = time.time()
+    scope._span = span
+    scope._group_name = group_name
+    scope._participants = participants
+    scope._server_address, scope._server_port = server_address, server_port
+    scope._sender_name = getattr(sender, "name", "Unknown") if sender else "Unknown"
+
+    # Add required model attributes for common_agent_logic
+    scope._request_model = request_model
+    scope._response_model = request_model
+
+    # Set agent name for groupchat
+    scope._agent_name = group_name
+
+    # Set GroupChat-specific attributes
+    span.set_attribute(
+        SemanticConvention.GEN_AI_GROUPCHAT_PARTICIPANTS, ",".join(participants)
+    )
+    span.set_attribute(
+        SemanticConvention.GEN_AI_WORKFLOW_AGENT_COUNT, len(participants)
+    )
+    span.set_attribute(SemanticConvention.GEN_AI_WORKFLOW_EXECUTION_TYPE, "groupchat")
+
+    if max_turns:
+        span.set_attribute(SemanticConvention.GEN_AI_GROUPCHAT_TURN_COUNT, max_turns)
+
+    # Content capture for GroupChat
+    if capture_message_content and messages:
+        span.set_attribute(SemanticConvention.GEN_AI_CONTENT_PROMPT, str(messages))
+
+    # Use framework operation type for GroupChat
+    common_agent_logic(
+        scope,
+        pricing_info,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        SemanticConvention.GEN_AI_OPERATION_TYPE_FRAMEWORK,
+    )
+
+
+def process_speaker_selection(
+    last_speaker,
+    selected_speaker,
+    selector,
+    agents,
+    request_model,
+    pricing_info,
+    server_port,
+    server_address,
+    environment,
+    application_name,
+    metrics,
+    start_time,
+    span,
+    capture_message_content=False,
+    disable_metrics=False,
+    version="1.0.0",
+    **kwargs,
+):
+    """
+    Process speaker selection and generate Telemetry
+    """
+
+    # Create scope object
+    scope = type("GenericScope", (), {})()
+
+    scope._start_time = start_time
+    scope._end_time = time.time()
+    scope._span = span
+    scope._last_speaker = last_speaker
+    scope._selected_speaker = selected_speaker
+    scope._server_address, scope._server_port = server_address, server_port
+
+    # Add required model attributes for common_agent_logic
+    scope._request_model = request_model
+    scope._response_model = request_model
+
+    # Set agent name for speaker selection
+    scope._agent_name = "speaker_selection"
+
+    # Set speaker selection attributes
+    span.set_attribute(
+        SemanticConvention.GEN_AI_GROUPCHAT_SPEAKER_SELECTION, selected_speaker
+    )
+    span.set_attribute(SemanticConvention.GEN_AI_AGENT_SENDER, last_speaker)
+
+    if selector:
+        span.set_attribute(SemanticConvention.GEN_AI_AGENT_ROLE, "selector")
+
+    # Set agent count
+    if agents:
+        span.set_attribute(SemanticConvention.GEN_AI_WORKFLOW_AGENT_COUNT, len(agents))
+
+    # Use agent operation type for speaker selection
+    common_agent_logic(
+        scope,
+        pricing_info,
+        environment,
+        application_name,
+        metrics,
+        capture_message_content,
+        disable_metrics,
+        version,
+        SemanticConvention.GEN_AI_OPERATION_TYPE_AGENT,
+    )

openlit/instrumentation/pydantic_ai/__init__.py

@@ -8,6 +8,12 @@ from wrapt import wrap_function_wrapper
 from openlit.instrumentation.pydantic_ai.pydantic_ai import (
     agent_create,
     agent_run,
+    graph_execution,
+    user_prompt_processing,
+    model_request_processing,
+    tool_calls_processing,
+)
+from openlit.instrumentation.pydantic_ai.async_pydantic_ai import (
     async_agent_run,
 )
 
@@ -77,6 +83,88 @@ class PydanticAIInstrumentor(BaseInstrumentor):
             ),
         )
 
+        # Enhanced instrumentation for richer span hierarchy
+        # These wrap internal Pydantic AI graph execution components
+        try:
+            # Agent.iter() - Graph execution iterator
+            wrap_function_wrapper(
+                "pydantic_ai.agent",
+                "Agent.iter",
+                graph_execution(
+                    version,
+                    environment,
+                    application_name,
+                    tracer,
+                    pricing_info,
+                    capture_message_content,
+                    metrics,
+                    disable_metrics,
+                ),
+            )
+        except Exception:
+            # If Agent.iter doesn't exist, skip this instrumentation
+            pass
+
+        try:
+            # UserPromptNode.run() - User prompt processing
+            wrap_function_wrapper(
+                "pydantic_ai._agent_graph",
+                "UserPromptNode.run",
+                user_prompt_processing(
+                    version,
+                    environment,
+                    application_name,
+                    tracer,
+                    pricing_info,
+                    capture_message_content,
+                    metrics,
+                    disable_metrics,
+                ),
+            )
+        except Exception:
+            # If UserPromptNode.run doesn't exist, skip this instrumentation
+            pass
+
+        try:
+            # ModelRequestNode.run() - Model request processing
+            wrap_function_wrapper(
+                "pydantic_ai._agent_graph",
+                "ModelRequestNode.run",
+                model_request_processing(
+                    version,
+                    environment,
+                    application_name,
+                    tracer,
+                    pricing_info,
+                    capture_message_content,
+                    metrics,
+                    disable_metrics,
+                ),
+            )
+        except Exception:
+            # If ModelRequestNode.run doesn't exist, skip this instrumentation
+            pass
+
+        try:
+            # CallToolsNode.run() - Tool calls processing
+            wrap_function_wrapper(
+                "pydantic_ai._agent_graph",
+                "CallToolsNode.run",
+                tool_calls_processing(
+                    version,
+                    environment,
+                    application_name,
+                    tracer,
+                    pricing_info,
+                    capture_message_content,
+                    metrics,
+                    disable_metrics,
+                ),
+            )
+        except Exception:
+            # If CallToolsNode.run doesn't exist, skip this instrumentation
+            pass
+
     def _uninstrument(self, **kwargs):
         # Proper uninstrumentation logic to revert patched methods
         pass

openlit/instrumentation/pydantic_ai/async_pydantic_ai.py

@@ -0,0 +1,38 @@
+"""
+Module for monitoring async Pydantic AI API calls.
+"""
+
+from openlit.instrumentation.pydantic_ai.utils import (
+    common_agent_run_async,
+)
+
+
+def async_agent_run(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
+    """
+    Generates a telemetry wrapper for async GenAI function call
+    """
+
+    async def wrapper(wrapped, instance, args, kwargs):
+        return await common_agent_run_async(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            tracer,
+            version,
+            environment,
+            application_name,
+            capture_message_content,
+            pricing_info=pricing_info,
+        )
+
+    return wrapper

openlit/instrumentation/pydantic_ai/pydantic_ai.py

@@ -5,6 +5,10 @@ Module for monitoring Pydantic AI API calls.
 from openlit.instrumentation.pydantic_ai.utils import (
     common_agent_run,
     common_agent_create,
+    common_graph_execution,
+    common_user_prompt_processing,
+    common_model_request_processing,
+    common_tool_calls_processing,
 )
 
 
@@ -23,7 +27,6 @@ def agent_create(
     """
 
     def wrapper(wrapped, instance, args, kwargs):
-        response = wrapped(*args, **kwargs)
         return common_agent_create(
             wrapped,
             instance,
@@ -34,7 +37,6 @@ def agent_create(
             environment,
             application_name,
             capture_message_content,
-            response=response,
         )
 
     return wrapper
@@ -55,7 +57,6 @@ def agent_run(
     """
 
     def wrapper(wrapped, instance, args, kwargs):
-        response = wrapped(*args, **kwargs)
         return common_agent_run(
             wrapped,
             instance,
@@ -66,13 +67,13 @@ def agent_run(
             environment,
             application_name,
             capture_message_content,
-            response=response,
+            pricing_info=pricing_info,
         )
 
     return wrapper
 
 
-def async_agent_run(
+def graph_execution(
     version,
     environment,
     application_name,
@@ -83,12 +84,101 @@ def async_agent_run(
     disable_metrics,
 ):
     """
-    Generates a telemetry wrapper for GenAI function call
+    Generates a telemetry wrapper for Pydantic AI graph execution
     """
 
-    async def wrapper(wrapped, instance, args, kwargs):
-        response = await wrapped(*args, **kwargs)
-        return common_agent_run(
+    def wrapper(wrapped, instance, args, kwargs):
+        return common_graph_execution(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            tracer,
+            version,
+            environment,
+            application_name,
+            capture_message_content,
+        )
+
+    return wrapper
+
+
+def user_prompt_processing(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
+    """
+    Generates a telemetry wrapper for Pydantic AI user prompt processing
+    """
+
+    def wrapper(wrapped, instance, args, kwargs):
+        return common_user_prompt_processing(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            tracer,
+            version,
+            environment,
+            application_name,
+            capture_message_content,
+        )
+
+    return wrapper
+
+
+def model_request_processing(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
+    """
+    Generates a telemetry wrapper for Pydantic AI model request processing
+    """
+
+    def wrapper(wrapped, instance, args, kwargs):
+        return common_model_request_processing(
+            wrapped,
+            instance,
+            args,
+            kwargs,
+            tracer,
+            version,
+            environment,
+            application_name,
+            capture_message_content,
+        )
+
+    return wrapper
+
+
+def tool_calls_processing(
+    version,
+    environment,
+    application_name,
+    tracer,
+    pricing_info,
+    capture_message_content,
+    metrics,
+    disable_metrics,
+):
+    """
+    Generates a telemetry wrapper for Pydantic AI tool calls processing
+    """
+
+    def wrapper(wrapped, instance, args, kwargs):
+        return common_tool_calls_processing(
             wrapped,
             instance,
             args,
@@ -98,7 +188,6 @@ def async_agent_run(
             environment,
             application_name,
             capture_message_content,
-            response=response,
         )
 
     return wrapper