from typing import List, TypedDict, Dict, Any, Literal from langgraph.graph import StateGraph, START, END from langgraph.types import Command from langchain_core.messages import HumanMessage, AIMessage, ToolMessage from langchain_anthropic import ChatAnthropic from langchain_core.tools import tool from langchain_core.prompts import ChatPromptTemplate from langgraph.prebuilt import ToolNode import os from dotenv import load_dotenv from datetime import datetime from tavily import TavilyClient from langfuse.callback import CallbackHandler import requests import json import time from daytona_sdk import Daytona, DaytonaConfig import yt_dlp import io import os import tempfile from pathlib import Path # Load environment variablesTuple load_dotenv() # Define the state schema with messages that ToolNode can use class AgentState(TypedDict): messages: List current_question: str final_answer: str validation_result: str worker_iterations: int supervisor_satisfaction: bool validator_approval: bool # Define tools following Langgraph guide @tool def search_web_tavily(query: str) -> str: """Search the web for information using the Tavily search API.""" # Initialize the Tavily client with API key from environment variables client = TavilyClient(os.getenv("TAVILY_API_KEY")) # Perform the search response = client.search(query=query) # Process the results into a readable format results = [] for i, result in enumerate(response.get("results", []), 1): results.append(f"{i}. {result.get('title')}\n URL: {result.get('url')}\n {result.get('content')}\n") # Format the final response formatted_response = f"Search results for '{query}':\n\n" + "\n".join(results) return formatted_response @tool def search_web_serper(query: str, result_limit: int = 5, search_type: str = "search") -> str: """Search the web for information using the Serper.dev API. This tool provides comprehensive search results including: 1. Knowledge Graph data when available (title, description, attributes) 2. Organic search results (titles, links, snippets) 3. Related questions from "People Also Ask" section 4. Top stories and news articles related to the query It's particularly useful for gathering factual information, current events, and general knowledge from across the web. The results are formatted in a readable structure with clear sections. Parameters: - query: The search query string - result_limit: Maximum number of results to return per section (default: 5) - search_type: Type of search ('search', 'news', 'places', 'images', 'shopping') """ # API URL and headers setup url = "https://google.serper.dev/search" headers = { 'X-API-KEY': os.getenv("SERPER_API_KEY"), 'Content-Type': 'application/json' } # Prepare the payload with the query and search type payload = json.dumps({ "q": query, "type": search_type }) try: # Make the API request response = requests.request("POST", url, headers=headers, data=payload, timeout=30) response.raise_for_status() # Raise exception for HTTP errors # Parse the JSON response data = response.json() # Format the results results = [] # Add knowledge graph if available if "knowledgeGraph" in data: kg = data["knowledgeGraph"] results.append(f"Knowledge Graph:\n{kg.get('title', 'Unknown')} - {kg.get('type', '')}") results.append(f"Description: {kg.get('description', 'No description available')}") if "attributes" in kg: results.append("Attributes:") for key, value in kg["attributes"].items(): results.append(f"- {key}: {value}") results.append("") # Empty line for separation # Add organic search results if "organic" in data: results.append("Organic Search Results:") for i, result in enumerate(data["organic"][:result_limit], 1): results.append(f"{i}. {result.get('title', 'No title')}") results.append(f" URL: {result.get('link', 'No link')}") results.append(f" {result.get('snippet', 'No snippet')}") results.append("") # Empty line for separation # Add people also ask if available if "peopleAlsoAsk" in data and data["peopleAlsoAsk"]: results.append("People Also Ask:") for i, qa in enumerate(data["peopleAlsoAsk"][:min(3, result_limit)], 1): results.append(f"{i}. Q: {qa.get('question', 'No question')}") results.append(f" A: {qa.get('snippet', 'No answer')}") results.append("") # Empty line for separation # Add top stories if available if "topStories" in data and data["topStories"]: results.append("Top Stories:") for i, story in enumerate(data["topStories"][:min(3, result_limit)], 1): results.append(f"{i}. {story.get('title', 'No title')}") results.append(f" Source: {story.get('source', 'Unknown source')}") if "date" in story: results.append(f" Published: {story.get('date')}") results.append(f" URL: {story.get('link', 'No link')}") results.append("") # Empty line for separation # Format the final response formatted_response = f"Search results for '{query}':\n\n" + "\n".join(results) return formatted_response except requests.exceptions.Timeout: return f"Error: Request to Serper API timed out after 30 seconds" except requests.exceptions.RequestException as e: return f"Error making request to Serper API: {str(e)}" except json.JSONDecodeError: return f"Error: Received invalid JSON response from Serper API" except Exception as e: return f"Error processing search results: {str(e)}" # Initialize a global Daytona sandbox for reuse _daytona_sandbox = None @tool def execute_code_securely(code: str, language: str = "python", timeout: int = 300) -> str: """Execute code securely in an isolated sandbox environment using Daytona. This tool runs code in a secure, isolated environment to prevent security issues. It's particularly useful for solving computational problems, data processing tasks, mathematical calculations, and other scenarios where code execution is needed. The tool supports multiple languages, with Python as the default. Parameters: - code: The code to execute - language: The programming language (default: "python") - timeout: Maximum execution time in seconds (default: 30) Returns: - The execution result or error message """ global _daytona_sandbox try: # Initialize Daytona client if not already done if _daytona_sandbox is None: api_key = os.getenv("DAYTONA_API_KEY") if not api_key: return "Error: DAYTONA_API_KEY environment variable not set" # Initialize the Daytona client and create a sandbox config = DaytonaConfig(api_key=api_key) daytona_client = Daytona(config) _daytona_sandbox = daytona_client.create() # Execute the code based on the specified language if language.lower() == "python": response = _daytona_sandbox.process.code_run(code, timeout=timeout) else: # For non-Python languages, create a temporary file and execute it timestamp = datetime.now().strftime("%Y%m%d%H%M%S") file_extension = { "javascript": "js", "nodejs": "js", "ruby": "rb", "php": "php", "bash": "sh", "shell": "sh", "powershell": "ps1", "c": "c", "cpp": "cpp", "java": "java", "go": "go", "rust": "rs", }.get(language.lower(), "txt") filename = f"/tmp/code_{timestamp}.{file_extension}" # Upload the code file to the sandbox _daytona_sandbox.fs.upload_file(filename, code.encode('utf-8')) # Prepare the execution command based on language exec_cmd = { "javascript": f"node {filename}", "nodejs": f"node {filename}", "ruby": f"ruby {filename}", "php": f"php {filename}", "bash": f"bash {filename}", "shell": f"sh {filename}", "powershell": f"pwsh {filename}", "c": f"gcc {filename} -o /tmp/prog_{timestamp} && /tmp/prog_{timestamp}", "cpp": f"g++ {filename} -o /tmp/prog_{timestamp} && /tmp/prog_{timestamp}", "java": f"javac {filename} && java -cp /tmp {os.path.basename(filename).split('.')[0]}", "go": f"go run {filename}", "rust": f"rustc {filename} -o /tmp/prog_{timestamp} && /tmp/prog_{timestamp}", }.get(language.lower(), f"cat {filename}") # Execute the command response = _daytona_sandbox.process.exec(exec_cmd, cwd="/tmp", timeout=timeout) # Extract and return the result if hasattr(response, 'result'): result = response.result elif hasattr(response, 'stdout'): result = response.stdout else: result = str(response) return f"Code Execution Result ({language}):\n{result}" except Exception as e: # Clean up on error try: if _daytona_sandbox is not None: _daytona_sandbox = None except: pass return f"Error executing code: {str(e)}" @tool def execute_shell_command(command: str, working_dir: str = "/tmp", timeout: int = 300) -> str: """Execute a shell command securely in an isolated sandbox environment using Daytona. This tool runs shell commands in a secure, isolated environment to prevent security issues. It's useful for file operations, system tasks, and other command-line operations. Parameters: - command: The shell command to execute - working_dir: The working directory (default: "/tmp") - timeout: Maximum execution time in seconds (default: 30) Returns: - The command execution output or error message """ global _daytona_sandbox try: # Initialize Daytona client if not already done if _daytona_sandbox is None: api_key = os.getenv("DAYTONA_API_KEY") if not api_key: return "Error: DAYTONA_API_KEY environment variable not set" # Initialize the Daytona client and create a sandbox config = DaytonaConfig(api_key=api_key) daytona_client = Daytona(config) _daytona_sandbox = daytona_client.create() # Execute the command response = _daytona_sandbox.process.exec(command, cwd=working_dir, timeout=timeout) # Extract and return the result if hasattr(response, 'result'): result = response.result elif hasattr(response, 'stdout'): result = response.stdout else: result = str(response) return f"Shell Command Execution Result:\n{result}" except Exception as e: # Clean up on error try: if _daytona_sandbox is not None: _daytona_sandbox = None except: pass return f"Error executing shell command: {str(e)}" @tool def sandbox_file_operation(operation: str, file_path: str, content: str = "", target_path: str = "") -> str: """Perform file operations in the secure sandbox environment. This tool allows secure file manipulation in an isolated sandbox. It supports creating, reading, writing, moving, copying and deleting files. Parameters: - operation: The operation to perform ('create', 'read', 'write', 'append', 'delete', 'move', 'copy', 'list') - file_path: Path to the file to operate on - content: Content to write (for 'create', 'write', 'append' operations) - target_path: Target path for 'move' and 'copy' operations Returns: - Operation result or file content """ global _daytona_sandbox try: # Initialize Daytona client if not already done if _daytona_sandbox is None: api_key = os.getenv("DAYTONA_API_KEY") if not api_key: return "Error: DAYTONA_API_KEY environment variable not set" # Initialize the Daytona client and create a sandbox config = DaytonaConfig(api_key=api_key) daytona_client = Daytona(config) _daytona_sandbox = daytona_client.create() # Perform the requested operation operation = operation.lower() if operation == "create" or operation == "write": # Create or overwrite file _daytona_sandbox.fs.upload_file(file_path, content.encode('utf-8')) return f"File {file_path} created/written successfully" elif operation == "append": # First try to read the existing content try: existing_content = _daytona_sandbox.fs.download_file(file_path).decode('utf-8') except: existing_content = "" # Append new content and write back new_content = existing_content + content _daytona_sandbox.fs.upload_file(file_path, new_content.encode('utf-8')) return f"Content appended to {file_path} successfully" elif operation == "read": # Read file content try: content = _daytona_sandbox.fs.download_file(file_path).decode('utf-8') return f"Content of {file_path}:\n{content}" except Exception as e: return f"Error reading {file_path}: {str(e)}" elif operation == "delete": # Delete file response = _daytona_sandbox.process.exec(f"rm -f {file_path}", cwd="/tmp") return f"File {file_path} deleted" elif operation == "move": # Move file if not target_path: return "Error: Target path required for move operation" response = _daytona_sandbox.process.exec(f"mv {file_path} {target_path}", cwd="/tmp") return f"File moved from {file_path} to {target_path}" elif operation == "copy": # Copy file if not target_path: return "Error: Target path required for copy operation" response = _daytona_sandbox.process.exec(f"cp {file_path} {target_path}", cwd="/tmp") return f"File copied from {file_path} to {target_path}" elif operation == "list": # List directory contents response = _daytona_sandbox.process.exec(f"ls -la {file_path}", cwd="/tmp") if hasattr(response, 'result'): result = response.result elif hasattr(response, 'stdout'): result = response.stdout else: result = str(response) return f"Directory listing of {file_path}:\n{result}" else: return f"Unsupported operation: {operation}" except Exception as e: return f"Error performing file operation: {str(e)}" def cleanup_daytona_sandbox(): """Clean up the Daytona sandbox when it's no longer needed.""" global _daytona_sandbox try: if _daytona_sandbox is not None: # Get the Daytona client api_key = os.getenv("DAYTONA_API_KEY") if api_key: config = DaytonaConfig(api_key=api_key) daytona_client = Daytona(config) # Remove the sandbox daytona_client.remove(_daytona_sandbox) _daytona_sandbox = None print("Daytona sandbox cleaned up successfully") except Exception as e: print(f"Error cleaning up Daytona sandbox: {str(e)}") # Track last execution time for rate limiting _last_extract_url_time = 0 @tool def extract_document_data(input_method: str, files: list, prompt: str, json_mode: bool = False) -> str: """Extract structured data from documents using Dumpling AI. This tool allows you to extract information from various document formats including PDFs, Office documents, images, and many other file types. It uses vision-capable Large Language Models (LLMs) to interpret and extract data based on your specific prompt. Parameters: - input_method: How to input files, either "url" or "base64" - files: List of file URLs or base64-encoded strings depending on input_method - prompt: Specific instructions for what data to extract from the document - json_mode: Whether to return structured JSON (true) or free text (false) Returns: - Extracted data from the document based on your prompt Supported file extensions include PDFs, Word docs, Excel files, PowerPoint, images, HTML, and many others. """ api_key = os.getenv("DUMPLING_API_KEY") if not api_key: return "Error: DUMPLING_API_KEY environment variable not set" try: url = "https://app.dumplingai.com/api/v1/extract-document" headers = { "Content-Type": "application/json", "Authorization": f"Bearer {api_key}" } data = { "inputMethod": input_method, "files": files, "prompt": prompt, "jsonMode": json_mode } response = requests.post(url, headers=headers, json=data, timeout=120) response.raise_for_status() result = response.json() # Format the response in a readable way formatted_response = f"Document Extraction Results:\n\n" formatted_response += f"Extracted Data:\n{result.get('results', 'No results found')}\n\n" formatted_response += f"Pages Processed: {result.get('pages', 'Unknown')}\n" formatted_response += f"Files Processed: {result.get('fileCount', 'Unknown')}\n" formatted_response += f"Credit Usage: {result.get('creditUsage', 'Unknown')}\n" return formatted_response except requests.exceptions.Timeout: return "Error: Request to Dumpling AI API timed out after 120 seconds" except requests.exceptions.HTTPError as e: error_detail = f"HTTP Error: {e.response.status_code}" try: error_json = e.response.json() error_detail += f" - {error_json.get('detail', error_json)}" except: error_detail += f" - {e.response.text[:500]}" return error_detail except requests.exceptions.RequestException as e: return f"Error making request to Dumpling AI API: {str(e)}" except Exception as e: return f"Error extracting document data: {str(e)}" @tool def extract_image_data(input_method: str, images: list, prompt: str, json_mode: bool = False) -> str: """Extract visual information from images using Dumpling AI. This tool allows you to extract detailed descriptions or specific information from images using vision-capable Large Language Models (LLMs). It can identify objects, scenes, text, and other visual elements based on your specific prompt. Parameters: - input_method: How to input images, either "url" or "base64" - images: List of image URLs or base64-encoded strings depending on input_method - prompt: Specific instructions for what information to extract from the image - json_mode: Whether to return structured JSON (true) or free text (false) Returns: - Extracted visual data from the image based on your prompt """ api_key = os.getenv("DUMPLING_API_KEY") if not api_key: return "Error: DUMPLING_API_KEY environment variable not set" try: url = "https://app.dumplingai.com/api/v1/extract-image" headers = { "Content-Type": "application/json", "Authorization": f"Bearer {api_key}" } data = { "inputMethod": input_method, "images": images, "prompt": prompt, "jsonMode": json_mode } response = requests.post(url, headers=headers, json=data, timeout=120) response.raise_for_status() result = response.json() # Format the response in a readable way formatted_response = f"Image Analysis Results:\n\n" formatted_response += f"Extracted Data:\n{result.get('results', 'No results found')}\n\n" formatted_response += f"Images Processed: {result.get('imageCount', 'Unknown')}\n" formatted_response += f"Credit Usage: {result.get('creditUsage', 'Unknown')}\n" return formatted_response except requests.exceptions.Timeout: return "Error: Request to Dumpling AI API timed out after 120 seconds" except requests.exceptions.HTTPError as e: error_detail = f"HTTP Error: {e.response.status_code}" try: error_json = e.response.json() error_detail += f" - {error_json.get('detail', error_json)}" except: error_detail += f" - {e.response.text[:500]}" return error_detail except requests.exceptions.RequestException as e: return f"Error making request to Dumpling AI API: {str(e)}" except Exception as e: return f"Error extracting image data: {str(e)}" @tool def extract_url_content(url: str) -> str: """Extract content from a URL using Diffbot API (supports webpages, articles, PDFs, etc.). This function is rate-limited to execute no more frequently than once every 20 seconds.""" global _last_extract_url_time # Check if we need to wait before executing current_time = time.time() time_since_last_call = current_time - _last_extract_url_time if time_since_last_call < 20 and _last_extract_url_time > 0: # Calculate how long to wait wait_time = 20 - time_since_last_call print(f"Rate limiting: waiting {wait_time:.2f} seconds before next API call") time.sleep(wait_time) current_time = time.time() # Update current time after sleeping # Update last execution time _last_extract_url_time = current_time # Diffbot token from environment or use the fallback token = os.getenv("DIFFBOT_TOKEN") if not token: return "Error: DIFFBOT_TOKEN environment variable not set" # Set up the API endpoint api_url = "https://api.diffbot.com/v3/article" # Parameters for the request params = { "token": token, "url": url } try: # Make the API request with a timeout response = requests.get(api_url, params=params, timeout=60) # 30 second timeout response.raise_for_status() # Raise exception for HTTP errors # Parse the response data = response.json() # Extract relevant information if "objects" in data and len(data["objects"]) > 0: obj = data["objects"][0] # Create a formatted result result = f"Title: {obj.get('title', 'No title')}\n\n" if "text" in obj: result += f"Content:\n{obj.get('text')}\n\n" #if "html" in obj: # result += f"HTML Content:\n{obj.get('html')}\n\n" if "categories" in obj and obj["categories"]: categories = ", ".join([f"{cat.get('name')} ({cat.get('score', 0):.2f})" for cat in obj["categories"]]) result += f"Categories: {categories}\n" result += f"Source: {obj.get('siteName', 'Unknown')}\n" result += f"URL: {obj.get('pageUrl', url)}" return result else: return f"No content could be extracted from {url}. Response: {data}" except requests.exceptions.Timeout: return f"Error: Request to extract content from {url} timed out after 30 seconds" except requests.exceptions.RequestException as e: return f"Error: Failed to extract content from {url}: {str(e)}" except Exception as e: return f"Error extracting content from {url}: {str(e)}" @tool def get_youtube_transcript(url: str) -> str: """Get the transcript (captions) from a YouTube video as text. This tool extracts the transcript text from YouTube videos, returns the transcript as a string. Parameters: - url: The YouTube video URL Returns: - The transcript as a string, or an error message if the transcript couldn't be obtained """ # Create a temporary directory to store subtitle files temp_dir = tempfile.mkdtemp() current_dir = os.getcwd() try: # Change to temp directory for download os.chdir(temp_dir) ydl_opts = { 'writesubtitles': True, # Download subtitles 'writeautomaticsub': True, # Download automatic subtitles 'subtitleslangs': ['en'], # Specify English language 'skip_download': True, # Skip downloading the video, only get subtitles 'outtmpl': 'subtitle', # Simple output template } # Download the subtitles with yt_dlp.YoutubeDL(ydl_opts) as ydl: info_dict = ydl.extract_info(url, download=True) video_title = info_dict.get('title', 'Unknown Title') # Look for subtitle files in the temp directory subtitle_content = "" subtitle_files = list(Path(temp_dir).glob("*.vtt")) + list(Path(temp_dir).glob("*.srt")) if subtitle_files: # Read the first subtitle file found with open(subtitle_files[0], 'r', encoding='utf-8') as f: subtitle_content = f.read() # Clean up the subtitle content to remove timestamps and formatting # This is a simple cleaning - more complex parsing may be needed for perfect results lines = subtitle_content.split('\n') cleaned_lines = [] for line in lines: # Skip time codes, numbering and empty lines if line.strip() and not line.strip().isdigit() and not '-->' in line and not line.startswith('WEBVTT'): cleaned_lines.append(line) subtitle_content = ' '.join(cleaned_lines) return f"Transcript from YouTube video: '{video_title}'\n\n{subtitle_content}" else: return f"No transcript found for YouTube video: '{video_title}'" except Exception as e: return f"Error retrieving YouTube transcript: {str(e)}" finally: # Change back to original directory and clean up os.chdir(current_dir) # Cleanup files (optional) try: for file in os.listdir(temp_dir): os.remove(os.path.join(temp_dir, file)) os.rmdir(temp_dir) except: pass class BasicAgent: def __init__(self): print("BasicAgent initialized.") # Initialize the Anthropic models # Standard model for supervisor and validator self.langfuse_handler = CallbackHandler() self.supervisor_model = ChatAnthropic( model="claude-3-7-sonnet-20250219", max_tokens=20000, anthropic_api_key=os.getenv("ANTHROPIC_API_KEY"), temperature=0.6, # thinking={ # "type": "enabled", # "budget_tokens": 5000 # } ) # Standard model for validator self.validator_model = ChatAnthropic( model="claude-3-7-sonnet-20250219", max_tokens=20000, temperature=0.5, # Lower temperature for more consistent validation anthropic_api_key=os.getenv("ANTHROPIC_API_KEY") ) # Tool-enabled model for worker self.worker_model_base = ChatAnthropic( model="claude-3-7-sonnet-20250219", max_tokens=20000, temperature=0.75, anthropic_api_key=os.getenv("ANTHROPIC_API_KEY") ) # Initialize tools self.tools = [search_web_tavily, search_web_serper, execute_code_securely, execute_shell_command, sandbox_file_operation, extract_document_data, extract_image_data, extract_url_content, get_youtube_transcript] # Bind tools only to the worker model self.worker_model = self.worker_model_base.bind_tools(self.tools) # Create the tool node for executing tools self.tool_node = ToolNode(self.tools) # Create the workflow self.app = self._create_workflow() def _process_messages_after_tools(self, messages): """Process messages to ensure tool calls and tool results are properly paired. This helps prevent the Anthropic error: unexpected `tool_use_id` found in `tool_result` blocks.""" # Create a mapping of tool_call_id to AIMessage index tool_call_map = {} for i, msg in enumerate(messages): if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None): for tool_call in msg.tool_calls: if "id" in tool_call: tool_call_map[tool_call["id"]] = i # Filter out ToolMessages that don't have a matching AIMessage with tool_calls processed_messages = [] for i, msg in enumerate(messages): if isinstance(msg, ToolMessage) and hasattr(msg, "tool_call_id"): # Only include if there is a matching AIMessage with this tool_call_id if msg.tool_call_id in tool_call_map: ai_msg_index = tool_call_map[msg.tool_call_id] # Make sure this tool message comes right after its AIMessage if i > ai_msg_index and not any( isinstance(messages[j], ToolMessage) and hasattr(messages[j], "tool_call_id") and messages[j].tool_call_id == msg.tool_call_id for j in range(ai_msg_index + 1, i) ): processed_messages.append(msg) else: processed_messages.append(msg) return processed_messages def _create_workflow(self): workflow = StateGraph(AgentState) # Add nodes workflow.add_node("supervisor", self._supervisor_agent) workflow.add_node("worker", self._worker_agent) workflow.add_node("tools", self._handle_tools) workflow.add_node("validator", self._validation_agent) # Add edges using the START and END constants workflow.add_edge(START, "supervisor") # All nodes use Command to specify their next destination, so we don't need conditional edges # Each node's Command(goto=...) specifies the next node # Compile the graph return workflow.compile() def _supervisor_agent(self, state: AgentState) -> Command: """Supervisor agent that coordinates the workflow.""" # Get the question from state current_question = state["current_question"] messages = state["messages"] worker_iterations = state.get("worker_iterations", 0) # If we have messages and this isn't the first iteration, evaluate worker's response if messages and worker_iterations > 0: # Find the last worker response worker_response = None for msg in reversed(messages): if isinstance(msg, AIMessage) and not getattr(msg, "tool_calls", None): worker_response = msg.content break if worker_response: # Evaluate the worker's response eval_prompt = ChatPromptTemplate.from_messages([ ("system", """You are a supervisor agent evaluating a worker's research report about user's question. Analyze whether the report with answer completely and accurately answers the question. Your evaluation criteria: - Completeness: Does the answer address all aspects of the question? - Accuracy: Are the facts, references and reasoning correct? - Path clarity: Is the path to the answer logical and well-explained? - Evidence quality: Are the references reliable and directly relevant? Worker has access to search and web content extraction tools, also python code execution tool. Tasks given to You are not casual questions by random humans, but tricky contest puzzles that test LLM capabilities. If all criteria are met, respond with "SATISFIED". If any criteria are not met, respond with "UNSATISFIED: [specific detailed feedback]". Be precise in your feedback so the worker knows exactly what to improve."""), ("human", f"Question: {current_question}\nWorker's report with answer: {worker_response}") ]) evaluation = self.supervisor_model.invoke(eval_prompt.format_prompt().to_messages()).content # Determine if supervisor is satisfied supervisor_satisfaction = evaluation.startswith("SATISFIED") if supervisor_satisfaction: # If satisfied, prepare to move to validator return Command( goto="validator", update={ "supervisor_satisfaction": True } ) else: # If not satisfied, give feedback to worker feedback = evaluation.replace("UNSATISFIED: ", "") prompt = ChatPromptTemplate.from_messages([ ("system", """You are a supervisor agent providing targeted feedback to the worker agent. Your role is to guide the worker to improve their research report by: 1) Highlighting specific areas that need improvement 2) Providing clear, actionable guidance on what additional research is needed 3) Explaining exactly how the worker should revise their approach 4) Reminding them of any specific formatting requirements in the original question Worker has access to the following tools: - Web search (using Tavily and Serper) - Web content extraction - Image analysis (can extract visual information from images) - Document data extraction (from PDFs, documents, etc.) - Secure code execution (for Python and other languages) - Secure shell command execution - Secure file operations For computational puzzles, math problems, data processing, or tasks requiring exact precision, recommend using the code execution tools rather than relying on reasoning alone. Tasks given to You are not casual questions by random humans, but tricky contest puzzles that test LLM capabilities. Focus on being constructive and precise. The worker should understand exactly what to do next."""), ("human", f"Question: {current_question}\nWorker's current response: {worker_response}\nImprovement needed: {feedback}") ]) feedback_message = self.supervisor_model.invoke(prompt.format_prompt().to_messages()).content # Update messages with feedback and increment worker iterations return Command( goto="worker", update={ "messages": messages + [HumanMessage(content=feedback_message)], "worker_iterations": worker_iterations + 1, "supervisor_satisfaction": False } ) # First iteration, provide initial instructions prompt = ChatPromptTemplate.from_messages([ ("system", """You are a supervisor agent responsible for coordinating a research workflow. Your responsibilities: 1) Analyze the question to identify required knowledge, tools, and research strategy 2) Provide clear, specific instructions to the worker agent 3) Specify exactly what information to gather and what analysis to perform The worker will prepare a concise research report containing: 1) Their research path - the logical sequence of steps taken to reach the answer 2) The specific references used with clear citations 3) A proposed final answer formatted EXACTLY as requested in the question in separate section Worker has access to the following powerful tools: - Web search (using Tavily and Serper) - Web content extraction - Image analysis (can extract visual information from images) - Document data extraction (can extract data from PDFs, documents, etc.) - Secure code execution (for Python and other languages) - Secure shell command execution - Secure file operations You must understand LLM limitations of solving puzzles that can be solved only by code execution, for example math problems, word character flipping, counting and similar tasks that typically plain LLM will fail at. In case of such tasks, worker should use the code execution tools to solve the puzzle. Tasks given to You are not casual questions by random humans, but tricky contest puzzles that test LLM capabilities. Worker should give You full report with all sections for You to evaluate.""" ), ("human", current_question) ]) response = self.supervisor_model.invoke(prompt.format_prompt().to_messages()).content # Use Command pattern to update state and move to worker return Command( goto="worker", update={ "messages": [HumanMessage(content=current_question), AIMessage(content=response)], "worker_iterations": 1, "supervisor_satisfaction": False } ) def _worker_agent(self, state: AgentState) -> Command: """Worker agent that performs the actual work using tools when needed.""" messages = state["messages"] # Process messages to ensure proper tool call-result pairing processed_messages = self._process_messages_after_tools(messages) # Filter out any ToolMessages that don't have a corresponding AIMessage with tool_calls # This helps prevent the "unexpected tool_use_id" error with Anthropic filtered_messages = [] tool_call_ids = set() # First pass: collect all tool_call_ids from AIMessages for msg in processed_messages: if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None): for tool_call in msg.tool_calls: if "id" in tool_call: tool_call_ids.add(tool_call["id"]) # Second pass: only include ToolMessages that have a corresponding tool_call_id for msg in processed_messages: if isinstance(msg, ToolMessage) and getattr(msg, "tool_call_id", None): if msg.tool_call_id in tool_call_ids: filtered_messages.append(msg) else: filtered_messages.append(msg) # If messages exist, use them directly with the tool-enabled model response = self.worker_model.invoke(filtered_messages) # Update messages - add the response to the original messages # We don't want to lose the original message history updated_messages = messages + [response] # Determine next step using Command pattern if response.tool_calls: # If tool calls are present, go to tools return Command( goto="tools", update={"messages": updated_messages} ) else: # No tool calls, return to supervisor for evaluation return Command( goto="supervisor", update={"messages": updated_messages} ) def _validation_agent(self, state: AgentState) -> Command: """Agent that validates the final answer.""" messages = state["messages"] question = state["current_question"] # Get the final answer from the last message final_answer = "" for msg in reversed(messages): if isinstance(msg, AIMessage) and not getattr(msg, "tool_calls", None): final_answer = msg.content break prompt = ChatPromptTemplate.from_messages([ ("system", """You are a quality assurance agent responsible for final verification of research reports and precise formatting of final answers. Your critical responsibilities: 1) Verify the factual accuracy and completeness of the report, ensuring you can extract and format the final answer exactly as requested in the question 2) Ensure EXACT compliance with any formatting instructions in the question by producing a properly structured final answer Pay extremely close attention to formatting requirements. The user may request: - Only specific parts of information (first/last names, specific data points, numerical values) - Particular ordering (alphabetical, chronological, size-based, relevance-based) - Special formatting (bullet points, numbered lists, specific separators, tables) - Exact text case, spacing, punctuation, or other presentational elements Exact formatting compliance is MANDATORY for this challenge evaluation. Your role is to ensure the final answer meets all specified requirements. If numerical values are requested, ensure they are formatted as numbers, not text. Remember that the worker had access to: - Web search tools - Web content extraction - Image analysis (can extract visual information from images) - Document data extraction (from PDFs, documents, etc.) - Secure code execution - Secure shell commands - Secure file operations For computational or precision-based questions, check if code execution was appropriately used and validate the results. When evaluating the answer: - Check if all required information is present and accurate - Verify that the answer directly addresses the specific question asked - Ensure any numerical values, dates, names, or technical terms are correct - Confirm that the formatting precisely matches what was requested - Do not add units to the final answer if not explicitly requested - Do not use money symbols like in the final answer if not explicitly requested - Dont use comma separators for integers like 1,000,000, just use 1000000 - Answers tend to be as short as possible, so do not add extra data unless explicitly requested If the answer report is correct, format it exactly as asked in the question, and respond with: "APPROVED: [THE PROPERLY FORMATTED ANSWER]" If there are issues with overall answer quality and you cannot format the final answer as requested, respond with: "REJECTED: [DETAILED EXPLANATION OF ISSUES]" Be extremely precise in your evaluation - the success of this task depends on your attention to detail. """ ), ("human", f"Question: {question}\nReport to validate: {final_answer}") ]) validation_result = self.validator_model.invoke(prompt.format_prompt().to_messages()).content validator_approval = validation_result.startswith("APPROVED") if validator_approval: # Approved - end the workflow return Command( goto=END, update={ "final_answer": validation_result[10:], # Remove "APPROVED: " prefix "validation_result": validation_result, "validator_approval": True } ) else: # Rejected - restart from supervisor with reset state return Command( goto="supervisor", update={ "messages": [HumanMessage(content=question)], "validation_result": validation_result, "validator_approval": False, "worker_iterations": 0, "supervisor_satisfaction": False } ) def _handle_tools(self, state: AgentState) -> Command: """Custom wrapper around ToolNode to ensure proper message handling.""" # Execute the tool using the tool node tool_result = self.tool_node.invoke(state) # Process the result to ensure proper message ordering if "messages" in tool_result: # Get original messages original_messages = state["messages"] # Get all existing AIMessages with tool calls and their indices ai_indices = {} for i, msg in enumerate(original_messages): if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None): for tool_call in msg.tool_calls: if "id" in tool_call: ai_indices[tool_call["id"]] = i # Add the new tool messages, ensuring they come right after their corresponding tool call updated_messages = list(original_messages) for msg in tool_result["messages"]: if isinstance(msg, ToolMessage) and hasattr(msg, "tool_call_id"): tool_id = msg.tool_call_id if tool_id in ai_indices: # Insert after the AIMessage with the matching tool call insert_idx = ai_indices[tool_id] + 1 # Move past any existing tool messages for this AI message while insert_idx < len(updated_messages) and \ isinstance(updated_messages[insert_idx], ToolMessage) and \ hasattr(updated_messages[insert_idx], "tool_call_id") and \ updated_messages[insert_idx].tool_call_id != tool_id: insert_idx += 1 updated_messages.insert(insert_idx, msg) # Update subsequent indices for id in ai_indices: if ai_indices[id] >= insert_idx: ai_indices[id] += 1 else: # No matching tool call found, just append updated_messages.append(msg) return Command( goto="worker", update={"messages": updated_messages} ) # If no message updates, just return the state return Command( goto="worker", update=tool_result ) def __call__(self, question: str) -> str: print(f"Agent received question (first 50 chars): {question[:50]}...") # Initialize the state initial_state = { "messages": [], "current_question": question, "final_answer": "", "validation_result": "", "worker_iterations": 0, "supervisor_satisfaction": False, "validator_approval": False } try: # Run the workflow final_state = self.app.invoke(initial_state, config={"callbacks": [self.langfuse_handler], "recursion_limit": 35}) # Return the final answer answer = final_state.get("final_answer", "") if not answer and final_state["messages"]: for msg in reversed(final_state["messages"]): if isinstance(msg, AIMessage) and not getattr(msg, "tool_calls", None): answer = msg.content break print(f"Agent returning answer: {answer[:50]}...") return answer except Exception as e: print(f"Error in agent processing: {str(e)}") # Fallback to basic workflow without tool calls if there's an error return f"I encountered an error while processing your question: {str(e)}. Please try reformulating your question." finally: # Clean up resources cleanup_daytona_sandbox()