ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/examples/imo/run.py

@@ -0,0 +1,306 @@
+import argparse
+import json
+import logging
+import os
+import re
+import sys
+import traceback
+from pathlib import Path
+from typing import Any, Dict, List
+
+from dotenv import load_dotenv
+from pathlib import Path
+
+from aworld.agents.llm_agent import Agent
+from aworld.config.conf import AgentConfig, TaskConfig
+from aworld.core.task import Task
+from aworld.runner import Runners
+from aworld.core.task import Task
+from prompt import system_prompt
+from utils import (
+    add_file_path,
+    load_dataset_meta,
+)
+from guard_tool_caller import GuardToolCaller
+
+# Create log directory if it doesn't exist
+if not os.path.exists(os.getenv("AWORLD_WORKSPACE", "~")):
+    os.makedirs(os.getenv("AWORLD_WORKSPACE", "~"))
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--q",
+    type=str,
+    help="Question Index, e.g., imo6. Highest priority: override other arguments if provided.",
+)
+parser.add_argument(
+    "--skip",
+    action="store_true",
+    help="Skip the question if it has been processed before.",
+)
+args = parser.parse_args()
+
+
+def setup_logging():
+    logging_logger = logging.getLogger()
+    logging_logger.setLevel(logging.INFO)
+
+    log_file_name = f"/solution_{args.q}.log" if args.q else f"/solution.log"
+    file_handler = logging.FileHandler(
+        os.getenv("AWORLD_WORKSPACE", "~") + log_file_name,
+        mode="a",
+        encoding="utf-8",
+    )
+    file_handler.setLevel(logging.INFO)
+
+    formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    file_handler.setFormatter(formatter)
+
+    logging_logger.addHandler(file_handler)
+
+
+class GuardRunner:
+    """The guard tool runner"""
+    
+    def __init__(self, super_agent: Agent, guard_tool_caller: GuardToolCaller, original_task: str):
+        self.super_agent = super_agent
+        self.guard_tool_caller = guard_tool_caller
+        self.original_task = original_task
+        self.conversation_history = []
+        self.max_iterations = 10  # max conversation rounds
+        
+    async def run_conversation(self, question: str) -> str:
+        """run the conversation"""
+        current_input = question
+        iteration = 0
+        
+        while iteration < self.max_iterations:
+            logging.info(f"=== The {iteration + 1} round of the conversation ===")
+            
+            # 1. Super-agent handles the current input
+            logging.info(f"Super-agent input: {current_input[:100]}...")
+            super_output = await self._call_super_agent(current_input)
+            logging.info(f"Super-agent output: {super_output[:200]}...")
+            
+            # 2. Check if the output contains the final answer
+            if self._has_final_answer(super_output) and iteration > 1:
+                logging.info("find the final answer, the conversation should stop")
+                return super_output
+            
+            if iteration == self.max_iterations - 1: 
+                logging.info("reach the max conversation rounds, the conversation should stop and return the final answer")
+                return super_output
+            
+            # 3. Call the guard tool
+            logging.info("guard tool is being called...")
+            guard_output = await self._call_guard_tool(super_output)
+            logging.info(f"guard tool output: {guard_output[:200]}...")
+            
+            # 4. Record the current round of the conversation history (before preparing the next input)
+            self.conversation_history.append({
+                "iteration": iteration + 1,
+                "super_input": current_input,
+                "super_output": super_output,
+                "guard_output": guard_output
+            })
+            
+            # 5. Prepare the next input
+            next_input = self._prepare_next_input(super_output, guard_output)
+            current_input = next_input
+            
+            iteration += 1
+            
+            # Check if the conversation should continue
+            if self._should_stop_conversation(super_output, guard_output):
+                logging.info("the conversation should stop")
+                break
+        
+        logging.info(f"the converstation finished, there are totally {iteration} rounds")
+        return super_output
+    
+    async def _call_super_agent(self, input_text: str) -> str:
+        """call the super-agent"""
+        try:
+            # create the task
+            task = Task(
+                input=input_text, 
+                agent=self.super_agent, 
+                conf=TaskConfig()
+            )
+            
+            # run the task
+            result = Runners.sync_run_task(task=task)
+            
+            # extract the answer
+            if result and task.id in result:
+                return result[task.id].answer
+            else:
+                return "Super-agent fail to return the result"
+                
+        except Exception as e:
+            logging.error(f"fail to call the super-agent: {e}")
+            return f"fail to call the super-agent: {str(e)}"
+    
+    async def _call_guard_tool(self, super_output: str) -> str:
+        """call the guard tool"""
+        try:
+            # call the guard tool
+            guard_result = await self.guard_tool_caller.call_guard_tool(super_output, self.original_task)
+            return guard_result
+            
+        except Exception as e:
+            logging.error(f"fail to call the guard tool: {e}")
+            return f"fail to call the guard tool: {str(e)}"
+    
+    def _has_final_answer(self, output: str) -> bool:
+        """check if the output contains the final answer"""
+        answer_patterns = [
+            r"<answer>.*?</answer>",
+            r"Final answer[：:]\s*",
+            r"final answer[：:]\s*",
+            r"The final answer is[：:]\s*",
+            r"the final answer is[：:]\s*",
+            r"the final answer is\s*",
+            r"I have successfully solved the problem."
+        ]
+        
+        for pattern in answer_patterns:
+            if re.search(pattern, output, re.IGNORECASE):
+                return True
+        return False
+    
+    def _prepare_next_input(self, super_output: str, guard_output: str) -> str:
+        """Prepare for the next round with complete conversation history"""
+        # build the complete conversation history
+        conversation_text = ""
+        
+        # add the original task
+        conversation_text += f"Original Task:\n{self.original_task}\n\n"
+        
+        # add all the history conversation rounds (not including the current round)
+        for i, history in enumerate(self.conversation_history):
+            round_num = history["iteration"]
+            conversation_text += f"=== Round {round_num} ===\n"
+            conversation_text += f"Previous solution:\n{history['super_output']}\n\n"
+            conversation_text += f"IMO grader review:\n{history['guard_output']}\n\n"
+        
+        # build the final input
+        combined_input = f"""
+{conversation_text}
+Please seriously consider all the above reviews from the IMO grader across all rounds, and provide a refined and improved solution that addresses all the issues identified.
+"""
+        return combined_input.strip()
+    
+    def _should_stop_conversation(self, super_output: str, guard_output: str) -> bool:
+        """check if the conversation should stop"""
+        # check if the guard tool thinks the conversation should stop
+        stop_indicators = [
+            "The answer is completed",
+            "No need to further refine",
+            "The question has been correctly solved",
+            "No loopholes or oversights found"
+        ]
+        
+        for indicator in stop_indicators:
+            if indicator in guard_output:
+                return True
+        
+        return False
+
+
+if __name__ == "__main__":
+    env_path = ".env"
+    load_dotenv(env_path, override=True, verbose=True)
+    setup_logging()
+
+    imo_dataset_path = os.getenv("IMO_DATASET_PATH", "./imo_dataset")
+    full_dataset = load_dataset_meta(imo_dataset_path)
+    logging.info(f"Total questions: {len(full_dataset)}")
+
+    # create the super-agent (without MCP tools)
+    agent_config = AgentConfig(
+        llm_provider="openai",
+        llm_model_name=os.getenv("LLM_MODEL_NAME", "gpt-4o"),
+        llm_api_key=os.getenv("LLM_API_KEY"),
+        llm_base_url=os.getenv("LLM_BASE_URL"),
+        llm_temperature=0.1,
+    )
+    super_agent = Agent(
+        conf=agent_config,
+        name="gaia_super_agent",
+        system_prompt=system_prompt,
+    )
+
+    # create the guard tool caller
+    guard_tool_caller = GuardToolCaller()
+
+    # load results from the checkpoint file
+    if os.path.exists(os.getenv("AWORLD_WORKSPACE", "~") + "/results.json"):
+        with open(os.getenv("AWORLD_WORKSPACE", "~") + "/results.json", "r", encoding="utf-8") as results_f:
+            results: List[Dict[str, Any]] = json.load(results_f)
+    else:
+        results: List[Dict[str, Any]] = []
+
+    try:
+        # appoint the task+id
+        if args.q is not None:
+            dataset_slice = [dataset_record for dataset_record in full_dataset if dataset_record["task_id"] == args.q]
+            if not dataset_slice:
+                logging.error(f"Task ID '{args.q}' not found in dataset")
+                sys.exit()
+        else:
+            logging.error("Please specify a task_id using --q parameter")
+            sys.exit()
+
+        # main loop to execute questions
+        for i, dataset_i in enumerate(dataset_slice):
+            # run
+            try:
+                logging.info(f"Start to process: {dataset_i['task_id']}")
+                logging.info(f"Question: {dataset_i['Question']}")
+
+                question = add_file_path(dataset_i, file_path=imo_dataset_path)["Question"]
+
+                # use the guard tool runner
+                guard_runner = GuardRunner(
+                    super_agent=super_agent,
+                    guard_tool_caller=guard_tool_caller,
+                    original_task=question
+                )
+                
+                # run the conversation
+                import asyncio
+                result = asyncio.run(guard_runner.run_conversation(question))
+
+                # Create the new result record
+                new_result = {
+                    "task_id": dataset_i["task_id"],
+                    "question": question,
+                    "response": result,
+                    "conversation_history": guard_runner.conversation_history,
+                }
+
+                # Check if this task_id already exists in results
+                existing_index = next(
+                    (i for i, result in enumerate(results) if result.get("task_id") == dataset_i["task_id"]),
+                    None,
+                )
+
+                if existing_index is not None:
+                    # Update existing record
+                    results[existing_index] = new_result
+                    logging.info(f"Updated existing record for task_id: {dataset_i['task_id']}")
+                else:
+                    # Append new record
+                    results.append(new_result)
+                    logging.info(f"Added new record for task_id: {dataset_i['task_id']}")
+
+            except Exception:
+                logging.error(f"Error processing {i}: {traceback.format_exc()}")
+                continue
+    except KeyboardInterrupt:
+        pass
+    finally:
+        # Save results to file
+        with open(os.getenv("AWORLD_WORKSPACE", "~") + "/results.json", "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=4, ensure_ascii=False)