Create agent.py

agent.py

@@ -0,0 +1,254 @@
+import base64
+import os
+import io
+import contextlib
+import requests
+from typing import TypedDict, Annotated
+from langchain_core.messages import HumanMessage, 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.document_loaders import WikipediaLoader
+from langchain_google_genai import ChatGoogleGenerativeAI
+from pathlib import Path
+import tempfile
+
+# 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"
+
+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", temeperature = 0, max_output_tokens = 1024)
+
+@tool
+def add_numbers(a: int, b: int) -> int:
+    """
+    Adds two numbers and return the result.
+    Args:
+        a (int)
+        b (int)
+    """
+    return a + b
+
+@tool
+def subtract_numbers(a: int, b: int) -> int:
+    """
+    Subtracts the second number from the first and return the result.
+    Args:
+        a (int)
+        b (int)
+    """
+    return a - b
+
+@tool
+def multiply_numbers(a: int, b: int) -> int:
+    """
+    Multiplies two numbers and return the result.
+    Args:
+        a (int)
+        b (int)
+    """
+    return a * b
+
+@tool
+def divide_numbers(a: int, b: int) -> float:
+    """
+    Divides the first number by the second and return the result.
+    Args:
+        a (int)
+        b (int)
+    """
+    return a / b
+
+@tool
+def search_web(query:str) -> str:
+        """
+        Searches the web for the answer to a given question or topic.
+        Args:
+            query (str): the question or topic to search for.
+        """
+        return DuckDuckGoSearchRun().run(query)
+
+@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)}"
+
+    return file_content
+
+@tool
+def wikipedia_search(query: str):
+    """
+    Search wikipedia for a query and return a max of three results.
+    Takes a string query as the search query
+    Args:
+        query (str): The search query.
+    """
+    try:
+        search_results = WikipediaLoader(query=query, load_max_docs=3).load()
+
+        if not search_results:
+            return f"No Wikipedia results found for {query}. Consider another query or try a web search."
+        return "\n\n---\n\n".join(
+            f"Title: {doc.metadata.get('title', 'Unknown')}\n"
+            f"Content: {doc.page_content}"
+            for doc in search_results
+        )
+    except Exception as e:
+        print(f"Wikipedia search failed for {query}: {e}")
+        return f"Wikipedia search failed for {query}. Try a web search instead."
+
+
+model = build_gemini_llm()
+tools = [
+    add_numbers,
+    subtract_numbers,
+    multiply_numbers,
+    divide_numbers,
+    extract_text_from_image,
+    wikipedia_search,
+    download_and_read_file,
+    search_web
+]
+model_with_tools = model.bind_tools(tools)
+
+def assistant(state: AgentState):
+    return {
+        "messages": state["messages"],
+        "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()