agent.py

import base64
import os
import io
import contextlib
import requests
from typing import TypedDict, Annotated
from langchain_core.messages import AnyMessage
from langgraph.graph import START, StateGraph, add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_community.tools import tool, DuckDuckGoSearchRun
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain_google_genai import ChatGoogleGenerativeAI
# from pathlib import Path
# import tempfile
from dotenv import load_dotenv
import time
import random

# constants
API_URL = "https://agents-course-unit4-scoring.hf.space"
QUESTIONS_URL = f"{API_URL}/questions"
FILES_URL = f"{API_URL}/files"
SUBMIT_URL = f"{API_URL}/submit"
load_dotenv()

class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    # file_path: str | None
    # task_id: str | None
    # url: str | None

def build_gemini_llm():
    if not os.environ.get("GOOGLE_API_KEY"):
        raise ValueError("GOOGLE_API_KEY environment variable is not set.")
    return ChatGoogleGenerativeAI(model = "gemini-3.7-flash", temperature = 0, max_output_tokens = 1025, include_thoughts=True)

# @tool
# def extract_text_from_image(img_path: str) -> str:
#     """
#     Describe the image and extract any text in it.
#     Args:
#         img_path (str): the path to the image file.
#     """
#     all_text = ""
#     try:
#         # Read image and encode as base64
#         with open(img_path, "rb") as image_file:
#             image_bytes = image_file.read()

#         image_base64 = base64.b64encode(image_bytes).decode("utf-8")

#         # Prepare the prompt including the base64 image data
#         message = [
#             HumanMessage(
#                 content=[
#                     {
#                         "type": "text",
#                         "text": (
#                             "Describe the image and extract any text in it."
#                         ),
#                     },
#                     {
#                         "type": "image_url",
#                         "image_url": {
#                             "url": f"data:image/png;base64,{image_base64}"
#                         },
#                     },
#                 ]
#             )
#         ]
#         response = model.invoke(message)
#         # Append extracted text
#         all_text += response.text + "\n\n"
#         return all_text.strip()
#     except Exception as e:
#         # A butler should handle errors gracefully
#         error_msg = f"Error extracting text: {str(e)}"
#         print(error_msg)
#         return ""

# @tool
# def download_and_read_file(task_id: str) -> str:
#     """
#     Download and read the file attached to the GAIA task its contents.
#     Always call this first if there is a file attached to a GAIA Task.
#     Args:
#         task_id (str): The ID of the GAIA task.
#     Returns:
#         str: The contents of the file as a string.
#     """

#     try:
#         # Download the file from the GAIA API
#         response = requests.get(f"{FILES_URL}/{task_id}", timeout = 10)
#         response.raise_for_status()

#         # Determine the file type and read its contents
#         content_disposition = response.headers.get("content-disposition", "")
#         content_type = response.headers.get("content-type", "")
#         filename = None
#         if "filename=" in content_disposition:
#             filename = content_disposition.split("filename=")[1].strip('"')

#         if not filename:
#             filename = f"{task_id}.bin"

#         ext = Path(filename).suffix.lower()

#         if ext in(".txt", ".py", ".json", ".md", ".ymal", ".html", ".xml", ""):
#             return response.text

#         if ext == ".xlsx" or "xlsx" in content_type:
#             import pandas as pd
#             with tempfile.NamedTemporaryFile(suffix=".xlsx", delete=False) as file:
#                 file.write(response.content)
#                 temp_path = file.name

#             read_file = pd.read_excel(temp_path)
#             return read_file.to_string()
#         if ext == ".csv" or "csv" in content_type:
#             import pandas as pd
#             with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as file:
#                 file.write(response.content)
#                 temp_path = file.name
#             read_file = pd.read_csv(temp_path)
#             return read_file.to_string()
#         if ext == ".csv" or "csv" in content_type:
#             import pandas as pd
#             with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as file:
#                 file.write(response.content)
#                 temp_path = file.name
#             read_file = pd.read_csv(temp_path)
#             return read_file.to_string()
#         if ext == ".jpg" or ext == ".jpeg" or ext == ".png" or "image" in content_type:
#             from PIL import Image
#             with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as file:
#                 file.write(response.content)
#                 temp_path = file.name
#             return extract_text_from_image(temp_path)
        
#         # Unsupported file type 
#         return ( 
#             f"Unsupported file type: {content_type}. "
#             "I downloaded the file successfully, but I don't know "
#             "how to extract its contents." 
#             ) 
#     except requests.RequestException as e: 
#         return f"Failed to download file: {e}" 
#     except Exception as e: 
#         return f"Failed to read file: {e}"
#     except Exception as e:
#         return f"error downloading or reading file: {str(e)}"

@tool
def wikipedia_search(query: str) -> str:
    """
     Search Wikipedia for factual and encyclopedic information.

    Use this tool FIRST for:
    - people
    - historical events
    - countries and places
    - musicians, artists, movies, books
    - scientific concepts
    - biographies
    - general factual knowledge

    Wikipedia is preferred for stable, well-known topics.
    Args:
        query: Keywords to search on Wikipedia.
    """

    url = "https://en.wikipedia.org/w/api.php"
    params = {
        "action": "query",
        "list": "search",
        "srsearch": query,
        "format": "json",
        "srlimit": 3,
        "utf8": 1,
    }
    headers = {
        "User-Agent": "MyLangGraphAgent/1.0"
    }

    max_retries = 3
    for attempt in range(max_retries):
        try:
            response = requests.get(
                url,
                params=params,
                headers=headers,
                timeout=10,
            )
            print(
                f"Wikipedia status={response.status_code}, "
                f"content-type={response.headers.get('content-type')}"
            )

            response.raise_for_status()
            
            data = response.json()
            results = data.get("query", {}).get("search", [])
            if not results:
                return f"No Wikipedia results found for '{query}'."
            MAX_CONTENT_LENGTH = 3000
            return "\n\n---\n\n".join(
                f"Title: {item['title']}\n"
                f"Snippet: {item.get('snippet', '')}"
                for item in results
            )
        except requests.exceptions.RequestException as e:
            print(
                f"Wikipedia request failed "
                f"attempt {attempt + 1}/{max_retries}: {e}"
            )
        except requests.exceptions.JSONDecodeError:
            print(
                f"Wikipedia returned non-JSON response. "
                f"Status={response.status_code}"
            )
            print("Response preview:")
            print(response.text[:500])

        if attempt < max_retries - 1:
            delay = 2 ** attempt + random.random()
            print(f"Retrying in {delay:.2f}s")
            time.sleep(delay)
    return (
        f"Wikipedia search temporarily failed for '{query}'. "
        "Please use another search source."
    )

# fix wikipedia engine builds invalid URL for region="wt-wt"(default) issue from duckducksearchrun
search_wrapper = DuckDuckGoSearchAPIWrapper(
    region="us-en",
    backend="duckduckgo",
)
search_ddgs = DuckDuckGoSearchRun(
    api_wrapper=search_wrapper
)

@tool
def search_web(query: str) -> str:
    """
    Search the public web for information not suitable for Wikipedia.

    Use this tool for:
    - recent news
    - current events
    - latest information
    - official websites
    - product information
    - information that may have changed recently

    Do NOT use this tool as the first choice for stable
    encyclopedic information that can be found on Wikipedia.
    Args:
        query: Keywords to search, only keywords and spaces.
    """
    try:
        result = search_ddgs.invoke(query)
        if not result:
            return f"No web results found for: {query}"
        print("ddgs result:")
        print(result)
        return result
    except Exception as e:
        print(f"DuckDuckGo search failed for {query}: {e}")
        return (
            f"Web search failed for query: {query}\n"
            f"Error: {type(e).__name__}: {e}\n"
            "Please try a different search query."
        )

model = build_gemini_llm()
tools = [
    search_web,
    wikipedia_search
]
model_with_tools = model.bind_tools(tools)

def assistant(state: AgentState):
    response = model_with_tools.invoke(state["messages"])
    print("\n===== Thinking Section =====")
    print(response.content[0].get("thinking").strip())
    print("===============================\n")
    print("\n===== ASSISTANT RESPONSE =====")
    print("TYPE:", type(response))
    print("CONTENT:", response.content)
    print("TOOL CALLS:", response.tool_calls)
    print("===============================\n")
    return {
        "messages": [response],
        # "file_path": state["file_path"],
        # "task_id": state["task_id"],
        # "url": state["url"]
    }

builder = StateGraph(AgentState)

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

graph = builder.compile()