[Bug]: Session not tracked for litellm + autogen-agentchat v4

[rapturt9]

Contact Details

ram@sitewiz.ai

📦 Package Version

0.3.24

🎞️ Framework Version

autogen-agentchat 0.4.2

🔎 Describe the Bug

Duration is tracked, but not cost or LLMs, eg

🖇 AgentOps: Session Stats - Duration: 24.6s | Cost: $0.00 | LLMs: 0 | Tools: 0 | Actions: 0 | Errors: 0

🤝 Contribution

• Yes, I'd be happy to submit a pull request with these changes.
• I need some guidance on how to contribute.
• I'd prefer the AgentOps team to handle this update.

[areibman]

Thanks for posting-- can you also share a code snippet? @rapturt9

[rapturt9]

@areibman I isolated the issue. The problem was actually the import of litellm that caused it. Eg the following code works, but if I remove the comment of litellm it stops working.

import os
import argparse
import json
import logging
from dotenv import load_dotenv
from pathlib import Path
from pydantic import BaseModel
from typing import Optional, Dict, Any
# import litellm
import asyncio


import os
import asyncio
import agentops
from dotenv import load_dotenv
from typing import Dict, Any, Optional
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.messages import TextMessage
from autogen_agentchat.ui import Console
from autogen_core import CancellationToken
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_agentchat.teams import MagenticOneGroupChat
from autogen_agentchat.conditions import MaxMessageTermination

load_dotenv()

def store_results(stream_key: str, timestamp: str, result: Dict[str, Any]):
    # Implement the logic to store the results
    pass

async def run_basic_agent(
    task: str,
    stream_key: str,
    options: Dict[str, Any],
    agent_instructions: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    Run a basic agent with agentops tracking.
    
    Args:
        task: The task to perform
        stream_key: Unique identifier for tracking
        options: Configuration options
        agent_instructions: Optional instructions for agent behavior
    """
    try:
        # Initialize agentops session
        session = agentops.init(
            default_tags=[stream_key, "basic_agent_test"],
            api_key=os.getenv("AGENTOPS_API_KEY")
        )

        # Configure the assistant
        assistant = AssistantAgent(
            name="assistant",
            model_client=OpenAIChatCompletionClient(
                model="gpt-4o",
                api_key=os.getenv("OPENAI_API_KEY"),
                temperature=0.7
            ),
            system_message="You are a helpful AI assistant."
        )

        # Configure group chat
        termination = MaxMessageTermination(100)
        magentic_one_group_chat = MagenticOneGroupChat(
            [assistant], 
            model_client=OpenAIChatCompletionClient(
                model="gpt-4o",
                api_key=os.getenv("OPENAI_API_KEY"),
                temperature=0.7
            ),
            termination_condition=termination
        )

        # Track the conversation with agentops
        # Run the chat
        stream = magentic_one_group_chat.run_stream(task=task)
        response = await Console(stream)

        session.end_session("Success")

        return {
            "status": "success",
            "message": "Agent conversation completed successfully",
            "response": response.messages[-1].content if response.messages else ""
        }

    except Exception as e:
        return {
            "status": "error",
            "message": f"Error running agent: {str(e)}"
        }

async def main():
    # Test the agent
    test_task = "Write a hello world program in Python"
    result = await run_basic_agent(
        task=test_task,
        stream_key="test_stream",
        options={"test": True},
        agent_instructions={"mode": "basic"}
    )
    print(f"Test result: {result}")



if __name__ == "__main__":
    asyncio.run(main())

Here is my requirements.txt:

aiofiles autogen-agentchat autogen-ext[web-surfer,magentic-one,openai] pydantic psycopg2-binary boto3 google-auth numpy pandas matplotlib scikit-learn requests urllib3 pillow pytest beautifulsoup4 lxml litellm agentops google-cloud-bigquery zep-cloud tavily-python

Using a docker with python:3.12-slim-bookworm