agent.py

import os
import json
import torch
import base64
import requests
import mimetypes
from PIL import Image
from tavily import TavilyClient
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.rate_limiters import InMemoryRateLimiter
from dataclasses import dataclass, field
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_experimental.utilities import PythonREPL
from langchain_community.document_loaders import ArxivLoader
from langchain_core.messages import SystemMessage, HumanMessage, AnyMessage, AIMessage, ToolMessage
from langchain_community.document_loaders import WikipediaLoader
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.graph import StateGraph, START, END, MessagesState
from typing import TypedDict, List, Dict, Any, Optional, Annotated



PLANNER_PROMPT = """You are a helpful assistant tasked with answering the question. When formulating your response, make sure that:

1. You follow the necessary steps to solve the problem, using available tools where required.
2. Once the final result is known, return that result in a clear and concise manner without further elaboration.

You have full access to these tools:

• python_repl_tool(query)  
  – For any Python query: math, data processing, unit conversions, file parsing (pandas, chess, etc.)  
  – Always wrap query in `print(...)` so PythonREPL returns the result.

• web_search(query, max_results)  
  – General factual lookup, definitions, current/historical data.

• wiki_search(query)  
  – Encyclopedic facts from Wikipedia.

• vision_caption(image_path)  
  – Describe images, diagrams, or board positions.

• arxiv_search(query)  
  – Locate scholarly PDF URLs on arXiv.

• download_file_from_url(url)  
  – Retrieve any PDF, text, or CSV file.

• file_reader(file_path)  
  – Load CSV/Excel/text into DataFrames or raw strings.

### Your Task

You will receive one JSON payload:

```json
{ "message": "<message>", "input_files": […]}



Your job is to plan, execute, reflect, and refine iteratively until you can confidently answer the question:

Plan Stage

Emit exactly one JSON array of step objects. Do not call any tools yet; this is just your plan.

Example plan:

- Emit exactly one JSON array of step objects, e.g.:
  ```json
  [
    {"tool":"web_search","args":{"query":"…","max_results":3}},
  ]
  ```

- Do **not** call any tools yet—this is just your plan.

Do not perform any tool calls within the planning stage. Simply structure your plan, ensuring it aligns with the best approach to answer the question.

Iterative Plan Review

After the initial plan, start executing the tools.

Review the results of each tool call. If the results are unsatisfactory, refine your plan, adjust arguments, and ensure you are on the right path.

Always analyze the results carefully, and if necessary, revise the plan or tools based on the insights you gather at each step.

Revisions Based on Results

If the previous step didn’t provide the expected results, revise the plan. Look at the question, tools used, and results from previous steps to improve your approach.

If the plan isn't generating the correct result after five iterations, pass the current plan to the next node (e.g., validator) for further review.

Final Answer Handling

Ensure that FINAL ANSWER is included in the result. If no solution is reached, provide the steps you have taken so far along with the tools used.

Your response should only contain the FINAL ANSWER. Do not include intermediate steps or re-plans once a solution has been found.

If the Answer Cannot Be Found

If the tools and plans fail to deliver a final answer, inform the user. Acknowledge that a solution could not be found and explain the steps taken.

Example message: "I couldn't solve this problem. Here's the sequence of steps I took: ..."

Answer Format

When returning the FINAL ANSWER, ensure it follows this format:

FINAL ANSWER: [Your final answer here]

Use minimal words or a number. If necessary, a comma-separated list is allowed.

Do not include units, unnecessary details, or explanations unless specifically requested.

Key Focus Areas
The assistant should not merely perform tool calls but should critically evaluate its process at every step.

Keep the overall approach flexible so that if errors or issues arise, the model can adapt and improve its strategy in real-time.

Prioritize efficiency. Once the correct answer is found, avoid unnecessary further actions.

Your answer must always start with "FINAL ANSWER:".


### Key Improvements:
1. **Clearer Instructions for Iteration**: The prompt now emphasizes revising the plan based on results and ensuring that each step is critically evaluated.
2. **Final Answer Verification**: We ensure that only the final answer is provided in the correct format, and intermediate steps are excluded.
3. **Revised Plan if Initial Attempt Fails**: If the initial plan doesn't work or if the results don't meet expectations, the plan is revised and tools are adjusted accordingly.
4. **Error Handling and Clarification**: If a solution can't be found, the model now acknowledges the steps taken and explains why the answer couldn't be solved, improving transparency.

This updated prompt should make your planner node more efficient and capable of adapting to complex scenarios while keeping the process structured and robust.

if you want to return FINAL ANSWER, 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 comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, Apply the rules above for each element (number or string), ensure there is exactly one space after each comma.
Your answer should only start with "FINAL ANSWER: ", then follows with the answer.

**Return exactly** either a `FINAL ANSWER:` line or a JSON array of tool calls—nothing else.
"""



# ===  TOOL STUBS ===
@tool
def python_repl_tool(query: str) -> str:
    """
    Execute Python query like calculations; return printed output.

    Args:
      query: The calculation formula.

    Returns:
      The result of the query exacuation
    """
    python_repl = PythonREPL()
    return python_repl.run(query)


@tool
def web_search(query: str) -> str:
    """Search Tavily for a query and return maximum 3 results.
    Args:
        query: The search query.
    Returns:
        A concise details about the query from internet."""
    response = TavilySearchResults().invoke({"query": query, "max_results": 3})
    content = []
    for doc in response:
      content.append(f'<Document source="{doc["url"]}" confidence score="{doc["score"]}"/>\n{doc["content"]}\n</Document>')
    return {"web_results": content}

@tool
def vision_caption(image_path: str) -> str:
    """
    Extract text from an image using OCR library pytesseract (if available).
    Args:
        image_path (str): the path to the image file.
    Returns:
      A concise, human‑readable caption describing the image content.
    """
    try:
        # Open the image
        image = Image.open(image_path)

        # Extract text from the image
        text = pytesseract.image_to_string(image)

        return f"Extracted text from image:\n\n{text}"
    except Exception as e:
        return f"Error extracting text from image: {str(e)}"


@tool
def wiki_search(query: str) -> str:
    """Search Wikipedia for a query and return maximum 2 results.

    Args:
        query: The search query.
    Returns:
        Search results about the query from the Wikipedia website
    """
    docs = WikipediaLoader(query=query, load_max_docs=2).load()
    return "\n\n---\n\n".join(f'{doc.page_content}' for doc in docs)

@tool
def arvix_search(query: str) -> str:
    """Search Arxiv for a query and return maximum 3 result.

    Args:
        query: The search query.
    Returns:
        Search results about the query from the Arxiv website
        """
    docs = ArxivLoader(query=query, load_max_docs=2).load()
    return "\n\n---\n\n".join(f'{doc.page_content}' for doc in docs)


@tool
def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
    """
    Download a file from a URL and save it to a temporary location.
    Args:
        url (str): the URL of the file to download.
        filename (str, optional): the name of the file. If not provided, a random name file will be created.
    Returns:
      Nothing
    """
    try:
        # Parse URL to get filename if not provided
        if not filename:
            path = urlparse(url).path
            filename = os.path.basename(path)
            if not filename:
                filename = f"downloaded_{uuid.uuid4().hex[:8]}"

        # Create temporary file
        temp_dir = tempfile.gettempdir()
        filepath = os.path.join(temp_dir, filename)

        # Download the file
        response = requests.get(url, stream=True)
        response.raise_for_status()

        # Save the file
        with open(filepath, "wb") as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        return f"File downloaded to {filepath}. You can read this file to process its contents."
    except Exception as e:
        return f"Error downloading file: {str(e)}"

@tool
def file_reader(file_path: str) -> str:
    """
    Analyze a text file using pandas and answer a question about it.
    The file can be any of excel, csv, tsv, xls, json, txt or log.
    Args:
        file_path (str): the path to the CSV file.
    Returns:
      The conetent inside the PDF/Text file
    """
    try:
        _, ext = os.path.splitext(file_path.lower())
        if ext in {".csv", ".tsv"}:
            sep = "\t" if ext == ".tsv" else ","
            data = pd.read_csv(file_path, sep=sep)
        if ext in {".xls", ".xlsx"}:
            data = pd.read_excel(file_path)
        if ext == ".json":
            data = pd.read_json(file_path)
        if ext in {".txt", ".log"}:
            with open(file_path, "r", encoding="utf-8") as f:
                data = f.read()

        # Run various analyses based on the query
        result = f"text file loaded with {len(data)} rows and {len(data.columns)} columns.\n"
        result += f"Columns: {', '.join(data.columns)}\n\n"+\
                    "Summary statistics:\n" + str(data.describe())
        # Add summary statistics
        return result

    except Exception as e:
        return f"Error analyzing text file: {str(e)}"

tools = [
        python_repl_tool,
        web_search,
        vision_caption,
        arvix_search,
        wiki_search,
        download_file_from_url,
        file_reader,
        ]




class MessagesState(TypedDict):
    input_file: Optional[str]
    messages: Annotated[list[AnyMessage], add_messages]
    final_answer: Optional[str]
    is_done: Optional[bool]
    assist_iter: Optional[int]


class EsiAgent:
    def __init__(self):
        llm = ChatOpenAI(
            base_url="https://api.groq.com/openai/v1",
            openai_api_key=os.environ["GPT_4o_API_KEY"],
            model_name="meta-llama/llama-4-scout-17b-16e-instruct",
            temperature=0,
            # max_tokens=128
            max_tokens=256
        )

        llm_with_tools = llm.bind_tools(tools)
        # 3) Single system message
        self.sys_msg = SystemMessage(content=PLANNER_PROMPT)

        # 4) Assistant node
        def assistant(state: MessagesState) -> dict:
          """
          Assistant node: Validates if a final answer exists or generates the plan.
          If a final answer is found, it returns it, otherwise, it generates a plan and iterates.
          """
          llm_input = [SystemMessage(content=PLANNER_PROMPT)]+state["messages"]
          ai_resp = llm_with_tools.invoke(llm_input)
          content = ai_resp.content.strip()
    
          if "FINAL ANSWER:" in content and state['assist_iter']<=20:
              final_answer = content
              is_done = True
          else:
              final_answer = None
              is_done = False

          return {
              "messages": state["messages"] + [ai_resp],
              "input_file": state.get("input_file", None),
              "final_answer": final_answer,
              'assist_iter': state['assist_iter'] + 1,
              "is_done": is_done,
          }

        # 5) Build the graph
        builder = StateGraph(MessagesState)
        builder.add_node("assistant", assistant)
        builder.add_node("tools",     ToolNode(tools))
        builder.add_edge(START, "assistant")
        builder.add_conditional_edges(
            source="assistant",
            path=lambda s: s.get("is_done", False),
            path_map={False: "tools", True: END})

        builder.add_edge("tools", "assistant")
        self.graph = builder.compile()


    def run(self, question: str, files=None) -> str:
        """
        Wraps graph.invoke to run end-to-end on a plain question.
        """
        # Build a minimal GAIA payload (if your MessagesState expects more, adjust here)
        initial_state = {
            "messages": [HumanMessage(content=question)],
            "input_file": files,
            "final_answer": "",
            'is_done': False,
            'assist_iter': 0
        }
        final_state = self.graph.invoke(initial_state)
        last_msg = final_state["messages"][-1].content

        # If you need to strip off "FINAL ANSWER: " prefix, do it here:
        idx = last_msg.find("FINAL ANSWER:")
        if idx >= 0:
          return last_msg[idx+len("FINAL ANSWER:"):].strip()
        return last_msg.strip()

if __name__ == "__main__":
    # question = "Where is the capital of the Iran?"
    agent = EsiAgent()
    answer = agent.run("What is 15% of 200 plus 7 squared?")