agents/reasoning_agent.py

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}")