monocle-apptrace 0.5.0b1__py3-none-any.whl → 0.5.1__py3-none-any.whl

monocle_apptrace/instrumentation/metamodel/agents/entities/inference.py

@@ -0,0 +1,196 @@
+from monocle_apptrace.instrumentation.common.constants import AGENT_REQUEST_SPAN_NAME, SPAN_SUBTYPES, SPAN_TYPES
+from monocle_apptrace.instrumentation.metamodel.agents import _helper
+
+AGENT = {
+    "type": SPAN_TYPES.AGENTIC_INVOCATION,
+    "subtype": SPAN_SUBTYPES.ROUTING,
+    "attributes": [
+        [
+            {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments: _helper.AGENTS_AGENT_NAME_KEY,
+            },
+            {
+                "_comment": "name of the agent",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_agent_name(arguments),
+            },
+            {
+                "_comment": "agent description",
+                "attribute": "description",
+                "accessor": lambda arguments: _helper.get_agent_description(arguments),
+            },
+            {
+                "_comment": "agent instructions",
+                "attribute": "instructions",
+                "accessor": lambda arguments: _helper.get_agent_instructions(arguments),
+            },
+        ]
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "this is Agent input",
+                    "attribute": "query",
+                    "accessor": lambda arguments: _helper.extract_agent_input(
+                        arguments
+                    ),
+                }
+            ],
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "this is response from Agent",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_agent_response(
+                        arguments["result"]
+                    ),
+                }
+            ],
+        },
+        {
+            "name": "metadata",
+            "attributes": [
+                {
+                    "_comment": "this is metadata from Agent response",
+                    "accessor": lambda arguments: _helper.update_span_from_agent_response(
+                        arguments["result"]
+                    ),
+                }
+            ],
+        },
+    ],
+}
+
+AGENT_REQUEST = {
+    "type": AGENT_REQUEST_SPAN_NAME,
+    "subtype": SPAN_SUBTYPES.PLANNING,
+    "attributes": [
+        [
+            {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments: _helper.AGENTS_AGENT_NAME_KEY,
+            }
+        ],
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "this is Agent input",
+                    "attribute": "input",
+                    "accessor": lambda arguments: _helper.extract_agent_input(
+                        arguments
+                    ),
+                }
+            ],
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "this is response from Agent",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_agent_response(
+                        arguments["result"]
+                    ),
+                }
+            ],
+        },
+    ],
+}
+
+TOOLS = {
+    "type": SPAN_TYPES.AGENTIC_TOOL_INVOCATION,
+    "subtype": SPAN_SUBTYPES.ROUTING,
+    "attributes": [
+        [
+            {
+                "_comment": "tool type",
+                "attribute": "type",
+                "accessor": lambda arguments: _helper.get_tool_type(arguments["span"]),
+            },
+            {
+                "_comment": "name of the tool",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_tool_name(
+                    arguments["instance"]
+                ),
+            },
+            {
+                "_comment": "tool description",
+                "attribute": "description",
+                "accessor": lambda arguments: _helper.get_tool_description(
+                    arguments["instance"]
+                ),
+            },
+        ],
+        [
+            {
+                "_comment": "name of the agent",
+                "attribute": "name",
+                "accessor": lambda arguments: _helper.get_source_agent(),
+            },
+            {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments: _helper.AGENTS_AGENT_NAME_KEY,
+            },
+        ],
+    ],
+    "events": [
+        {
+            "name": "data.input",
+            "attributes": [
+                {
+                    "_comment": "this is Tool input",
+                    "attribute": "Inputs",
+                    "accessor": lambda arguments: _helper.extract_tool_input(arguments),
+                },
+            ],
+        },
+        {
+            "name": "data.output",
+            "attributes": [
+                {
+                    "_comment": "this is response from Tool",
+                    "attribute": "response",
+                    "accessor": lambda arguments: _helper.extract_tool_response(
+                        arguments["result"]
+                    ),
+                }
+            ],
+        },
+    ],
+}
+
+AGENT_DELEGATION = {
+    "type": SPAN_TYPES.AGENTIC_DELEGATION,
+    "subtype": SPAN_SUBTYPES.ROUTING,
+    "attributes": [
+        [
+            {
+                "_comment": "agent type",
+                "attribute": "type",
+                "accessor": lambda arguments: _helper.AGENTS_AGENT_NAME_KEY,
+            },
+            {
+                "_comment": "name of the source agent",
+                "attribute": "from_agent",
+                "accessor": lambda arguments: _helper.get_source_agent(),
+            },
+            {
+                "_comment": "name of the target agent",
+                "attribute": "to_agent",
+                "accessor": lambda arguments: _helper.extract_handoff_target(arguments),
+            },
+        ]
+    ],
+}

monocle_apptrace/instrumentation/metamodel/agents/methods.py

@@ -0,0 +1,55 @@
+from monocle_apptrace.instrumentation.common.wrapper import task_wrapper, atask_wrapper
+from monocle_apptrace.instrumentation.metamodel.agents.entities.inference import (
+    AGENT,
+    AGENT_DELEGATION,
+    TOOLS,
+    AGENT_REQUEST
+)
+from monocle_apptrace.instrumentation.metamodel.agents.agents_processor import (
+    constructor_wrapper,
+    handoff_constructor_wrapper,
+)
+
+AGENTS_METHODS = [
+    # Main agent runner methods
+    {
+        "package": "agents.run",
+        "object": "Runner",
+        "method": "run",
+        "wrapper_method": atask_wrapper,
+        "span_handler": "agents_agent_handler",
+        "output_processor": AGENT_REQUEST,
+    },
+    {
+        "package": "agents.run",
+        "object": "Runner",
+        "method": "run_sync",
+        "wrapper_method": task_wrapper,
+        "span_handler": "agents_agent_handler",
+        "output_processor": AGENT_REQUEST,
+    },
+    # AgentRunner class methods (internal runner)
+    {
+        "package": "agents.run",
+        "object": "AgentRunner",
+        "method": "_run_single_turn",
+        "wrapper_method": atask_wrapper,
+        "span_handler": "agents_agent_handler",
+        "output_processor": AGENT,
+    },
+    # Function tool decorator - wrap the function_tool function directly
+    {
+        "package": "agents.tool",
+        "object": "FunctionTool",
+        "method": "__init__",  # Empty string means wrap the function itself
+        "wrapper_method": constructor_wrapper,
+        "output_processor": TOOLS,
+    },
+    {
+        "package": "agents.handoffs",
+        "object": "Handoff",
+        "method": "__init__",  # Empty string means wrap the function itself
+        "wrapper_method": handoff_constructor_wrapper,
+        "output_processor": AGENT_DELEGATION,
+    },
+]

monocle_apptrace/instrumentation/metamodel/aiohttp/entities/http.py

@@ -1,6 +1,7 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
 from monocle_apptrace.instrumentation.metamodel.aiohttp import _helper
 AIO_HTTP_PROCESSOR = {
-    "type": "http.process",
+    "type": SPAN_TYPES.HTTP_PROCESS,
     "attributes": [
         [
             {

monocle_apptrace/instrumentation/metamodel/anthropic/_helper.py

@@ -5,6 +5,7 @@ and assistant messages from various input formats.
 
 import json
 import logging
+from opentelemetry.context import get_value
 from monocle_apptrace.instrumentation.common.utils import (
     Option,
     get_json_dumps,
@@ -15,6 +16,7 @@ from monocle_apptrace.instrumentation.common.utils import (
     get_exception_message,
 )
 from monocle_apptrace.instrumentation.metamodel.finish_types import map_anthropic_finish_reason_to_finish_type
+from monocle_apptrace.instrumentation.common.constants import AGENT_PREFIX_KEY, INFERENCE_AGENT_DELEGATION, INFERENCE_TURN_END, INFERENCE_TOOL_CALL
 
 
 logger = logging.getLogger(__name__)
@@ -60,10 +62,41 @@ def extract_assistant_message(arguments):
         if status == 'success':
             messages = []
             role = response.role if hasattr(response, 'role') else "assistant"
-            if response is not None and hasattr(response,"content") and len(response.content) >0:
-                if hasattr(response.content[0],"text"):
-                    messages.append({role: response.content[0].text})
-            # return first message if list is not empty
+            
+            # Handle tool use content blocks
+            if hasattr(response, "content") and response.content:
+                tools = []
+                text_content = []
+                
+                for content_block in response.content:
+                    if hasattr(content_block, "type"):
+                        if content_block.type == "tool_use":
+                            # Extract tool use information
+                            tool_info = {
+                                "tool_id": getattr(content_block, "id", ""),
+                                "tool_name": getattr(content_block, "name", ""),
+                                "tool_arguments": getattr(content_block, "input", "")
+                            }
+                            tools.append(tool_info)
+                        elif content_block.type == "text":
+                            # Extract text content
+                            if hasattr(content_block, "text"):
+                                text_content.append(content_block.text)
+                
+                # If we have tools, add them to the message
+                if tools:
+                    messages.append({"tools": tools})
+                
+                # If we have text content, add it to the message
+                if text_content:
+                    messages.append({role: " ".join(text_content)})
+                
+                # Fallback to original logic if no content blocks were processed
+                if not messages and len(response.content) > 0:
+                    if hasattr(response.content[0], "text"):
+                        messages.append({role: response.content[0].text})
+            
+            # Return first message if list is not empty
             return get_json_dumps(messages[0]) if messages else ""
         else:
             if arguments["exception"] is not None:
@@ -103,4 +136,48 @@ def extract_finish_reason(arguments):
 
 def map_finish_reason_to_finish_type(finish_reason):
     """Map Anthropic stop_reason to finish_type, similar to OpenAI mapping."""
-    return map_anthropic_finish_reason_to_finish_type(finish_reason)
+    return map_anthropic_finish_reason_to_finish_type(finish_reason)
+
+def agent_inference_type(arguments):
+    """Extract agent inference type from Anthropic response"""
+    try:
+        status = get_status_code(arguments)
+        if status == 'success' or status == 'completed':
+            response = arguments["result"]
+            
+            # Check if stop_reason indicates tool use
+            if hasattr(response, "stop_reason") and response.stop_reason == "tool_use":
+                # Check if this is agent delegation by looking at tool names
+                if hasattr(response, "content") and response.content:
+                    agent_prefix = get_value(AGENT_PREFIX_KEY)
+                    for content_block in response.content:
+                        if (hasattr(content_block, "type") and 
+                            content_block.type == "tool_use" and
+                            hasattr(content_block, "name")):
+                            tool_name = content_block.name
+                            if agent_prefix and tool_name.startswith(agent_prefix):
+                                return INFERENCE_AGENT_DELEGATION
+                    # If we found tool use but no agent delegation, it's a regular tool call
+                    return INFERENCE_TOOL_CALL
+            
+            # Fallback: check the extracted message for tool content
+            assistant_message = extract_assistant_message(arguments)
+            if assistant_message:
+                try:
+                    message = json.loads(assistant_message)
+                    if message and isinstance(message, dict):
+                        assistant_content = message.get("assistant", "")
+                        if assistant_content:
+                            agent_prefix = get_value(AGENT_PREFIX_KEY)
+                            if agent_prefix and agent_prefix in assistant_content:
+                                return INFERENCE_AGENT_DELEGATION
+                except (json.JSONDecodeError, TypeError):
+                    # If JSON parsing fails, fall back to string analysis
+                    agent_prefix = get_value(AGENT_PREFIX_KEY)
+                    if agent_prefix and agent_prefix in assistant_message:
+                        return INFERENCE_AGENT_DELEGATION
+        
+        return INFERENCE_TURN_END
+    except Exception as e:
+        logger.warning("Warning: Error occurred in agent_inference_type: %s", str(e))
+        return INFERENCE_TURN_END

monocle_apptrace/instrumentation/metamodel/anthropic/entities/inference.py

@@ -1,10 +1,11 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
 from monocle_apptrace.instrumentation.metamodel.anthropic import (
     _helper,
 )
 from monocle_apptrace.instrumentation.common.utils import (get_error_message, resolve_from_alias)
 
 INFERENCE = {
-    "type": "inference",
+    "type": SPAN_TYPES.INFERENCE,
     "attributes": [
         [
             {
@@ -79,6 +80,10 @@ INFERENCE = {
                     "_comment": "finish type mapped from finish reason",
                     "attribute": "finish_type",
                     "accessor": lambda arguments: _helper.map_finish_reason_to_finish_type(_helper.extract_finish_reason(arguments))
+                },
+                {
+                    "attribute": "inference_sub_type",
+                    "accessor": lambda arguments: _helper.agent_inference_type(arguments)
                 }
             ]
         }

monocle_apptrace/instrumentation/metamodel/azfunc/entities/http.py

@@ -1,6 +1,7 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
 from monocle_apptrace.instrumentation.metamodel.azfunc import _helper
 AZFUNC_HTTP_PROCESSOR = {
-    "type": "http.process",
+    "type": SPAN_TYPES.HTTP_PROCESS,
     "attributes": [
         [
             {

monocle_apptrace/instrumentation/metamodel/azureaiinference/entities/inference.py

@@ -1,6 +1,7 @@
 import logging
 import time
 from types import SimpleNamespace
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
 from monocle_apptrace.instrumentation.metamodel.azureaiinference import _helper
 from monocle_apptrace.instrumentation.common.utils import (
     get_error_message,
@@ -133,7 +134,7 @@ def process_stream(to_wrap, response, span_processor):
 
 
 INFERENCE = {
-    "type": "inference",
+    "type": SPAN_TYPES.INFERENCE,
     "is_auto_close": lambda kwargs: kwargs.get("stream", False) is False,
     "response_processor": process_stream,
     "attributes": [

monocle_apptrace/instrumentation/metamodel/botocore/entities/inference.py

@@ -1,9 +1,10 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
 from monocle_apptrace.instrumentation.metamodel.botocore import (
     _helper,
 )
 from monocle_apptrace.instrumentation.common.utils import (get_error_message, get_llm_type, get_status,)
 INFERENCE = {
-    "type": "inference",
+    "type": SPAN_TYPES.INFERENCE,
     "attributes": [
         [
             {

monocle_apptrace/instrumentation/metamodel/fastapi/entities/http.py

@@ -1,7 +1,8 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
 from monocle_apptrace.instrumentation.metamodel.fastapi import _helper
 
 FASTAPI_HTTP_PROCESSOR = {
-    "type": "http.process",
+    "type": SPAN_TYPES.HTTP_PROCESS,
     "attributes": [
         [
             {

monocle_apptrace/instrumentation/metamodel/fastapi/methods.py

@@ -2,22 +2,22 @@ from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper
 from monocle_apptrace.instrumentation.metamodel.fastapi.entities.http import FASTAPI_HTTP_PROCESSOR, FASTAPI_RESPONSE_PROCESSOR
 
 FASTAPI_METHODS = [
-    {
-        "package": "fastapi",
-        "object": "FastAPI",
-        "method": "__call__",
-        "wrapper_method": atask_wrapper,
-        "span_name": "fastapi.request",
-        "span_handler": "fastapi_handler",
-        "output_processor": FASTAPI_HTTP_PROCESSOR,
-    },
-    {
-        "package": "starlette.responses",
-        "object": "Response",
-        "method": "__call__",
-        "span_name": "fastapi.response",
-        "wrapper_method": atask_wrapper,
-        "span_handler": "fastapi_response_handler",
-        "output_processor": FASTAPI_RESPONSE_PROCESSOR
-    }
+    # {
+    #     "package": "fastapi",
+    #     "object": "FastAPI",
+    #     "method": "__call__",
+    #     "wrapper_method": atask_wrapper,
+    #     "span_name": "fastapi.request",
+    #     "span_handler": "fastapi_handler",
+    #     "output_processor": FASTAPI_HTTP_PROCESSOR,
+    # },
+    # {
+    #     "package": "starlette.responses",
+    #     "object": "Response",
+    #     "method": "__call__",
+    #     "span_name": "fastapi.response",
+    #     "wrapper_method": atask_wrapper,
+    #     "span_handler": "fastapi_response_handler",
+    #     "output_processor": FASTAPI_RESPONSE_PROCESSOR
+    # }
 ]

monocle_apptrace/instrumentation/metamodel/finish_types.py

@@ -11,6 +11,7 @@ class FinishType(Enum):
     TRUNCATED = "truncated"
     CONTENT_FILTER = "content_filter"
     ERROR = "error"
+    TOOL_CALL_ERROR = "tool_call_error"
     REFUSAL = "refusal"
     RATE_LIMITED = "rate_limited"
 
@@ -39,6 +40,7 @@ GEMINI_FINISH_REASON_MAPPING = {
     "MAX_TOKENS": FinishType.TRUNCATED.value,
     "SAFETY": FinishType.CONTENT_FILTER.value,
     "RECITATION": FinishType.CONTENT_FILTER.value,
+    "MALFORMED_FUNCTION_CALL": FinishType.TOOL_CALL_ERROR.value,
     "OTHER": FinishType.ERROR.value,
     "FINISH_REASON_UNSPECIFIED": None
 }
@@ -231,6 +233,48 @@ TEAMSAI_FINISH_REASON_MAPPING = {
     "rate_limited": FinishType.RATE_LIMITED.value,
     "invalid_response": FinishType.ERROR.value,
 }
+# Haystack finish reason mapping
+HAYSTACK_FINISH_REASON_MAPPING = {
+    # Standard completion reasons
+    "stop": FinishType.SUCCESS.value,
+    "complete": FinishType.SUCCESS.value,
+    "finished": FinishType.SUCCESS.value,
+
+    # Token limits
+    "length": FinishType.TRUNCATED.value,
+    "max_tokens": FinishType.TRUNCATED.value,
+    "token_limit": FinishType.TRUNCATED.value,
+
+    # Tool/function calling
+    "tool_calls": FinishType.SUCCESS.value,
+    "function_call": FinishType.SUCCESS.value,
+
+    # Content filtering and safety
+    "content_filter": FinishType.CONTENT_FILTER.value,
+    "safety": FinishType.CONTENT_FILTER.value,
+    "filtered": FinishType.CONTENT_FILTER.value,
+
+    # Errors
+    "error": FinishType.ERROR.value,
+    "failed": FinishType.ERROR.value,
+    "exception": FinishType.ERROR.value,
+
+    # Provider-specific reasons that might pass through LangChain
+    # OpenAI reasons
+    "stop": FinishType.SUCCESS.value,  # Already defined above
+
+    # Anthropic reasons
+    "end_turn": FinishType.SUCCESS.value,
+    "stop_sequence": FinishType.SUCCESS.value,
+
+    # Gemini reasons
+    "STOP": FinishType.SUCCESS.value,
+    "SAFETY": FinishType.CONTENT_FILTER.value,
+    "RECITATION": FinishType.CONTENT_FILTER.value,
+    "OTHER": FinishType.ERROR.value,
+}
+
+ADK_FINISH_REASON_MAPPING = GEMINI_FINISH_REASON_MAPPING
 
 def map_openai_finish_reason_to_finish_type(finish_reason):
     """Map OpenAI finish_reason to standardized finish_type."""
@@ -368,6 +412,34 @@ def map_bedrock_finish_reason_to_finish_type(finish_reason):
     
     return None
 
+def map_haystack_finish_reason_to_finish_type(finish_reason):
+    """Map Haystack finish_reason to standardized finish_type."""
+    if not finish_reason:
+        return None
+
+    # Convert to lowercase for case-insensitive matching
+    finish_reason_lower = finish_reason.lower() if isinstance(finish_reason, str) else str(finish_reason).lower()
+
+    # Try direct mapping first
+    if finish_reason in HAYSTACK_FINISH_REASON_MAPPING:
+        return HAYSTACK_FINISH_REASON_MAPPING[finish_reason]
+
+    # Try lowercase mapping
+    if finish_reason_lower in HAYSTACK_FINISH_REASON_MAPPING:
+        return HAYSTACK_FINISH_REASON_MAPPING[finish_reason_lower]
+
+    # If no direct mapping, try to infer from common patterns
+    if any(keyword in finish_reason_lower for keyword in ['stop', 'complete', 'success', 'done']):
+        return FinishType.SUCCESS.value
+    elif any(keyword in finish_reason_lower for keyword in ['length', 'token', 'limit', 'truncat']):
+        return FinishType.TRUNCATED.value
+    elif any(keyword in finish_reason_lower for keyword in ['filter', 'safety', 'block']):
+        return FinishType.CONTENT_FILTER.value
+    elif any(keyword in finish_reason_lower for keyword in ['error', 'fail', 'exception']):
+        return FinishType.ERROR.value
+
+    return None
+
 def map_teamsai_finish_reason_to_finish_type(finish_reason):
     """Map TeamsAI finish_reason to standardized finish_type."""
     if not finish_reason:
@@ -384,4 +456,10 @@ def map_teamsai_finish_reason_to_finish_type(finish_reason):
     if finish_reason_lower in TEAMSAI_FINISH_REASON_MAPPING:
         return TEAMSAI_FINISH_REASON_MAPPING[finish_reason_lower]
 
-    return None
+    return None
+
+def map_adk_finish_reason_to_finish_type(finish_reason):
+    """Map ADK finish_reason to standardized finish_type."""
+    if not finish_reason:
+        return None
+    return ADK_FINISH_REASON_MAPPING.get(finish_reason, None)

monocle_apptrace/instrumentation/metamodel/flask/entities/http.py

@@ -1,6 +1,7 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
 from monocle_apptrace.instrumentation.metamodel.flask import _helper
 FLASK_HTTP_PROCESSOR = {
-    "type": "http.process",
+    "type": SPAN_TYPES.HTTP_PROCESS,
     "attributes": [
         [
             {

monocle_apptrace/instrumentation/metamodel/gemini/entities/inference.py

@@ -1,10 +1,11 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
 from monocle_apptrace.instrumentation.metamodel.gemini import (
     _helper,
 )
 from monocle_apptrace.instrumentation.common.utils import get_error_message
 
 INFERENCE = {
-    "type": "inference",
+    "type": SPAN_TYPES.INFERENCE,
     "attributes": [
         [
             {
@@ -15,8 +16,11 @@ INFERENCE = {
             {
                 "attribute": "inference_endpoint",
                 "accessor": lambda arguments: _helper.extract_inference_endpoint(arguments['instance'])
-            }
-        ],
+            },
+            {
+                "attribute": "provider_name",
+                "accessor": lambda arguments: 'gcp'
+            }        ],
         [
             {
                 "_comment": "LLM Model",

monocle_apptrace/instrumentation/metamodel/gemini/entities/retrieval.py

@@ -1,9 +1,10 @@
+from monocle_apptrace.instrumentation.common.constants import SPAN_TYPES
 from monocle_apptrace.instrumentation.metamodel.gemini import (
     _helper,
 )
 
 RETRIEVAL = {
-    "type": "retrieval",
+    "type": SPAN_TYPES.RETRIEVAL,
     "attributes": [
         [
             {

monocle_apptrace/instrumentation/metamodel/gemini/methods.py

@@ -1,4 +1,4 @@
-from monocle_apptrace.instrumentation.common.wrapper import task_wrapper
+from monocle_apptrace.instrumentation.common.wrapper import atask_wrapper, task_wrapper
 from monocle_apptrace.instrumentation.metamodel.gemini.entities.inference import (
     INFERENCE,
 )
@@ -14,6 +14,13 @@ GEMINI_METHODS = [
       "wrapper_method": task_wrapper,
       "output_processor": INFERENCE,
     },
+    {
+      "package": "google.genai.models",
+      "object": "AsyncModels",
+      "method": "generate_content",
+      "wrapper_method": atask_wrapper,
+      "output_processor": INFERENCE,
+    },
     {
       "package": "google.genai.models",
       "object": "Models",