import os import logging from llama_index.core.agent.workflow import ReActAgent from llama_index.core.tools import FunctionTool from llama_index.llms.google_genai import GoogleGenAI from llama_index.llms.openai import OpenAI # Setup logging logger = logging.getLogger(__name__) # Helper function to load prompt from file def load_prompt_from_file(filename: str, default_prompt: str) -> str: """Loads a prompt from a text file.""" try: # Assuming the prompt file is in the same directory as the agent script script_dir = os.path.dirname(__file__) prompt_path = os.path.join(script_dir, filename) with open(prompt_path, "r") as f: prompt = f.read() logger.info(f"Successfully loaded prompt from {prompt_path}") return prompt except FileNotFoundError: logger.warning(f"Prompt file {filename} not found at {prompt_path}. Using default.") return default_prompt except Exception as e: logger.error(f"Error loading prompt file {filename}: {e}", exc_info=True) return default_prompt # --- Tool Function --- def reasoning_tool_fn(context: str) -> str: """ Perform chain-of-thought reasoning over the provided context using a dedicated LLM. Args: context (str): The conversation/workflow history and current problem statement. Returns: str: A structured reasoning trace and conclusion, or an error message. """ logger.info(f"Executing reasoning tool with context length: {len(context)}") # Configuration for the reasoning LLM (OpenAI in the original) reasoning_llm_model = os.getenv("REASONING_LLM_MODEL", "gpt-4o-mini") # Use gpt-4o-mini as default openai_api_key = os.getenv("OPENAI_API_KEY") if not openai_api_key: logger.error("ALPAFLOW_OPENAI_API_KEY not found for reasoning tool LLM.") return "Error: ALPAFLOW_OPENAI_API_KEY must be set to use the reasoning tool." # Define the prompt for the reasoning LLM reasoning_prompt = f"""You are an expert reasoning engine. Analyze the following workflow context and problem statement: --- CONTEXT START --- {context} --- CONTEXT END --- Perform the following steps: 1. **Comprehension**: Identify the core question/problem and key constraints from the context. 2. **Decomposition**: Break the problem into logical sub-steps. 3. **Chain-of-Thought**: Reason through each sub-step, stating assumptions and deriving implications. 4. **Verification**: Check conclusions against constraints. 5. **Synthesis**: Integrate results into a cohesive answer/recommendation. 6. **Clarity**: Use precise language. Respond with your numbered reasoning steps followed by a concise final conclusion or recommendation. """ try: # Note: Original used OpenAI with a specific key and model. Retaining that. # Consider adding `reasoning_effort="high"` if supported and desired. llm = OpenAI( model=reasoning_llm_model, api_key=openai_api_key, reasoning_effort="high", temperature=0.055, max_tokens=16384 ) logger.info(f"Using reasoning LLM: {reasoning_llm_model}") response = llm.complete(reasoning_prompt) logger.info("Reasoning tool execution successful.") return response.text except Exception as e: logger.error(f"Error during reasoning tool LLM call: {e}", exc_info=True) return f"Error during reasoning: {e}" def answer_question(question: str) -> str: """ Answer any question by following this strict format: 1. Include your chain of thought (your reasoning steps). 2. End your reply with the exact template: FINAL ANSWER: [YOUR FINAL ANSWER] YOUR FINAL ANSWER must be: - A number, or - As few words as possible, or - A comma-separated list of numbers and/or strings. Formatting rules: * If asked for a number, do not use commas or units (e.g., $, %), unless explicitly requested. * If asked for a string, do not include articles or abbreviations (e.g., city names), and write digits in plain text. * If asked for a comma-separated list, apply the above rules to each element. This tool should be invoked immediately after completing the final planning sub-step. """ logger.info(f"Answering question: {question[:100]}") gemini_api_key = os.getenv("GEMINI_API_KEY") if not gemini_api_key: logger.error("GEMINI_API_KEY not set for answer_question tool.") return "Error: GEMINI_API_KEY not set." model_name = os.getenv("ANSWER_TOOL_LLM_MODEL", "gemini-2.5-pro-preview-03-25") # Build the assistant prompt enforcing the required format assistant_prompt = ( "You are a general AI assistant. I will ask you a question. " "Report your thoughts, and finish your answer with the following template: " "FINAL ANSWER: [YOUR FINAL ANSWER]. " "YOUR FINAL ANSWER should be a number OR as few words as possible " "OR a comma separated list of numbers and/or strings. " "If you are asked for a number, don't use commas for thousands or any units like $ or % unless specified. " "If you are asked for a string, omit articles and abbreviations, and write digits in plain text. " "If you are asked for a comma separated list, apply these rules to each element.\n\n" f"Question: {question}\n" "Answer:" ) try: llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05) logger.info(f"Using answer LLM: {model_name}") response = llm.complete(assistant_prompt) logger.info("Answer generated successfully.") return response.text except Exception as e: logger.error(f"LLM call failed during answer generation: {e}", exc_info=True) return f"Error during answer generation: {e}" # --- Tool Definition --- reasoning_tool = FunctionTool.from_defaults( fn=reasoning_tool_fn, name="reasoning_tool", description=( "Applies detailed chain-of-thought reasoning to the provided workflow context using a dedicated LLM. " "Input: context (str). Output: Reasoning steps and conclusion (str) or error message." ), ) answer_question = FunctionTool.from_defaults( fn=answer_question, name="answer_question", description=( "Use this tool to answer any question, reporting your reasoning steps and ending with 'FINAL ANSWER: ...'. " "Invoke this tool immediately after the final sub-step of planning is complete." ), ) # --- Agent Initialization --- def initialize_reasoning_agent() -> ReActAgent: """Initializes the Reasoning Agent.""" logger.info("Initializing ReasoningAgent...") # Configuration for the agent's main LLM (Google GenAI) agent_llm_model = os.getenv("REASONING_AGENT_LLM_MODEL", "gemini-2.5-pro-preview-03-25") gemini_api_key = os.getenv("GEMINI_API_KEY") if not gemini_api_key: logger.error("GEMINI_API_KEY not found for ReasoningAgent.") raise ValueError("GEMINI_API_KEY must be set for ReasoningAgent") try: llm = GoogleGenAI(api_key=gemini_api_key, model="gemini-2.5-pro-preview-03-25", temperature=0.05) logger.info(f"Using agent LLM: {agent_llm_model}") # Load system prompt default_system_prompt = ("You are ReasoningAgent... [Default prompt content - replace with actual]" # Placeholder ) system_prompt = load_prompt_from_file("../prompts/reasoning_agent_prompt.txt", default_system_prompt) if system_prompt == default_system_prompt: logger.warning("Using default/fallback system prompt for ReasoningAgent.") agent = ReActAgent( name="reasoning_agent", description=( "An autonomous reasoning specialist that applies `reasoning_tool` to perform " "in-depth chain-of-thought analysis on incoming queries or contexts, " "then seamlessly delegates the synthesized insights to `planner_agent` " "or `long_context_management_agent` for subsequent task orchestration." ), tools=[reasoning_tool], llm=llm, system_prompt=system_prompt, can_handoff_to=[ "code_agent", "research_agent", "math_agent", "role_agent", "image_analyzer_agent", "text_analyzer_agent", "planner_agent", "long_context_management_agent", "advanced_validation_agent", "video_analyzer_agent" ], ) return agent except Exception as e: logger.error(f"Error during ReasoningAgent initialization: {e}", exc_info=True) raise # Example usage (for testing if run directly) if __name__ == "__main__": logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger.info("Running reasoning_agent.py directly for testing...") # Check required keys required_keys = ["GEMINI_API_KEY", "ALPAFLOW_OPENAI_API_KEY"] missing_keys = [key for key in required_keys if not os.getenv(key)] if missing_keys: print(f"Error: Required environment variable(s) not set: {', '.join(missing_keys)}. Cannot run test.") else: try: # Test the reasoning tool directly print("\nTesting reasoning_tool_fn...") test_context = "User asked: What is the capital of France? ResearchAgent found: Paris. VerifierAgent confirmed: High confidence." reasoning_output = reasoning_tool_fn(test_context) print(f"Reasoning Tool Output:\n{reasoning_output}") # Initialize the agent (optional) # test_agent = initialize_reasoning_agent() # print("\nReasoning Agent initialized successfully for testing.") # Example chat (would require context passing mechanism) # result = test_agent.chat("Synthesize the findings about the capital of France.") # print(f"Agent chat result: {result}") except Exception as e: print(f"Error during testing: {e}")