Enhance AG2 instrumentation (#1146)

* Enhance AG2 instrumentation with additional wrappers for async operations and improved context handling.
Refactor span attributes for consistency and clarity across agent interactions.

* Enhance AG2 Async Agent functionality and documentation

---------

Co-authored-by: Pratyush Shukla <[email protected]>
Co-authored-by: Alex Reibman <[email protected]>

Author: 3 people

agentops/instrumentation/__init__.py

@@ -332,7 +332,7 @@ def _perform_instrumentation(package_name: str):
             _has_agentic_library = True
 
         # Special case: If mem0 is instrumented, also instrument concurrent.futures
-        if package_name == "mem0" and is_newly_added:
+        if (package_name == "mem0" or package_name == "autogen") and is_newly_added:
             try:
                 # Check if concurrent.futures module is available
 

agentops/instrumentation/agentic/ag2/instrumentor.py

@@ -1,176 +1,185 @@
 ---
-title: 'AG2 with Mem0 Example'
-description: 'Observe an AG2 Agent with memory powered by Mem0 using AgentOps'
+title: 'AG2'
+description: 'AG2 Async Agent Chat'
 ---
-{/*  SOURCE_FILE: examples/ag2/agentchat_with_memory.ipynb  */}
+{/*  SOURCE_FILE: examples/ag2/async_human_input.ipynb  */}
 
-_View Notebook on <a href={'https://github.com/AgentOps-AI/agentops/blob/main/examples/ag2/agentchat_with_memory.ipynb'} target={'_blank'}>Github</a>_
+_View Notebook on <a href={'https://github.com/AgentOps-AI/agentops/blob/main/examples/ag2/ag2_async_agent.ipynb'} target={'_blank'}>Github</a>_
 
-# Observe an Agent with memory powered by Mem0
+# AG2 Async Agent Chat with Automated Responses
 
-This notebook demonstrates an intelligent customer service chatbot system that combines:
+This notebook demonstrates how to leverage asynchronous programming with AG2 agents 
+to create automated conversations between AI agents, eliminating the need for human 
+input while maintaining full traceability.
 
-- AG2 for conversational agents
-- Mem0 for memory management
+# Overview
+This notebook demonstrates a practical example of automated AI-to-AI communication where we:
 
-[Mem0](https://www.mem0.ai/) provides a smart, self-improving memory layer for Large Language Models (LLMs), enabling developers to create personalized AI experiences that evolve with each user interaction. Refer [docs](https://docs.mem0.ai/overview) for more information.
+1. Initialize AG2 agents with OpenAI's GPT-4o-mini model
+2. Create custom async agents that simulate human-like responses and processing delays
+3. Automate the entire conversation flow without requiring manual intervention
+4. Track all interactions using AgentOps for monitoring and analysis
 
-The implementation showcases how to initialize agents, manage conversation memory, and facilitate multi-agent conversations for enhanced problem-solving in customer support scenarios.
-
-With AgentOps, you can observe the agent's memory and interactions in real-time, providing insights into how the agent learns and adapts over time.
-
-## Pre-requisites
-- AgentOps API key from [AgentOps](https://app.agentops.ai/).
-- Mem0 API key from [Mem0 Platform](https://app.mem0.ai/).
-- OpenAI API key from [OpenAI](https://platform.openai.com/).
+By using async operations and automated responses, you can create fully autonomous 
+agent conversations that simulate real-world scenarios. This is particularly useful 
+for testing, prototyping, and creating demos where you want to showcase agent 
+capabilities without manual input.
 
 ## Installation
-
-Install required dependencies:
 <CodeGroup>
   ```bash pip
-  pip install agentops "ag2[openai]" mem0ai python-dotenv
+  pip install ag2 agentops nest-asyncio
   ```
   ```bash poetry
-  poetry add agentops ag2 mem0ai python-dotenv
-  # Note: For ag2[openai] with poetry, you might need to specify openai as an extra or directly.
-  # poetry add ag2 -E openai
+  poetry add ag2 agentops nest-asyncio
   ```
   ```bash uv
-  uv add agentops "ag2[openai]" mem0ai python-dotenv
+  uv add ag2 agentops nest-asyncio
   ```
 </CodeGroup>
 
-## Setup
-
-```python
+```
+import asyncio
+from typing import Dict, Optional, Union
 import os
 from dotenv import load_dotenv
+import nest_asyncio
 import agentops
-from mem0 import MemoryClient
-from autogen import ConversableAgent
+from autogen import AssistantAgent
+from autogen.agentchat.user_proxy_agent import UserProxyAgent
+```
+
 
+```
+# Load environment variables for API keys
 load_dotenv()
-os.environ["AGENTOPS_API_KEY"] = os.getenv("AGENTOPS_API_KEY", "your_agentops_api_key_here")
+os.environ["AGENTOPS_API_KEY"] = os.getenv("AGENTOPS_API_KEY", "your_api_key_here")
 os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your_openai_api_key_here")
-os.environ["MEM0_API_KEY"] = os.getenv("MEM0_API_KEY", "your_mem0_api_key_here")
+# Initialize AgentOps for tracking and monitoring
+agentops.init(auto_start_session=False, trace_name="AG2 Async Demo")
+tracer = agentops.start_trace(trace_name="AG2 Async Agent Demo", tags=["ag2-async-demo", "agentops-example"])
 ```
 
-## Initialize Agent and Memory
-
-The conversational agent is set up using the 'gpt-4o' model and a mem0 client. We'll utilize the client's methods for storing and accessing memories.
-
-```python
-agentops.init(auto_start_session=False)
-tracer = agentops.start_trace(trace_name="AG2 Agent using Mem0", tags=["ag2-mem0-example", "agentops-example"])
 
-agent = ConversableAgent(
-    "chatbot",
-    llm_config={"config_list": [{"model": "gpt-4o", "api_key": os.environ.get("OPENAI_API_KEY")}]},
-    code_execution_config=False,
-    function_map=None,
-    human_input_mode="NEVER",
-)
-
-memory = MemoryClient()
 ```
-
-Initialize a conversation history for a Best Buy customer service chatbot. It contains a list of message exchanges between the user and the assistant, structured as dictionaries with 'role' and 'content' keys. The entire conversation is then stored in memory using the `memory.add()` method, associated with the identifier "customer_service_bot".
-
-```python
-conversation = [
-    {
-        "role": "assistant",
-        "content": "Hi, I'm Best Buy's chatbot!\n\nThanks for being a My Best Buy TotalTM member.\n\nWhat can I help you with?",
-    },
-    {
-        "role": "user",
-        "content": 'Seeing horizontal lines on our tv. TV model: Sony - 77" Class BRAVIA XR A80K OLED 4K UHD Smart Google TV',
-    },
-    {
-        "role": "assistant",
-        "content": "Thanks for being a My Best Buy Total™ member. I can connect you to an expert immediately - just one perk of your membership!\n\nSelect the button below when you're ready to chat.",
-    },
-    {
-        "role": "assistant",
-        "content": "Good evening, thank you for choosing Best Buy, Fnu. My name is Lovely. I hope you are doing well. I'm sorry to hear that you're seeing horizontal lines on your TV.\n\nI'm absolutely committed to exploring all possible ways to assist you to fix this issue.\n\nTo ensure that we are on the right account, may I please have your email address registered with your Best Buy account?",
-    },
-    {"role": "user", "content": "[email protected]"},
-    {
-        "role": "assistant",
-        "content": "Perfect! Thank you for providing all the details, surely you have made my job easier by doing this. I really appreciate it.\n\nI also want to take a moment to express our heartfelt appreciation for your trust and loyalty. Thank you for being an amazing customer of BestBuy Total.\n\nCould you please help me with the order number or product's details to check it quickly?\n\nSamsung - 49\\" Odyssey OLED G9 (G95SC) DQHD 240Hz 0.03ms G-Sync Compatible Curved Smart Gaming Monitor - Silver - just to confirm this is the item, right?",
-    },
-    {"role": "user", "content": "Order number: 112217629"},
-    {
-        "role": "assistant",
-        "content": "Superb! Thank you for confirmation.\n\nThank you for your patience. After exploring all possible solutions, I can help you to arrange a home repair appointment for your device. Our Geek Squad experts will visit your home to inspect and fix your device.\n\nIt's great that you have a protection plan - rest assured, we've got your back! As a valued Total member, you can avail this service at a minimal service fee. This fee, applicable to all repairs, covers the cost of diagnosing the issue and any small parts needed for the repair. It's part of our 24-month free protection plan.\n\nPlease click here to review the service fee and plan coverage details -\n\nhttps://www.bestbuy.com/site/best-buy-membership/best-buy-protection/pcmcat1608643232014.c?id=pcmcat1608643232014#jl-servicefees\n\nFnu - just to confirm shall I proceed to schedule the appointment?",
-    },
-    {"role": "user", "content": "Yes please"},
-    {"role": "assistant", "content": "When should I schedule the appointment?"},
-    {"role": "user", "content": "Schedule it for tomorrow please"},
-]
-
-memory.add(messages=conversation, user_id="customer_service_bot")
+# Define an asynchronous function that simulates async processing 
+async def simulate_async_processing(task_name: str, delay: float = 1.0) -> str:
+    """
+    Simulate some asynchronous processing (e.g., API calls, file operations, etc.)
+    """
+    print(f"🔄 Starting async task: {task_name}")
+    await asyncio.sleep(delay)  # Simulate async work
+    print(f"✅ Completed async task: {task_name}")
+    return f"Processed: {task_name}"
 ```
 
-## Agent Inference
-
-We ask a question to the agent, utilizing mem0 to retrieve relevant memories. The agent then formulates a response based on both the question and the retrieved contextual information.
 
-```python
-data = "I forgot the order number, can you quickly tell me?"
-
-relevant_memories = memory.search(data, user_id="customer_service_bot")
-flatten_relevant_memories = "\n".join([m["memory"] for m in relevant_memories])
-
-prompt = f\"\"\"Answer the user question considering the memories. Keep answers clear and concise.
-Memories:
-{flatten_relevant_memories}
-\n\n
-Question: {data}
-\"\"\"
-
-reply = agent.generate_reply(messages=[{"content": prompt, "role": "user"}])
-print(reply)
+```
+# Define a custom UserProxyAgent that simulates automated user responses
+class AutomatedUserProxyAgent(UserProxyAgent):
+    def __init__(self, name: str, **kwargs):
+        super().__init__(name, **kwargs)
+        self.response_count = 0
+        self.predefined_responses = [
+            "Yes, please generate interview questions for these topics.",
+            "The questions look good. Can you make them more specific to senior-level positions?",
+            "Perfect! These questions are exactly what we need. Thank you!",
+        ]
+
+    async def a_get_human_input(self, prompt: str) -> str:
+        # Simulate async processing before responding
+        await simulate_async_processing(f"Processing user input #{self.response_count + 1}")
+
+        if self.response_count < len(self.predefined_responses):
+            response = self.predefined_responses[self.response_count]
+            self.response_count += 1
+            print(f"👤 User: {response}")
+            return response
+        else:
+            print("👤 User: TERMINATE")
+            return "TERMINATE"
+
+    async def a_receive(
+        self,
+        message: Union[Dict, str],
+        sender,
+        request_reply: Optional[bool] = None,
+        silent: Optional[bool] = False,
+    ):
+        await super().a_receive(message, sender, request_reply, silent)
 ```
 
-## Multi Agent Conversation
-
-Initialize two AI agents: a "manager" for resolving customer issues and a "customer_bot" for gathering information on customer problems, both using GPT-4. It then retrieves relevant memories for a given question, combining them with the question into a prompt. This prompt can be used by either the manager or customer_bot to generate a contextually informed response.
-
-```python
-manager = ConversableAgent(
-    "manager",
-    system_message="You are a manager who helps in resolving customer issues.",
-    llm_config={"config_list": [{"model": "gpt-4o-mini", "api_key": os.environ.get("OPENAI_API_KEY")}]},
-    human_input_mode="NEVER",
-)
-
-customer_bot = ConversableAgent(
-    "customer_bot",
-    system_message="You are a customer service bot who gathers information on issues customers are facing. Keep answers clear and concise.",
-    llm_config={"config_list": [{"model": "gpt-4", "api_key": os.environ.get("OPENAI_API_KEY")}]},
-    human_input_mode="NEVER",
-)
 
-data = "When is the appointment?"
+```
+# Define an AssistantAgent that simulates async processing before responding
+class AsyncAssistantAgent(AssistantAgent):
+    async def a_receive(
+        self,
+        message: Union[Dict, str],
+        sender,
+        request_reply: Optional[bool] = None,
+        silent: Optional[bool] = False,
+    ):
+        # Simulate async processing before responding
+        await simulate_async_processing("Analyzing request and preparing response", 0.5)
+        await super().a_receive(message, sender, request_reply, silent)
+```
 
-relevant_memories = memory.search(data, user_id="customer_service_bot")
-flatten_relevant_memories = "\n".join([m["memory"] for m in relevant_memories])
 
-prompt = f\"\"\"
-Context:
-{flatten_relevant_memories}
-\n\n
-Question: {data}
-\"\"\"
+```
+async def main():
+    print("🚀 Starting AG2 Async Demo")
+
+    # Create agents with automated behavior
+    user_proxy = AutomatedUserProxyAgent(
+        name="hiring_manager",
+        human_input_mode="NEVER",  # No human input required
+        max_consecutive_auto_reply=3,
+        code_execution_config=False,
+        is_termination_msg=lambda msg: "TERMINATE" in str(msg.get("content", "")),
+    )
+
+    assistant = AsyncAssistantAgent(
+        name="interview_consultant",
+        system_message="""You are an expert interview consultant. When given interview topics, 
+        you create thoughtful, relevant questions. You ask for feedback and incorporate it.
+        When the user is satisfied with the questions, end with 'TERMINATE'.""",
+        llm_config={"config_list": [{"model": "gpt-4o-mini", "api_key": os.environ.get("OPENAI_API_KEY")}]},
+        is_termination_msg=lambda msg: "TERMINATE" in str(msg.get("content", "")),
+    )
+
+    try:
+        print("🤖 Initiating automated conversation...")
+        # Start the automated chat between the user and assistant
+        await user_proxy.a_initiate_chat(
+            assistant,
+            message="""I need help creating interview questions for these topics:
+            - Resume Review
+            - Technical Skills Assessment  
+            - Project Discussion
+            - Job Role Expectations
+            - Closing Remarks
+            
+            Please create 2-3 questions for each topic.""",
+            max_turns=6,
+        )
+    except Exception as e:
+        print(f"\n❌ Error occurred: {e}")
+    finally:
+        agentops.end_trace(tracer, end_state="Success")
+
+    print("\n🎉 Demo completed successfully!")
+```
 
-result = manager.send(prompt, customer_bot, request_reply=True)
 
-agentops.end_trace(tracer, end_state="Success")
 ```
+# Run the main async demo
+nest_asyncio.apply()
+asyncio.run(main())
+```
+
 
 <script type="module" src="/scripts/github_stars.js"></script>
 <script type="module" src="/scripts/scroll-img-fadein-animation.js"></script>
 <script type="module" src="/scripts/button_heartbeat_animation.js"></script>
-<script type="module" src="/scripts/adjust_api_dynamically.js"></script>
+<script type="module" src="/scripts/adjust_api_dynamically.js"></script>