| |
| |
| |
|
|
| import os |
| from dotenv import load_dotenv |
| from langgraph.graph import START, StateGraph, MessagesState |
| from langgraph.prebuilt import tools_condition, ToolNode |
| from langchain_google_genai import ChatGoogleGenerativeAI |
| from langchain_groq import ChatGroq |
| from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings |
| from langchain_community.document_loaders import WikipediaLoader, ArxivLoader |
| from langchain_community.vectorstores import SupabaseVectorStore |
| from langchain_core.messages import SystemMessage, HumanMessage |
| from langchain_core.tools import tool |
| from langchain_tavily import TavilySearch |
| from langchain.tools.retriever import create_retriever_tool |
| from supabase.client import Client, create_client |
|
|
| load_dotenv() |
|
|
| |
| url = os.getenv("SUPABASE_URL") |
| key = os.getenv("SUPABASE_KEY") |
| supabase: Client = create_client(url, key) |
|
|
| |
| @tool |
| def multiply(a: int, b: int) -> int: |
| """Multiply two numbers and return the result.""" |
| return a * b |
|
|
| @tool |
| def add(a: int, b: int) -> int: |
| """Add two numbers and return the result.""" |
| return a + b |
|
|
| @tool |
| def subtract(a: int, b: int) -> int: |
| """Subtract second number from first and return the result.""" |
| return a - b |
|
|
| @tool |
| def divide(a: int, b: int) -> float: |
| """Divide first number by second and return the result.""" |
| if b == 0: |
| raise ValueError("Cannot divide by zero.") |
| return a / b |
|
|
| @tool |
| def modulus(a: int, b: int) -> int: |
| """Return the modulus (remainder) of two numbers.""" |
| return a % b |
|
|
| @tool |
| def wiki_search(query: str) -> str: |
| """Search Wikipedia and return 2 results.""" |
| docs = WikipediaLoader(query=query, load_max_docs=2).load() |
| return "\n\n---\n\n".join(doc.page_content for doc in docs) |
|
|
| @tool |
| def web_search(query: str) -> str: |
| """Search the web using Tavily and return 3 results.""" |
| docs = TavilySearch(max_results=3).invoke(query) |
| return "\n\n---\n\n".join(doc.page_content for doc in docs) |
|
|
| @tool |
| def arvix_search(query: str) -> str: |
| """Search Arxiv for academic papers and return 3 results.""" |
| docs = ArxivLoader(query=query, load_max_docs=3).load() |
| return "\n\n---\n\n".join(doc.page_content[:1000] for doc in docs) |
|
|
| |
| with open("system_prompt.txt", "r") as f: |
| system_prompt = f.read() |
|
|
| sys_msg = SystemMessage(content=system_prompt) |
|
|
| |
| embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") |
| 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 vector DB.", |
| ) |
|
|
| |
| tools = [ |
| multiply, add, subtract, divide, modulus, |
| wiki_search, web_search, arvix_search, |
| retriever_tool, |
| ] |
|
|
| |
|
|
| def build_graph(provider: str = "groq"): |
| if provider == "google": |
| llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0) |
| elif provider == "groq": |
| llm = ChatGroq(model="qwen-qwq-32b", temperature=0, api_key=os.getenv("GROQ_API")) |
| 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") |
|
|
| llm_with_tools = llm.bind_tools(tools) |
|
|
| def assistant(state: MessagesState): |
| return {"messages": [llm_with_tools.invoke(state["messages"])]} |
|
|
| def retriever(state: MessagesState): |
| docs = vector_store.similarity_search(state["messages"][0].content) |
| if not docs: |
| return {"messages": [sys_msg] + state["messages"]} |
| similar_msg = HumanMessage(content=f"Reference: {docs[0].page_content}") |
| return {"messages": [sys_msg] + state["messages"] + [similar_msg]} |
|
|
| builder = StateGraph(MessagesState) |
| builder.add_node("retriever", retriever) |
| 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() |
|
|
|
|
| |
| |
| |
|
|
