agent.py

import os
from dotenv import load_dotenv
from langgraph.graph import START, StateGraph, MessagesState
from langgraph.prebuilt import tools_condition, ToolNode
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.document_loaders import WikipediaLoader, JSONLoader
from langchain_community.vectorstores import Chroma
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.tools import tool
from langchain.tools.retriever import create_retriever_tool
from PIL import Image
import pandas as pd
import numpy as np
import google.generativeai as genai
import subprocess


load_dotenv()

genai.configure(api_key=os.environ["GOOGLE_API_KEY"])


@tool
def wiki_search(query: str) -> str:
    """
    Search Wikipedia for the content of a specific article. Use this tool to get the current information, facts, and data from a Wikipedia page.
    This tool is NOT for finding out about Wikipedia's edit history or discussions; use web_search for that.

    Args:
        query: The search query, ideally the exact title of the Wikipedia page.
    """
    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
            for doc in search_docs
        ])
    return formatted_search_docs


tavily_search = TavilySearchResults(max_results=3)


@tool
def web_search(query: str) -> str:
    """
    Search the general web for information, news, or to answer questions that require up-to-date information or knowledge about a page's history (like edit history).
    Use this tool when you need to find information that is not the content of a specific page, for example, 'who nominated a specific article' or 'when was a page created'.

    Args:
        query: The search query.
    """
    search_docs = tavily_search.invoke(query)
    formatted_search_docs = "\n\n---\n\n".join(
        [
            f'<Document source="{doc["url"]}">{doc["content"]}</Document>'
            for doc in search_docs
        ]
    )
    return formatted_search_docs


@tool
def process_image(image_path: str, question: str) -> str:
    """
    Process an image file and answer questions about it. Use this when a question refers to an image.
    The `image_path` is the name of the file (e.g., 'image.png') provided in the question's 'file_name'.

    Args:
        image_path: Path to the image file.
        question: Question about the image.
    """
    try:
        model = genai.GenerativeModel('gemini-2.5-flash')
        image = Image.open(image_path)
        response = model.generate_content([question, image])
        return response.text
    except Exception as e:
        return f"Error processing image: {str(e)}"


@tool
def process_audio(audio_path: str, question: str) -> str:
    """
    Process an audio file and extract information. Use this when a question refers to an audio file (e.g., .mp3, .wav).
    The `audio_path` is the name of the file (e.g., 'audio.mp3') provided in the question's 'file_name'.

    Args:
        audio_path: Path to the audio file.
        question: Question about the audio content.
    """
    try:
        model = genai.GenerativeModel('gemini-2.5-flash')
        audio_file = genai.upload_file(path=audio_path)
        response = model.generate_content([question, audio_file])
        return response.text
    except Exception as e:
        return f"Error processing audio: {str(e)}"


@tool
def process_excel_file(file_path: str, question: str) -> str:
    """
    Process an Excel file by loading it into a pandas DataFrame and using code to answer a question about it.
    Use this for complex queries about data in an Excel file (.xlsx).
    The `file_path` is the name of the file provided in the question's 'file_name'.

    Args:
        file_path: Path to the Excel file.
        question: The question to answer about the Excel data.
    """
    try:
        df = pd.read_excel(file_path)
        code_gen_llm = ChatGoogleGenerativeAI(
            model="gemini-2.5-flash", temperature=0)

        prompt = f"""
You are an expert in pandas. You are given a pandas DataFrame named `df`.
The user has the following question about the data:
"{question}"

The DataFrame has the following columns: {list(df.columns)}
And here is the head of the DataFrame:
{df.head().to_string()}

Write a short Python script that uses the `df` DataFrame to answer the question.
The script must calculate the answer and print it.
Your code must not contain any explanation or markdown formatting.
Assume `df` is already loaded. For example: `print(df['Sales'].sum())`
"""

        code_response = code_gen_llm.invoke(prompt)
        generated_code = code_response.content.strip().replace(
            "```python", "").replace("```", "")

        from io import StringIO
        import sys

        old_stdout = sys.stdout
        redirected_output = sys.stdout = StringIO()

        local_scope = {'df': df, 'pd': pd, 'np': np}
        exec(generated_code, globals(), local_scope)

        sys.stdout = old_stdout

        result = redirected_output.getvalue().strip()
        if not result:
            return "The code executed but produced no output."
        return f"The answer to '{question}' is: {result}"

    except Exception as e:
        return f"Error processing Excel file '{file_path}': {str(e)}"


@tool
def execute_python_code(code_path: str) -> str:
    """
    Execute a Python file and return the output. Use this when a question refers to a Python code file (.py).
    The `code_path` is the name of the file provided in the question's 'file_name'.

    Args:
        code_path: Path to the Python file.
    """
    try:
        result = subprocess.run(['python', code_path],
                                capture_output=True, text=True, check=True)
        return f"Output: {result.stdout}\nErrors: {result.stderr}"
    except Exception as e:
        return f"Error executing Python code: {str(e)}"


@tool
def reverse_text(text: str) -> str:
    """
    Reverse the given text. Use this for questions that require reversing a string.

    Args:
        text: Text to reverse.
    """
    return text[::-1]


@tool
def analyze_text_pattern(text: str) -> str:
    """
    Analyze text patterns and solve text puzzles. Use this for complex text manipulation that is not simple reversal.

    Args:
        text: Text to analyze.
    """
    words = text.split()
    reversed_words = [word[::-1] for word in words]
    reversed_sentence = ' '.join(reversed_words[::-1])

    analysis = f"Original: {text}\n"
    analysis += f"Reversed sentence: {reversed_sentence}\n"
    analysis += f"Word-by-word reverse: {' '.join(reversed_words)}\n"

    return analysis


@tool
def youtube_video_info(video_url: str, question: str) -> str:
    """
    Get information about a YouTube video, like its transcript or a summary, by searching the web.
    Use this tool when the question involves watching or analyzing a YouTube video.

    Args:
        video_url: YouTube video URL.
        question: Question about the video.
    """
    try:
        search_query = f"transcript of youtube video {video_url} {question}"
        search_results = tavily_search.invoke(search_query)
        if search_results:
            formatted_search_docs = "\n\n---\n\n".join(
                [
                    f'<Document source="{doc["url"]}">{doc["content"]}</Document>'
                    for doc in search_results
                ]
            )
            return f"Found information about the video {video_url}:\n\n{formatted_search_docs}"
        return f"Could not find a transcript or information for the video: {video_url}"
    except Exception as e:
        return f"Error processing video URL {video_url}: {str(e)}"


@tool
def mathematical_analysis(expression: str) -> str:
    """
    Evaluates a mathematical expression and returns the result.
    Can perform basic arithmetic. For more complex problems like analyzing tables,
    the model should break down the problem into smaller calculations.

    Args:
        expression: A string containing a mathematical expression to be evaluated.
    """
    try:
        result = eval(expression)
        return f"The result of the expression '{expression}' is: {result}"
    except Exception as e:
        return f"Could not evaluate the mathematical expression. Error: {str(e)}. Please provide a standard Python mathematical expression."


# load the system prompt from the file
with open("system_prompt.txt", "r", encoding="utf-8") as f:
    system_prompt = f.read()

# System message
sys_msg = SystemMessage(content=system_prompt)

# # build a retriever
# embeddings = GoogleGenerativeAIEmbeddings(
#     model="models/embedding-001")

# # Load documents from metadata.jsonl
# loader = JSONLoader(
#     file_path='./metadata.jsonl',
#     jq_schema='.',
#     json_lines=True,
#     text_content=False)
# documents = loader.load()

# # Create or load the Chroma vector store
# persist_directory = "chroma_db_persistent"

# if os.path.exists(persist_directory) and os.listdir(persist_directory):
#     vector_store = Chroma(
#         persist_directory=persist_directory,
#         embedding_function=embeddings
#     )
# else:
#     vector_store = Chroma.from_documents(
#         documents=documents,
#         embedding=embeddings,
#         persist_directory=persist_directory
#     )

# retriever_tool = create_retriever_tool(
#     retriever=vector_store.as_retriever(),
#     name="question_search",
#     description="A tool to retrieve similar questions from a vector store to use as reference.",
# )


tools = [
    wiki_search,
    web_search,
    process_image,
    process_audio,
    process_excel_file,
    execute_python_code,
    reverse_text,
    analyze_text_pattern,
    youtube_video_info,
    mathematical_analysis,
    # retriever_tool,
]


# Build graph function
def build_graph():
    """Build the graph"""
    llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0)
    llm_with_tools = llm.bind_tools(tools)

    def assistant(state: MessagesState):
        """Assistant node"""
        return {"messages": [llm_with_tools.invoke(state["messages"])]}

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

    return builder.compile()


# # Wrapper class to comply with the Hugging Face evaluation format
# class Agent:
#     """
#     Wrapper for the LangGraph agent to be compatible with the evaluation script.
#     The evaluation script expects an object named 'agent' with a 'run' method.
#     """

#     def __init__(self):
#         self.graph = build_graph()

#     def run(self, question: str, file_name: str = None) -> str:
#         """
#         Invokes the graph and returns the final answer as a string.
#         """
#         try:
#             # Append file information to the question if a file_name is provided
#             if file_name:
#                 question += f"\n\n[Additional context: The question refers to the file named '{file_name}']"

#             messages = [sys_msg, HumanMessage(content=question)]
#             result = self.graph.invoke({"messages": messages})
#             final_answer = result["messages"][-1].content
#             return final_answer
#         except Exception as e:
#             print(f"Error during agent execution: {e}")
#             return "This is a default answer due to an error."


# # Instantiate the agent for the evaluation script to import
# agent = Agent()


# # test
# if __name__ == "__main__":
#     # Example question about an image
#     # image_question = "How many at bats did the Yankee with the most walks in the 1977 regular season have that same season?"
#     # image_question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
#     image_question = "Who nominated the only Featured Article on English Wikipedia about a dinosaur that was promoted in November 2016?"
#     # image_file = "cca530fc-4052-43b2-b130-b30968d8aa44.png"

#     print("--- Running Test ---")
#     # The evaluation script would call agent.run(question, file_name)
#     # We simulate that here.
#     answer = agent.run(question=image_question)
#     print(f"Question: {image_question}")
#     # print(f"File: {image_file}")
#     print(f"Answer: {answer}")
#     print("--- Test Complete ---")