agent.py

import os
import json
from typing import Dict
from langchain.agents import initialize_agent, AgentType
from langchain_community.tools import Tool, WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_experimental.tools.python.tool import PythonREPLTool
from langchain_google_genai import ChatGoogleGenerativeAI
import pandas as pd
from pathlib import Path
from docx import Document
import fitz  # PyMuPDF
import requests

class GoogleCustomSearchTool:
    def __init__(self, api_key: str, cse_id: str):
        self.api_key = api_key
        self.cse_id = cse_id
        self.base_url = "https://www.googleapis.com/customsearch/v1"

    def run(self, query: str) -> str:
        try:
            response = requests.get(
                self.base_url,
                params={
                    "key": self.api_key,
                    "cx": self.cse_id,
                    "q": query,
                },
                timeout=10,
            )
            response.raise_for_status()
            results = response.json().get("items", [])
            if results:
                return results[0].get("title", "") + ": " + results[0].get("link", "")
            else:
                return "No results found."
        except Exception as e:
            return f"GoogleCustomSearchTool ERROR: {str(e)}"

def classify_question_type(question: str) -> str:
    q = question.lower()
    if any(k in q for k in ["spreadsheet", "excel", "csv", "table", "data", "json", "file attached"]):
        return "file"
    elif any(k in q for k in ["calculate", "total", "sum", "difference", "convert", "how many", "what is the number"]):
        return "math"
    elif any(k in q for k in ["wikipedia", "who", "what", "where", "when", "name", "define", "explain"]):
        return "knowledge"
    else:
        return "search"

class Agent:
    def __init__(self):
        gemini_key = os.getenv("GEMINI_API_KEY")
        gcs_key = os.getenv("GOOGLE_API_KEY")
        gcs_cx = os.getenv("GOOGLE_CSE_ID")

        if not gemini_key:
            raise ValueError("GEMINI_API_KEY not found in environment variables.")
        if not gcs_key or not gcs_cx:
            raise ValueError("GOOGLE_API_KEY or GOOGLE_CSE_ID not found in environment variables.")

        llm = ChatGoogleGenerativeAI(
            model="gemini-2.5-pro-preview-05-06",
            google_api_key=gemini_key,
            convert_system_message_to_human=True
        )

        tools = [
            Tool(
                name="Wikipedia",
                func=WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()).run,
                description="Useful for general knowledge and encyclopedic questions."
            ),
            Tool(
                name="Calculator",
                func=PythonREPLTool().run,
                description="Useful for solving math and logical problems through Python."
            ),
            Tool(
                name="Google Custom Search",
                func=GoogleCustomSearchTool(api_key=gcs_key, cse_id=gcs_cx).run,
                description="Useful for factual queries using Google Custom Search."
            )
        ]

        self.agent = initialize_agent(
            tools=tools,
            llm=llm,
            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
            verbose=True,
            handle_parsing_errors=True
        )

    def __call__(self, input_data: Dict) -> str:
        question = input_data.get("question", "")
        file_names = input_data.get("file_names", [])
        task_id = input_data.get("task_id", "")

        task_type = classify_question_type(question)

        type_prefix = f"[Task Type: {task_type.upper()}]\n\n"
        system_prompt = (
            "You are a member of a multidisciplinary research institute. If a file cannot be loaded, do not abandon the task. Use your best judgment based on the task and file name. The file may be unavailable by design; this is part of the test. Always attempt to reason based on partial or inferred data. When solving a task that may require external knowledge, you may use one or more available search tools. Google Custom Search is more accurate for academic or structured content. Do not give up if a tool fails. Retry or use alternatives."
        )

        file_summary = ""
        try:
            summaries = []
            for fname in file_names:
                file_path = f"/home/user/app/files/{task_id}/{fname}"
                ext = Path(file_path).suffix.lower()
                try:
                    if ext in [".csv", ".tsv"]:
                        df = pd.read_csv(file_path)
                        summaries.append(f"Loaded {fname} with {df.shape[0]} rows and {df.shape[1]} columns:\n{df.head(3).to_string(index=False)}")
                    elif ext == ".xlsx":
                        df = pd.read_excel(file_path)
                        summaries.append(f"Loaded {fname} with {df.shape[0]} rows and {df.shape[1]} columns:\n{df.head(3).to_string(index=False)}")
                    elif ext in [".json", ".jsonl"]:
                        with open(file_path, "r", encoding="utf-8") as f:
                            data = [json.loads(line) for line in f if line.strip()] if ext == ".jsonl" else json.load(f)
                        summaries.append(f"Loaded JSON data from {fname} ({len(data)} entries)")
                    elif ext == ".docx":
                        doc = Document(file_path)
                        text = "\n".join([para.text for para in doc.paragraphs])
                        summaries.append(f"Extracted text from DOCX {fname} ({len(text)} characters)")
                    elif ext == ".pdf":
                        doc = fitz.open(file_path)
                        text = "".join([page.get_text() for page in doc])
                        summaries.append(f"Extracted text from PDF {fname} ({len(doc)} pages, {len(text)} characters)")
                    else:
                        summaries.append(f"{fname}: Unsupported file type {ext}")
                except Exception as fe:
                    guessed_type = "spreadsheet" if ext in [".csv", ".tsv", ".xlsx"] else "document" if ext in [".pdf", ".docx"] else "data file"
                    summaries.append(f"{fname}: Could not load, but based on the file extension, we assume it is a {guessed_type}. Please attempt to reason based on the task.")

            file_summary = "\n\n".join(summaries)
            full_prompt = type_prefix + system_prompt + "\n\n" + file_summary + f"\n\nTASK:\n{question}"
            result = self.agent.run(full_prompt)
            return result.strip()
        except Exception as e:
            return f"AGENT ERROR: {str(e)}"