import asyncio from openai import AsyncOpenAI from collections import defaultdict import weave from pydantic import BaseModel from abc import ABC, abstractmethod import json from typing import Dict, List from datetime import datetime import backoff from openai import APITimeoutError, APIError, RateLimitError class FailureCategory(BaseModel): category_id: int category_name: str description: str class FailureCategories(BaseModel): failure_categories: list[FailureCategory] class TaskSummary(BaseModel): task_id: str summary: str class TaskClassification(BaseModel): task_id: str category_id: str category_name: str explanation: str class OverallAnalysis(BaseModel): failure_categories: List[Dict] task_classifications: Dict[str, Dict] summary: str class AsyncLLMClient(ABC): @abstractmethod async def generate_text(self, prompt, system_message=None, response_format=None): pass # class AsyncOpenAIClient(AsyncLLMClient): # def __init__(self, model="gpt-4o-mini"): # self.model = model # self.client = AsyncOpenAI() # async def generate_text(self, prompt, system_message=None, response_format=None): # messages = [ # {"role": "system", "content": system_message or "You are a helpful AI assistant."}, # {"role": "user", "content": prompt} # ] # if response_format: # response = await self.client.beta.chat.completions.parse(model=self.model, messages=messages, response_format=response_format) # else: # response = await self.client.chat.completions.create(model=self.model, messages=messages) # return response.choices[0].message.content class AsyncOpenAIClient(AsyncLLMClient): def __init__(self, model="gpt-4o-mini", max_tries=5, max_time=300): self.model = model self.client = AsyncOpenAI() self.max_tries = max_tries self.max_time = max_time @backoff.on_exception( backoff.expo, (APITimeoutError, APIError, RateLimitError), max_tries=10, max_time=300 ) async def _make_request(self, messages, response_format=None): if response_format: return await self.client.beta.chat.completions.parse( model=self.model, messages=messages, response_format=response_format ) else: return await self.client.chat.completions.create( model=self.model, messages=messages ) async def generate_text(self, prompt, system_message=None, response_format=None): messages = [ {"role": "system", "content": system_message or "You are a helpful AI assistant."}, {"role": "user", "content": prompt} ] try: response = await self._make_request(messages, response_format) return response.choices[0].message.content except Exception as e: raise Exception(f"Failed after {self.max_tries} attempts or {self.max_time} seconds: {str(e)}") def get_weave_calls(client): calls = client.calls() processed_calls = [] for call in calls: ChatCompletion = weave.ref(call.output).get() choices = [choice.message.content for choice in ChatCompletion.choices] output = { 'weave_task_id': call.attributes['weave_task_id'], 'trace_id': call.trace_id, 'project_id': call.project_id, 'created_timestamp': ChatCompletion.created, 'inputs': dict(call.inputs), 'id': call.id, 'outputs': {'choices' : choices}, 'exception': call.exception, 'summary': call.summary, 'display_name': call.display_name, 'attributes': dict(call.attributes), "_children": call._children, '_feedback': call._feedback, } processed_calls.append(output) return processed_calls async def analyze_agent_performance(processed_calls, failed_tasks: list, llm_client): task_calls = defaultdict(list) for call in processed_calls: if call['weave_task_id'] in failed_tasks: task_calls[call['weave_task_id']].append(call) for task_id in task_calls: task_calls[task_id].sort(key=lambda x: x['created_timestamp']) task_summaries = await asyncio.gather(*[summarize_task(task_id, calls, llm_client) for task_id, calls in task_calls.items()]) failure_categories = await identify_failure_categories(task_summaries, llm_client) task_classifications = await classify_tasks(task_summaries, failure_categories, llm_client) overall_summary = await generate_overall_summary(failure_categories, task_classifications, llm_client) task_classifications = {tc["task_id"]: tc for tc in task_classifications} return dict(OverallAnalysis( failure_categories=failure_categories, task_classifications=task_classifications, summary=overall_summary )) async def summarize_task(task_id, calls, llm_client): calls_summary = "" for i, call in enumerate(calls, 1): calls_summary += f""" Step {i}: Input: {call['inputs']} Output: {call['outputs']} Timestamp: {datetime.fromtimestamp(call['created_timestamp'])} """ prompt = f""" Summarize the AI agent's performance on the following task: Task ID: {task_id} Number of steps: {len(calls)} Detailed steps: {calls_summary} Provide a brief summary of: 1. The main goal of the task (inferred from the inputs and outputs) 2. The agent's approach, including key steps and decisions made 3. Any significant challenges or errors encountered during the task 4. The final outcome why the task failed. Be detailed about the reason for failure. Keep the summary concise (around 200 words) but include specific details about the agent's performance and any notable aspects of its problem-solving process. """ system_message = "You are an AI performance analyst tasked with summarizing an AI agent's performance on individual tasks. Focus on the most important aspects of the agent's approach and performance." summary = await llm_client.generate_text(prompt, system_message, response_format=TaskSummary) return json.loads(summary) async def identify_failure_categories(task_summaries, llm_client): summaries_text = "\n\n".join([f"Task {s['task_id']}:\n{s['summary']}" for s in task_summaries]) prompt = f""" Analyze the following summaries of an AI agent's performance across multiple tasks: {summaries_text} Identify recurring categories of failures that the agent faces across these tasks. For each category: 1. Provide a short, descriptive name (max 5 words) 2. Write a brief description explaining the nature of this failure or challenge category Focus on patterns that appear across multiple tasks and represent specific errors that impacted the agent's performance. Make sure that your categories are distinct and cover a range of recurring issues. The categories should not bee too general. Examples for categories could include: Incorrect Implementation - The agent made a change to a reasonable area but their solution didn’t correctly address the issue. Gave Up Prematurely - The agent decides to stop solving the task after encountering some difficulty. Failed Edit Recovery - The agent went into an loop, making recurrent failing edits without recovering. """ system_message = "You are an expert in AI agent analysis, tasked with identifying recurring patterns in agent performance across multiple tasks." categories = await llm_client.generate_text(prompt, system_message, response_format=FailureCategories) return [dict(category) for category in json.loads(categories)['failure_categories']] async def classify_tasks(task_summaries, failure_categories, llm_client): categories_text = "\n".join([f"{cat['category_id']}. {cat['category_name']}: {cat['description']}" for i, cat in enumerate(failure_categories)]) classifications = [] for task in task_summaries: prompt = f""" Failure Categories: {categories_text} Task Summary: {task['summary']} Classify this task into one of the failure categories listed above. Provide: 1. The number of the chosen category 2. A brief explanation of why this category best fits the task's outcome If the task doesn't clearly fit any category, you may classify it as "0. Other" and explain why. """ system_message = "You are an AI performance analyst tasked with classifying task outcomes into predefined categories." classification = await llm_client.generate_text(prompt, system_message, response_format=TaskClassification) classification = json.loads(classification) category_number = classification['category_id'] if str(category_number) == "0": category_name = "Other" else: for cat in failure_categories: if str(cat['category_id']) == str(category_number): category_name = cat['category_name'] break else: category_name = "Other" explanation = classification['explanation'] classifications.append(dict(TaskClassification( task_id=task['task_id'], category_id=category_number, category_name=category_name, explanation=explanation ))) return classifications async def generate_overall_summary(failure_categories, task_classifications, llm_client): categories_text = "\n".join([f"{cat['category_name']}: {cat['description']}" for cat in failure_categories]) classifications_text = "\n".join([f"Task {tc['task_id']}: {tc['category_name']}" for tc in task_classifications]) prompt = f""" Failure Categories: {categories_text} Task Classifications: {classifications_text} Based on the failure categories identified and the classification of tasks, provide an overall summary of the AI agent's performance across all tasks. Include: 1. The most common types of failures or challenges 2. Any patterns in the agent's performance across different tasks 3. Suggestions for areas of improvement in the agent's design or training Keep the summary concise but insightful, focusing on the most significant findings and their implications for AI agent development. Do only return the summary itself without any preceding context etc. """ system_message = "You are a senior AI researcher tasked with providing a high-level analysis of an AI agent's performance across multiple tasks." return await llm_client.generate_text(prompt, system_message) async def main(): client = weave.init("citp_agent_eval/usaco_1723148990") processed_calls = get_weave_calls(client) weave.finish() openai_client = AsyncOpenAIClient(model="gpt-4o-mini") overall_analysis = await analyze_agent_performance(processed_calls, openai_client) with open("agent_performance_analysis.json", "w") as f: json.dump(overall_analysis.model_dump(), f, indent=4) if __name__ == "__main__": asyncio.run(main())