Upload agent.py

agent.py

@@ -0,0 +1,242 @@
+"""LangGraph Agent"""
+import os
+import requests as req
+from dotenv import load_dotenv
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_openai import ChatOpenAI
+from langchain_groq import ChatGroq
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
+from langchain_community.vectorstores import SupabaseVectorStore
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import tool
+from langchain.tools.retriever import create_retriever_tool
+from supabase.client import Client, create_client
+import pytesseract
+from PIL import Image
+
+load_dotenv()
+
+# ─── Math tools ────────────────────────────────────────────────────────────────
+
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> float:
+    """Divide two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Get the modulus of two numbers.
+    Args:
+        a: first int
+        b: second int
+    """
+    return a % b
+
+# ─── Search tools ──────────────────────────────────────────────────────────────
+
+@tool
+def wiki_search(query: str) -> dict:
+    """Search Wikipedia for a query and return maximum 2 results.
+    Args:
+        query: The search query.
+    """
+    search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+    formatted = "\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 {"wiki_results": formatted}
+
+@tool
+def web_search(query: str) -> dict:
+    """Search Tavily for a query and return maximum 3 results.
+    Args:
+        query: The search query.
+    """
+    search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+    formatted = "\n\n---\n\n".join(
+        f'<Document source="{doc["url"]}"/>\n{doc["content"]}\n</Document>'
+        for doc in search_docs
+    )
+    return {"web_results": formatted}
+
+@tool
+def arxiv_search(query: str) -> dict:
+    """Search Arxiv for a query and return maximum 3 results.
+    Args:
+        query: The search query.
+    """
+    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+    formatted = "\n\n---\n\n".join(
+        f'<Document source="{doc.metadata["source"]}"/>\n{doc.page_content[:1000]}\n</Document>'
+        for doc in search_docs
+    )
+    return {"arxiv_results": formatted}
+
+# ─── Image / OCR tools ─────────────────────────────────────────────────────────
+
+@tool
+def download_image(task_id: str) -> dict:
+    """Download an image file attached to a task from the evaluation API.
+    Use this first when the question mentions an image or attached file.
+    Args:
+        task_id: The task ID whose file should be downloaded.
+    """
+    api_url = "https://agents-course-unit4-scoring.hf.space"
+    file_url = f"{api_url}/files/{task_id}"
+    try:
+        response = req.get(file_url, timeout=30)
+        response.raise_for_status()
+        # Detect extension from Content-Type header
+        content_type = response.headers.get("content-type", "")
+        ext_map = {
+            "image/jpeg": ".jpg",
+            "image/png": ".png",
+            "image/gif": ".gif",
+            "image/webp": ".webp",
+        }
+        ext = next((v for k, v in ext_map.items() if k in content_type), ".jpg")
+        save_path = f"/tmp/{task_id}{ext}"
+        with open(save_path, "wb") as f:
+            f.write(response.content)
+        return {"image_path": save_path, "message": f"Image saved to {save_path}"}
+    except Exception as e:
+        return {"error": str(e)}
+
+@tool
+def extract_text_from_image(image_path: str, language: str = "eng") -> dict:
+    """Extract text from a local image file using OCR (pytesseract).
+    Use this after download_image to read text inside an image.
+    Args:
+        image_path: Local path to the image file (e.g. /tmp/task123.jpg).
+        language: Tesseract language code — 'eng' for English, 'vie' for Vietnamese.
+    """
+    try:
+        image = Image.open(image_path)
+        text = pytesseract.image_to_string(image, lang=language)
+        return {"ocr_result": text.strip()}
+    except Exception as e:
+        return {"error": str(e)}
+
+# ─── System prompt ─────────────────────────────────────────────────────────────
+
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
+
+sys_msg = SystemMessage(content=system_prompt)
+
+# ─── Vector store / retriever ──────────────────────────────────────────────────
+
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+supabase: Client = create_client(
+    os.environ.get("SUPABASE_URL"),
+    os.environ.get("SUPABASE_SERVICE_KEY"),
+)
+vector_store = SupabaseVectorStore(
+    client=supabase,
+    embedding=embeddings,
+    table_name="documents",
+    query_name="match_documents_langchain",
+)
+retriever_tool = create_retriever_tool(
+    retriever=vector_store.as_retriever(),
+    name="Question Search",
+    description="Retrieve similar questions from the vector store.",
+)
+
+# ─── Tool list ─────────────────────────────────────────────────────────────────
+
+tools = [
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulus,
+    wiki_search,
+    web_search,
+    arxiv_search,
+    download_image,
+    extract_text_from_image,
+]
+
+# ─── Graph builder ─────────────────────────────────────────────────────────────
+
+def build_graph(provider: str = "groq"):
+    """Build the LangGraph agent graph."""
+    if provider == "google":
+        llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
+    elif provider == "groq":
+        llm = ChatGroq(model="qwen-qwq-32b", temperature=0)
+    elif provider == "openai":
+        llm = ChatOpenAI(model="gpt-4o", temperature=0)
+    elif provider == "huggingface":
+        llm = ChatHuggingFace(
+            llm=HuggingFaceEndpoint(
+                url="https://api-inference.huggingface.co/models/Meta-DeepLearning/llama-2-7b-chat-hf",
+                temperature=0,
+            ),
+        )
+    else:
+        raise ValueError("Invalid provider. Choose 'google', 'groq', 'openai', or 'huggingface'.")
+
+    llm_with_tools = llm.bind_tools(tools)
+
+    def assistant(state: MessagesState):
+        return {"messages": [llm_with_tools.invoke(state["messages"])]}
+
+    def retriever_node(state: MessagesState):
+        results = vector_store.similarity_search(state["messages"][0].content)
+        example_msg = HumanMessage(
+            content=f"Here I provide a similar question and answer for reference:\n\n{results[0].page_content}"
+        )
+        return {"messages": [sys_msg] + state["messages"] + [example_msg]}
+
+    builder = StateGraph(MessagesState)
+    builder.add_node("retriever", retriever_node)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(tools))
+    builder.add_edge(START, "retriever")
+    builder.add_edge("retriever", "assistant")
+    builder.add_conditional_edges("assistant", tools_condition)
+    builder.add_edge("tools", "assistant")
+
+    return builder.compile()