Upload 2 files

gemini_agent.py

@@ -0,0 +1,131 @@
+import os
+import time
+from dotenv import load_dotenv
+from typing import TypedDict, Annotated, Optional
+
+from langgraph.prebuilt import ToolNode, tools_condition
+from langgraph.graph import StateGraph, START
+from langgraph.graph.message import add_messages
+
+from langchain_core.messages import AnyMessage, SystemMessage, HumanMessage
+from langchain_google_genai import ChatGoogleGenerativeAI
+
+from tools import *
+
+load_dotenv()
+DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+
+class AgentState(TypedDict):
+    """Agent state for the graph."""
+    input_file: Optional[str]
+    messages: Annotated[list[AnyMessage], add_messages]
+
+
+class GEMINI_AGENT:
+    def __init__(self):
+        self.llm = ChatGoogleGenerativeAI(
+            model="gemini-2.0-flash-lite",
+            temperature=0,
+            max_tokens=1024,
+            google_api_key=os.getenv("GEMINI_API_KEY"),
+        )
+        
+        self.tools = [
+            duckduck_websearch,
+            serper_websearch,
+            visit_webpage,
+            wiki_search,
+            youtube_viewer,
+            text_splitter,
+            read_file,
+            excel_read,
+            csv_read,
+            mp3_listen,
+            image_caption,
+            run_python,
+            multiply,
+            add,
+            subtract,
+            divide
+        ]
+
+        self.llm_with_tools = self.llm.bind_tools(self.tools)
+        self.app = self._graph_compile()
+
+    def _graph_compile(self):
+        builder = StateGraph(AgentState)
+        # Define nodes: these do the work
+        builder.add_node("assistant", self._assistant)
+        builder.add_node("tools", ToolNode(self.tools))
+        # Define edges: these determine how the control flow moves
+        builder.add_edge(START, "assistant")
+        builder.add_conditional_edges(
+            "assistant",
+            tools_condition,
+        )
+        builder.add_edge("tools", "assistant")
+        react_graph = builder.compile()
+        return react_graph
+
+    def _assistant(self, state: AgentState):
+        sys_msg = SystemMessage(
+            content=
+            """
+            You are a helpful assistant tasked with answering questions using a set of tools. When given a question, follow these steps:
+            1. Create a clear, step-by-step plan to solve the question.
+            2. If a tool is necessary, select the most appropriate tool based on its functionality. If one tool isn't working, use another with similar functionality.
+            3. Execute your plan and provide the response in the following format:
+
+            FINAL ANSWER: [YOUR FINAL ANSWER]
+
+            Your final answer should be:
+
+            - A number (without commas or units unless explicitly requested),
+            - A short string (avoid articles, abbreviations, and use plain text for digits unless otherwise specified),
+            - A comma-separated list (apply the formatting rules above for each element, with exactly one space after each comma).
+
+            Ensure that your answer is concise and follows the task instructions strictly. If the answer is more complex, break it down in a way that follows the format.
+            Begin your response with "FINAL ANSWER: " followed by the answer, and nothing else.
+            """
+        )
+
+        return {
+            "messages": [self.llm_with_tools.invoke([sys_msg] + state["messages"])],
+            "input_file": state["input_file"]
+        }
+
+    def extract_after_final_answer(self, text):
+        keyword = "FINAL ANSWER: "
+        index = text.find(keyword)
+        if index != -1:
+            return text[index + len(keyword):]
+        else:
+            return ""
+
+    def run(self, task: dict):
+        task_id, question, file_name = task["task_id"], task["question"], task["file_name"]
+        print(f"Agent received question (first 50 chars): {question[:50]}...")
+
+        if file_name == "" or file_name is None:
+            question_text = question
+        else:
+            question_text = f'{question} with TASK-ID: {task_id}'
+        messages = [HumanMessage(content=question_text)]
+
+        max_retries = 5
+        base_sleep = 1
+        for attempt in range(max_retries):
+            try:
+                response = self.app.invoke({"messgae": messages, "input_file": None})
+                final_ans = self.extract_after_final_answer(response)
+                time.sleep(60) # avoid rate limit
+                return final_ans
+            except Exception as e:
+                sleep_time = base_sleep * (attempt + 1)
+                if attempt < max_retries - 1:
+                    print(str(e))
+                    print(f"Attempt {attempt + 1} failed. Retrying in {sleep_time} seconds...")
+                    time.sleep(sleep_time)
+                    continue
+                return f"Error processing query after {max_retries} attempts: {str(e)}"
+        return "This is a default answer."

tools.py

@@ -0,0 +1,327 @@
+import os
+import re
+import requests
+import pandas as pd
+from typing import List
+from dotenv import load_dotenv
+
+from google import genai
+from google.genai import types
+
+from langchain_core.tools import tool
+from langchain.document_loaders import WebBaseLoader
+from langchain_experimental.tools import PythonREPLTool
+from langchain.text_splitter import CharacterTextSplitter
+from langchain_community.tools import DuckDuckGoSearchResults
+from langchain_community.retrievers import WikipediaRetriever
+from langchain_community.utilities import GoogleSerperAPIWrapper
+from langchain_community.document_loaders import ImageCaptionLoader, AssemblyAIAudioTranscriptLoader
+
+
+load_dotenv()
+DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+
+
+@tool
+def duckduck_websearch(query: str) -> str:
+    """
+    Performs a web search using the given query, downloads the content of two relevant web pages,
+    and returns their combined content as a raw string.
+
+    This is useful when the task requires analysis of web page content, such as retrieving poems, 
+    changelogs, or other textual resources.
+
+    Args:
+        query (str): The search query.
+
+    Returns:
+        str: The combined raw text content of the two retrieved web pages.
+    """
+    search_engine = DuckDuckGoSearchResults(output_format="list", num_results=2)
+    page_urls = [url["link"] for url in search_engine(query)]
+
+    loader = WebBaseLoader(web_paths=(page_urls))
+    docs = loader.load()
+
+    combined_text = "\n\n".join(doc.page_content[:15000] for doc in docs)
+
+    # Clean up excessive newlines, spaces and strip leading/trailing whitespace
+    cleaned_text = re.sub(r'\n{3,}', '\n\n', combined_text).strip()
+    cleaned_text = re.sub(r'[ \t]{6,}', ' ', cleaned_text)
+
+    # Strip leading/trailing whitespace
+    cleaned_text = cleaned_text.strip()
+    return cleaned_text
+
+
+@tool
+def serper_websearch(query: str) -> str:
+    """
+    Performs a web search using the given query with SERPER Search Engine
+
+    Args:
+        query (str): The search query.
+    
+    Returns:
+        str: the search result
+    """
+    search = GoogleSerperAPIWrapper(serper_api_key=os.getenv("SERPER_API_KEY"))
+    results = search.run(query)
+    return results
+
+@tool
+def visit_webpage(url: str) -> str:
+    """
+    Fetches raw HTML content of a web page.
+    
+    Args:
+        url: the webpage url
+    
+    Returns:
+        str: The combined raw text content of the webpage
+    """
+    try:
+        response = requests.get(url, timeout=5)
+        return response.text[:5000]
+    except Exception as e:
+        return f"[ERROR fetching {url}]: {str(e)}"
+
+@tool
+def wiki_search(query: str) -> str:
+    """
+    Searches for a Wikipedia articles using the provided query and returns the content of the corresponding Wikipedia pages.
+
+    Args:
+        query (str): The search term to look up on Wikipedia.
+
+    Returns:
+        str: The text content of the Wikipedia articles related to the query.
+    """
+    retriever = WikipediaRetriever()
+    docs = retriever.invoke(query)
+    combined_text = "\n\n".join(doc.page_content for doc in docs)
+    return combined_text
+
+@tool
+def youtube_viewer(youtube_url: str, question: str) -> str:
+    """
+    Analyzes a YouTube video from the provided URL and returns an answer 
+    to the given question based on the analysis results.
+
+    Args:
+        youtube_url (str): The URL of the YouTube video, in the format 
+            "https://www.youtube.com/...".
+        question (str): A question related to the content of the video.
+
+    Returns:
+        str: An answer to the question based on the video's content.
+    """
+    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+    response = client.models.generate_content(
+        model='models/gemini-2.5-flash-preview-04-17',
+        contents=types.Content(
+            parts=[
+                types.Part(
+                    file_data=types.FileData(file_uri=youtube_url)
+                ),
+                types.Part(text=question)
+            ]
+        )
+    )
+    return response.text
+
+@tool
+def text_splitter(text: str) -> List[str]:
+    """
+    Splits text into chunks using LangChain's CharacterTextSplitter.
+    
+    Args:
+        text: A string of text to split.
+    
+    Returns:
+        List[str]: a list of split text
+    """
+    splitter = CharacterTextSplitter(chunk_size=450, chunk_overlap=10)
+    return splitter.split_text(text)
+
+@tool
+def read_file(task_id: str) -> str:
+    """
+    First download the file, then read its content
+    
+    Args:
+        dir: the task_id
+    
+    Returns:
+        str: the file content
+    """
+    file_url = f'{DEFAULT_API_URL}/files/{task_id}'
+    r = requests.get(file_url, timeout=15, allow_redirects=True)
+    with open('temp', "wb") as fp:
+        fp.write(r.content)
+    with open('temp') as f:
+        return f.read()
+
+@tool
+def excel_read(task_id: str) -> str:
+    """
+    First download the excel file, then read its content
+    
+    Args:
+        dir: the task_id
+    
+    Returns:
+        str: the content of excel file
+    """
+    try:
+        file_url = f'{DEFAULT_API_URL}/files/{task_id}'
+        r = requests.get(file_url, timeout=15, allow_redirects=True)
+        with open('temp.xlsx', "wb") as fp:
+            fp.write(r.content)
+        # Read the Excel file
+        df = pd.read_excel('temp.xlsx')
+        # Run various analyses based on the query
+        result = (
+            f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
+        )
+        result += f"Columns: {', '.join(df.columns)}\n\n"
+        # Add summary statistics
+        result += "Summary statistics:\n"
+        result += str(df.describe())
+        return result
+    except Exception as e:
+        return f"Error analyzing Excel file: {str(e)}"
+   
+@tool
+def csv_read(task_id: str) -> str:
+    """
+    First download the csv file, then read its content
+    
+    Args:
+        dir: the task_id
+    
+    Returns:
+        str: the content of csv file
+    """
+    try:
+        file_url = f'{DEFAULT_API_URL}/files/{task_id}'
+        r = requests.get(file_url, timeout=15, allow_redirects=True)
+        with open('temp.csv', "wb") as fp:
+            fp.write(r.content)
+        # Read the CSV file
+        df = pd.read_csv('temp.csv')
+        # Run various analyses based on the query
+        result = (
+            f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
+        )
+        result += f"Columns: {', '.join(df.columns)}\n\n"
+        # Add summary statistics
+        result += "Summary statistics:\n"
+        result += str(df.describe())
+        return result
+    except Exception as e:
+        return f"Error analyzing CSV file: {str(e)}"
+
+
+@tool
+def mp3_listen(task_id: str) -> str:
+    """
+    First download the mp3 file, then listen to it
+    
+    Args:
+        dir: the task_id
+    
+    Returns:
+        str: the content of mp3 file
+    """
+    file_url = f'{DEFAULT_API_URL}/files/{task_id}'
+    r = requests.get(file_url, timeout=15, allow_redirects=True)
+    with open('temp.mp3', "wb") as fp:
+        fp.write(r.content)
+    loader = AssemblyAIAudioTranscriptLoader(file_path="temp.mp3", api_key=os.getenv("AssemblyAI_API_KEY"))
+    docs = loader.load()
+    contents = [doc.page_content for doc in docs]
+    return "\n".join(contents)
+    
+
+@tool
+def image_caption(dir: str) -> str:
+    """
+    Understand the content of the provided image
+    
+    Args:
+        dir: the image url link
+    
+    Returns:
+        str: the image caption
+    """
+    loader = ImageCaptionLoader(images=[dir])
+    metadata = loader.load()
+    return metadata[0].page_content
+
+
+@tool
+def run_python(code: str):
+    """ Run the given python code
+    
+    Args:
+        code: the python code
+    """
+    return PythonREPLTool().run(code)
+
+@tool
+def multiply(a: float, b: float) -> float:
+    """
+    Multiply two numbers.
+    
+    Args:
+        a: first float
+        b: second float
+    
+    Returns:
+        float: the multiplication of a and b
+    """
+    return a * b
+
+@tool
+def add(a: float, b: float) -> float:
+    """
+    Add two numbers.
+    
+    Args:
+        a: first float
+        b: second float
+    
+    Returns:
+        float: the sum of a and b
+    """
+    return a + b
+
+@tool
+def subtract(a: float, b: float) -> float:
+    """
+    Subtract two numbers.
+    
+    Args:
+        a: first float
+        b: second float
+    
+    Returns:
+        float: the result after a subtracted by b
+    """
+    return a - b
+
+@tool
+def divide(a: float, b: float) -> float:
+    """Divide two numbers.
+    
+    Args:
+        a: first float
+        b: second float
+    
+    Returns:
+        float: the result after a divided by b
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b