Tools, Retriever, Systemp prompt and Agent creation

agent.py

@@ -0,0 +1,120 @@
+import os
+from dotenv import load_dotenv
+from typing import TypedDict, List, Dict, Any, Optional
+
+from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
+from langchain_huggingface.chat_models import ChatHuggingFace
+from langchain_groq.chat_models import ChatGroq
+
+from langgraph.graph.message import add_messages
+from langgraph.graph import StateGraph, START, END, MessagesState
+from langgraph.prebuilt import ToolNode, tools_condition
+
+from tools import (
+    add, 
+    subtract, multiply, div, modulus, power,
+    wikipedia_search, search_web, arxiv_search, 
+    save_and_read_file, download_file_from_url, extract_text_from_image,
+    pdf_loader
+)
+from retriever import get_retriever_tool
+
+load_dotenv(dotenv_path = ".env")
+
+# Configurations
+SYSTEM_PROMPT_PATH = "system_prompt.txt"
+DEFAULT_PROVIDER = "groq"
+MODEL_NAME = "llama3-70b-8192"
+
+def load_system_prompt(path: str = SYSTEM_PROMPT_PATH) -> str:
+    if not os.path.exists(path):
+        raise ValueError(f"System prompt file not foud at: {path}")
+    with open(path, "r", encoding = "utf-8") as f:
+        return f.read()
+
+
+system_prompt = load_system_prompt()
+sys_msg = SystemMessage(content = system_prompt)
+
+# Load tools
+vector_store, vector_retriever, retriever_tool = get_retriever_tool()
+
+TOOLS = [
+    # Math
+    add, subtract, multiply, div, modulus, power,
+    # Documents Search
+    wikipedia_search, search_web, arxiv_search, 
+    # Process Files
+    save_and_read_file, download_file_from_url, extract_text_from_image,
+    pdf_loader,
+    # Retriever
+    retriever_tool
+]
+
+def get_llm(provider: str = DEFAULT_PROVIDER):
+    if provider == "groq":
+        return ChatGroq(model = MODEL_NAME, temperature = 0)
+    elif provider == "huggingface":
+        raise NotImplementedError("HuggingFace support not yet implemented.")
+    else:
+        raise ValueError("Invalid LLM provider. Choose 'groq' or 'huggingface'")
+    
+
+def build_graph(provider: str = DEFAULT_PROVIDER):
+    """
+    Builds LangGraph graph
+    """
+    llm = get_llm(provider)
+    
+    # Add tools to the LLM
+    llm_with_tools = llm.bind_tools(TOOLS)
+
+    def assistant(state: MessagesState):
+        return {"messages": llm_with_tools.invoke(state["messages"])}
+    
+    def retriever(state: MessagesState):
+        query = state["messages"][0].content
+        similar_qas = vector_store.similarity_search(query)
+
+        if similar_qas:
+            reference = similar_qas[0].page_content
+            example_qa = HumanMessage(
+                content = f"I provide a similar question and answer for reference:\n\n{reference}"
+            )
+            return {"messages": [sys_msg] + state["messages"] + [example_qa]}
+        else:
+            return {"messages": [sys_msg] + state["messages"]}
+    
+    # Graph
+    builder = StateGraph(MessagesState)
+    
+    # Nodes
+    builder.add_node("retriever", retriever)
+    builder.add_node("assistant", assistant)
+    builder.add_node("tools", ToolNode(TOOLS))
+    
+    # Edges
+    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()
+
+if __name__ == "__main__":
+    import random
+    import json
+
+    with open("metadata.jsonl") as dataset_file:
+        json_list = list(dataset_file)
+
+    QAs = [json.loads(qa) for qa in json_list]
+    question = random.choice(QAs)["Question"]
+    graph = build_graph()
+    messages = [HumanMessage(content = question)]
+    messages = graph.invoke({"messages": messages})
+    for m in messages["messages"]:
+        m.pretty_print()

retriever.py

@@ -0,0 +1,44 @@
+import os
+from dotenv import load_dotenv
+
+from huggingface_hub import login
+from supabase import Client, create_client
+from supabase.client import ClientOptions
+
+from langchain_community.vectorstores import SupabaseVectorStore
+from langchain_huggingface.embeddings import HuggingFaceEmbeddings
+from langchain.tools.retriever import create_retriever_tool
+
+load_dotenv()
+
+MODEL_NAME = "BAAI/bge-base-en-v1.5"
+TBL_NAME = "documents_tbl"
+QUERY_NAME = "match_documents"
+
+def get_retriever_tool():
+    embedding_model = HuggingFaceEmbeddings(model_name = MODEL_NAME)
+    DIMS_EMBEDDING = embedding_model._client.get_sentence_embedding_dimension()
+
+    # Supabase client
+    supabase: Client = create_client(
+        os.environ.get("SUPABASE_URL"), 
+        os.environ.get("SUPABASE_ANON_KEY"), 
+        options = ClientOptions(schema = "public")
+    )
+
+    # Vector Store
+    vector_store = SupabaseVectorStore(
+        client = supabase,
+        embedding = embedding_model,
+        table_name = TBL_NAME,
+        query_name = QUERY_NAME
+    )
+
+    vector_retriever = vector_store.as_retriever()
+
+    retriever_tool = create_retriever_tool(
+        retriever = vector_retriever,
+        name = "question_search_tool",
+        description = "A tool to retrieve similar questions based on embedding from Supabase vector store."
+    )
+    return vector_store, vector_retriever, retriever_tool

system_prompt.txt

@@ -0,0 +1,12 @@
+You are a general AI assistant tasked with answering questions using a set of tools.
+I will ask you a question. Report your thoughts, and finish your answer with the following template:
+FINAL ANSWER: [YOUR FINAL ANSWER].
+YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
+If you are asked for a number, don’t use comma to write your number neither use units such as $ or
+percent sign unless specified otherwise.
+If you are asked for a string, don’t use articles, neither abbreviations (e.g. for cities), and write the
+digits in plain text unless specified otherwise.
+If you are asked for a comma separated list, apply the above rules depending of whether the element
+to be put in the list is a number or a string, ensure there is exactly one space after each comma.
+Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
+

tools.py

@@ -0,0 +1,300 @@
+# - `Search Engine` (arXiv, Wikipedia, DuckDuckGo)
+# - `Calculator` (add, substract, divide, multiply, modulus, etc.)
+# - `Access` and `Download Files` from Web
+# - `Excel`/`Google Sheets`: Process Downloaded files
+
+import os
+import requests
+import tempfile
+import uuid
+
+import pytesseract
+
+from datetime import datetime, timezone
+from PIL import Image
+from urllib.parse import urlparse
+from typing import Optional
+
+from langchain_core.tools import tool
+from langchain_community.document_loaders import (
+    WikipediaLoader, ArxivLoader, PyPDFLoader
+)
+from langchain_community.tools import DuckDuckGoSearchResults
+
+#* === MATH TOOLS ===
+@tool
+def add(a: float, b: float) -> int:
+    """
+    Adds multple integers.
+    
+    Args:
+        a (float): Number to add.
+        b (float): Number to add.
+
+    Returns:
+        int: Sum of the two provided integers.
+    """
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """
+    Subtracts one integer from another.
+
+    Args:
+        a (int): The number from which to subtract.
+        b (int): The number to subtract.
+
+    Returns:
+        int: The result of a - b.
+    """
+    return a - b
+
+@tool
+def multiply(a: float, b: float) -> int:
+    """
+    Multiplies multple integers.
+    
+    Args:
+        a (float): First number to multiply.
+        b (float): Second number to multiply.
+
+    Returns:
+        int: Multiplication of the two provided floats or integers.
+    """
+    return a * b
+
+@tool
+def div(a: float, b: float):
+    """
+    Divides two numbers.
+
+    Args:
+        a (int or float): The dividend.
+        b (int or float): The divisor.
+
+    Returns:
+        float: The result of dividing a by b.
+
+    Raises:
+        ZeroDivisionError: If b is zero.
+    """
+    return a / b
+
+@tool
+def modulus(a: int, b: int):
+    """
+    Computes the modulus (remainder) of dividing two integers.
+
+    Args:
+        a (int): The dividend.
+        b (int): The divisor.
+
+    Returns:
+        int: The remainder when a is divided by b.
+
+    Raises:
+        ZeroDivisionError: If b is zero.
+    """
+    return a % b
+
+@tool
+def power(a: float, b: float) -> float:
+    """
+    Raises a number `a` to the power of `b`.
+
+    Args:
+        a (float): Base number.
+        b (float): Exponent.
+
+    Returns:
+        float: Result of a ** b.
+    """
+    return a**b
+
+#* === SEARCH TOOLS ===
+
+@tool
+def wikipedia_search(query: str) -> dict:
+    """
+    Search Wikipedia for a query and return up to 3 formatted results.
+
+    Args:
+        query (str): The topic to search for.
+
+    Returns:
+        dict: A dictionary with the key 'wikipedia_results' containing the formatted documents.
+    """
+    search_docs = WikipediaLoader(query = query,load_max_docs = 3).load()
+    formatted_docs = "\n\n---\n\n".join(
+        [
+            f"Document source='{doc.metadata['source']}' page={doc.metadata.get('page', '')}/>\n"
+            f"{doc.page_content}\n</Document>"
+            for doc in search_docs
+        ]
+    )
+    return {"wikipedia_results": formatted_docs}
+
+@tool
+def search_web(query: str) -> dict:
+    """
+    Performs a web search using DuckDuckGo and returns up to 4 formatted results.
+
+    Args:
+        query (str): The search query to submit to DuckDuckGo.
+
+    Returns:
+        dict: A dictionary with a single key "web_results" containing the formatted search results
+            as a string. Each result includes the document source and content, separated by "---".
+    """
+    search_docs = DuckDuckGoSearchResults(max_results = 4).invoke(query)
+    formatted_docs = "\n\n---\n\n".join(
+        [
+            f"Document source='{doc.metadata['source']}' page={doc.metadata.get('page', '')}/>\n"
+            f"{doc.page_content}\n</Document>"
+            for doc in search_docs
+        ]
+    )
+    return {"web_results": formatted_docs}
+
+@tool
+def arxiv_search(query: str) -> dict:
+    """
+    Perform a search on the arXiv academic paper repository and return the top results.
+
+    Args:
+        query (str): The search query to use on arXiv.
+
+    Returns:
+        dict: A dictionary containing a string under the key "arxiv_results", which includes
+            a formatted summary of the top retrieved documents. Each entry contains the 
+            document's source, optional page number, and the first 1000 characters of the content.
+    """
+    search_docs = ArxivLoader(query = query, load_max_docs = 3).load()
+    formatted_docs = "\n\n---\n\n".join(
+        [
+            f"Document source='{doc.metadata['source']}' page={doc.metadata.get('page', '')}/>\n"
+            f"{doc.page_content[:1000]}\n</Document>"
+            for doc in search_docs
+        ]
+    )
+    return {"arxiv_results": formatted_docs}
+
+#* === FILE PROCESSING TOOLS ===
+
+@tool
+def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
+    """
+    Saves the provided text content to a temporary file and returns its path.
+
+    If no filename is provided, a random temporary filename will be generated.
+    The file is saved in the system's temporary directory.
+
+    Args:
+        content (str): The text content to be written to the file.
+        filename (Optional[str]): Optional name for the file. If not provided, a temporary name is used.
+
+    Returns:
+        str: A message with the path to the saved file, indicating it is ready for processing.
+    """
+    try:
+        temp_dir = tempfile.gettempdir()
+        if filename is None:
+            temp_file = tempfile.NamedTemporaryFile(delete = False, dir = temp_dir)
+            filepath = temp_file.name
+        else:
+            filepath = os.path.join(temp_dir, filename)
+        
+        with open(filepath, "w", encoding = "utf-8") as f:
+            f.write(content)
+        return f"File saved to {filepath}. It is available to read for processing its contents."
+    except Exception as e:
+        return f"Error saving file: {str(e)}"
+
+@tool
+def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
+    """
+    Downloads a file from a given URL and saves it to a temporary directory.
+
+    If no filename is provided, it attempts to extract it from the URL. If the URL
+    does not contain a valid filename, a temporary unique filename will be generated.
+
+    Args:
+        url (str): The URL of the file to download.
+        filename (Optional[str]): Optional name for the downloaded file.
+
+    Returns:
+        str: A string indicating the path to the downloaded file, or an error message.
+    """
+    try:
+        # Parse URL to get filename if not provided
+        if not filename:
+            path = urlparse(url).path
+            filename = os.path.basename(path)
+            if not filename:
+                ts = datetime.now(timezone.utc).strftime('%Y%m%d_%H%M%S')
+                filename = f"downloaded_{ts}_{uuid.uuid4().hex[:8]}.tmp"
+        
+        # Create temporary file
+        temp_dir = tempfile.gettempdir()
+        filepath = os.path.join(temp_dir, filename)
+
+        # Download the file
+        response = requests.get(url, stream = True)
+        response.raise_for_status()
+
+        # Save the file
+        with open(filepath, "wb") as f:
+            for chunk in response.iter_content(chunk_size = 8192):
+                if chunk:
+                    f.write(chunk)
+        return f"File downloaded to {filepath}. It is available to read for processing its contents."
+    except Exception as e:
+        return f"Error downloading file: {str(e)}"
+
+@ tool
+def extract_text_from_image(image_path: str) -> str:
+    """
+    Extracts text content from an image file using Optical Character Recognition (OCR).
+
+    Args:
+        image_path (str): The path to the image file from which text will be extracted.
+
+    Returns:
+        str: Extracted text content. If extraction fails, returns an error message.
+    """
+    try:
+        # Open image
+        image = Image.open(image_path)
+        # Extract text from image
+        text = pytesseract.image_to_string(image)
+        return f"Text extracted from image:\n\n{text.strip()}"
+    except Exception as e:
+        return f"Error extracting text from image '{image_path}': {str(e)}"
+
+@tool
+def pdf_loader(filepath: str) -> dict:
+    """
+    Loads a PDF file from the given file path, parses its contents,
+    and returns a preview of each page's content (up to 1000 characters per page).
+
+    Args:
+        filepath (str): The full path to the PDF file.
+
+    Returns:
+        dict: A dictionary containing formatted PDF page previews under the key 'pdf_results'.
+            Each page is separated by "\n\n---\n\n".
+    """
+    try:
+        pdf_content = PyPDFLoader(file_path=filepath).load()
+        formatted_content = "\n\n---\n\n".join(
+            [
+                f"Document source='{doc.metadata['source']}' page={doc.metadata.get('page', '')}/>\n"
+                f"{doc.page_content[:1000]}\n</Document>"
+                for doc in pdf_content
+            ]
+        )
+        return {"pdf_results": formatted_content}
+    except Exception as e:
+        return {"pdf_results": f"Error reading PDF file: {str(e)}"}
+