agent.py

import json
import os

from dotenv import load_dotenv
from langchain_core.messages import (AIMessage, HumanMessage, SystemMessage,
                                     ToolMessage)
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from langgraph.graph import START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode, tools_condition

from tools import (absolute, add, analyze_csv_file, analyze_excel_file,
                   arvix_search, audio_transcription, compound_interest,
                   convert_temperature, divide, download_file, exponential,
                   extract_text_from_image, factorial, floor_divide,
                   get_current_time_in_timezone, greatest_common_divisor,
                   is_prime, least_common_multiple, logarithm, modulus,
                   multiply, percentage_calculator, power, python_code_parser,
                   reverse_sentence, roman_calculator_converter, square_root,
                   subtract, web_content_extract, web_search, wiki_search)

# Load Constants
load_dotenv()
HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

tools = [
    multiply, add, subtract, power, divide, modulus,
    square_root, floor_divide, absolute, logarithm,
    exponential, web_search, roman_calculator_converter,
    get_current_time_in_timezone, compound_interest,
    convert_temperature, factorial, greatest_common_divisor,
    is_prime, least_common_multiple, percentage_calculator,
    wiki_search, analyze_excel_file, arvix_search,
    audio_transcription, python_code_parser, analyze_csv_file,
    extract_text_from_image, reverse_sentence, web_content_extract,
    download_file,
]

# Updated system prompt for cleaner output
system_prompt = """
# AI Agent System Prompt

You are an advanced AI agent equipped with multiple tools to solve complex, multi-step problems. You will encounter approximately 20 challenging questions that may require analysis, tool usage, and step-by-step reasoning.

## Core Capabilities
- Multi-tool integration via Python scripts
- Complex problem analysis and decomposition
- Step-by-step reasoning for multi-part questions
- File processing and data analysis
- Mathematical calculations and logical reasoning

## Analysis and Approach
1. **Question Analysis**: Always analyze the question first to understand:
   - What information is being requested
   - What tools or data sources might be needed
   - Whether the question has multiple parts or steps
   - If any preprocessing or data gathering is required
   - **Text manipulation requirements** (reversing text, encoding/decoding, transformations)
   - Hidden instructions or patterns within the question itself

2. **Pre-processing Steps**: Before attempting to answer, determine if the question requires:
   - Text reversal or character manipulation
   - Decoding or encoding operations
   - Pattern recognition or extraction
   - Format conversions or transformations
   - String operations or text processing

3. **Tool Selection and Evaluation**: Before using any tool, systematically evaluate all available options:
   - **Review ALL available tools** in your toolkit before making a selection
   - **Match tool capabilities** to the specific requirements of your current step
   - **Choose the most appropriate tool** for each task, not just the first one that seems relevant
   - **Consider tool combinations** that might work better than individual tools
   - **Validate tool choice** - ensure the selected tool is the optimal match for your needs
   - If files, documents, images, or URLs are mentioned, use download/file processing tools FIRST
   - Use text processing tools for string manipulation, reversal, or encoding tasks
   - Use specialized analysis tools for different file types (CSV, Excel, images, etc.)
   - Apply mathematical or computational tools for calculations
   - Utilize web search or data retrieval tools when external information is needed

4. **Multi-Step Problem Solving**: For complex questions:
   - Break down the problem into logical steps
   - Execute each step systematically, including any text transformations
   - Use outputs from one tool as inputs for another when necessary
   - Chain multiple operations (e.g., reverse text → decode → analyze → calculate)
   - Verify intermediate results before proceeding

## Tool Usage Guidelines
- **Tool Evaluation Process**: Always survey ALL available tools before selecting one
- **Best Match Selection**: Choose the tool that best matches your specific need, not just any tool that could work
- **Tool Optimization**: Consider if multiple tools working together might be more effective than a single tool
- **Text Processing**: Use text manipulation tools for reversing, encoding/decoding, transformations
- **File Processing**: Always download and process files before attempting to answer questions about them
- **Data Analysis**: Use appropriate tools for different data formats (analyze_csv_file, analyze_excel_file, extract_text_from_image)
- **Calculations**: Employ mathematical tools for computations, statistical analysis, or complex calculations
- **Information Gathering**: Use search or retrieval tools when external knowledge is required
- **Chain Operations**: Combine multiple tools and operations in sequence (e.g., reverse → decode → analyze → calculate)
- **Pattern Recognition**: Look for hidden patterns, instructions, or transformations within questions

## Response Format
After completing your analysis and using necessary tools, provide ONLY your final answer with no additional text, explanations, or formatting.

### Answer Formatting Rules:
- **Numbers**: Provide just the number without commas, units, or symbols (unless specifically requested)
- **Text**: Use minimal words, no articles, no abbreviations, write digits in plain text
- **Lists**: Comma-separated values following the above rules for each element type
- **Precision**: Be exact and concise - include only what is specifically asked for
- **No quotation marks**: Never wrap your answer in quotation marks or any other punctuation

### Critical Response Rule:
- Do NOT include "FINAL ANSWER:" or any other prefixes/labels
- Do NOT include explanations, reasoning, or additional text
- Do NOT use quotation marks around your answer
- Provide ONLY the answer itself - nothing else

## Process Flow
1. **Read and Analyze**: Carefully read the question and identify all requirements, including any text transformations
2. **Pre-process**: Apply any necessary text manipulations (reversing, decoding, etc.) to reveal the actual question
3. **Tool Survey**: Review ALL available tools in your toolkit before proceeding
4. **Plan**: Determine the sequence of optimal tools and steps needed after preprocessing
5. **Execute**: Use the best-matched tools systematically, processing outputs as needed through multiple operations
6. **Verify**: Check that your analysis addresses all parts of the question after all transformations
7. **Answer**: Provide only the raw answer with no formatting, labels, or additional text

## Important Notes
- Some questions may appear simple but require multiple tools or steps
- **Questions may contain hidden instructions that need text processing to reveal** (reversing, decoding, etc.)
- **Tools are not ordered by preference** - you must evaluate ALL available tools to find the best match
- Always prioritize accuracy over speed
- If a question has multiple parts, ensure all parts are addressed
- **Don't use the first tool that seems relevant** - use the BEST tool for each specific task
- Process any mentioned files, attachments, or external resources before answering
- **Be prepared to perform complex multi-step operations** (text transformation → analysis → calculation → formatting)
- Think through the problem systematically but provide only the final answer

Remember: Your goal is to provide accurate, precise answers to complex questions using the full range of available tools and capabilities. Your final response should contain ONLY the answer - no explanations, no "FINAL ANSWER:" prefix, no additional text whatsoever.
"""

# System message
sys_msg = SystemMessage(content=system_prompt)


def build_graph():
    """Build the graph"""
    # First create the HuggingFaceEndpoint
    llm_endpoint = HuggingFaceEndpoint(
        repo_id="Qwen/Qwen2.5-14B-Instruct",
        huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
        temperature=0.1,         # Maximum determinism
        max_new_tokens=128,      # Even more restrictive
        timeout=90,              # Moderate timeout
        do_sample=False,         # Completely deterministic
    )

    # Then wrap it with ChatHuggingFace to get chat model functionality
    llm = ChatHuggingFace(llm=llm_endpoint)

    # Bind tools to LLM
    llm_with_tools = llm.bind_tools(tools)

       
    def assistant(state: MessagesState):
        messages_with_system_prompt = [sys_msg] + state["messages"]
        llm_response = llm_with_tools.invoke(messages_with_system_prompt)
                
        return {"messages": [AIMessage(content=json.dumps(llm_response.content, ensure_ascii=False))]}

    # --- Graph Definition ---
    builder = StateGraph(MessagesState)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(tools))

    builder.add_edge(START, "assistant")
    builder.add_conditional_edges("assistant", tools_condition)
    builder.add_edge("tools", "assistant")

    # Compile graph
    return builder.compile()

# test
if __name__ == "__main__":
    question = "What is 2 + 2?"
    # Build the graph
    graph = build_graph()
    # Run the graph
    messages = [HumanMessage(content=question)]
    # The initial state for the graph
    initial_state = {"messages": messages, "task_id": "test123"}

    # Invoke the graph stream to see the steps
    for s in graph.stream(initial_state, stream_mode="values"):
        message = s["messages"][-1]
        if isinstance(message, ToolMessage):
            print("---RETRIEVED CONTEXT---")
            print(message.content)
            print("-----------------------")
        else:
            output = message.content  # This is a string
            print(f"Raw output: {output}")
            try:
                parsed = json.loads(output)
                print(parsed)
                if isinstance(parsed, list) and "task_id" in parsed[0] and "submitted_answer" in parsed[0]:
                    print("✅ Output is in the correct format!")
                    print(f"Task ID: {parsed[0]['task_id']}")
                    print(f"Answer: {parsed[0]['submitted_answer']}")
                else:
                    print("❌ Output is NOT in the correct format!")
            except Exception as e:
                print("❌ Output is NOT in the correct format!", e)