Update agent.py

agent.py

@@ -1,497 +1,76 @@
-import os
-from dotenv import load_dotenv
-from langgraph.graph import START, END, StateGraph, MessagesState
-from langgraph.prebuilt import tools_condition, ToolNode
-from langchain_google_genai import ChatGoogleGenerativeAI
-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, AIMessage
-from langchain_core.tools import tool
-from langchain.tools.retriever import create_retriever_tool
+import os 
+from dotenv import load_dotenv 
 from supabase.client import Client, create_client
-
+from langchain_huggingface import HuggingFaceEmbeddings 
+from langchain_community.vectorstores import SupabaseVectorStore 
+from langgraph.graph import StateGraph, MessageState
+from langgraph.prebuilt import ToolNode
+from langchain_core.messages import HumanMessage, AIMessage
+from langchain_core.tools import tool 
 load_dotenv()
 
-@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) -> int:
-    """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
-
-@tool
-def wiki_search(query: str) -> str:
-    """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_search_docs = "\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_search_docs}
-
-# @tool
-# def web_search(query: str) -> str:
-#     """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_search_docs = "\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 {"web_results": formatted_search_docs}
-# from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
-
-@tool
-def web_search(query: str) -> str:
-    """Search the web for a query and return maximum 3 results.
-    
-    Args:
-        query: The search query."""
-    search = DuckDuckGoSearchAPIWrapper()
-    results = search.results(query, 3)
-    formatted_results = "\n\n---\n\n".join(
-        [
-            f"Title: {res['title']}\nURL: {res['link']}\nSnippet: {res['snippet']}"
-            for res in results
-        ]
-    )
-    return formatted_results
-@tool
-def arvix_search(query: str) -> str:
-    """Search Arxiv for a query and return maximum 3 result.
-    
-    Args:
-        query: The search query."""
-    search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-    formatted_search_docs = "\n\n---\n\n".join(
-        [
-            f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
-            for doc in search_docs
-        ])
-    return {"arvix_results": formatted_search_docs}
-
-
-
-# load the system prompt from the file
-with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    system_prompt = f.read()
+supabase: Client = create_client(
+    os.environ["SUPABASE_URL"],
+    os.environ["SUPABASE_SERVICE_KEY"]
+)
 
-# System message
-sys_msg = SystemMessage(content=system_prompt)
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
-# build a retriever
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2") #  dim=768
-supabase: Client = create_client(
-    os.environ.get("SUPABASE_URL"), 
-    os.environ.get("SUPABASE_SERVICE_KEY"))
-vector_store = SupabaseVectorStore(
+vector_search = SupabaseVectorStore(
     client=supabase,
     embedding= embeddings,
     table_name="docs",
-    query_name="match_documents_langchain",
+    query_name="match_documents_langchain"
 )
-create_retriever_tool = create_retriever_tool(
-    retriever=vector_store.as_retriever(),
-    name="question_search",
-    description="A tool to retrieve similar questions from a vector store.",
-)
-
-
-
-tools = [
-    multiply,
-    add,
-    subtract,
-    divide,
-    modulus,
-    wiki_search,
-    web_search,
-    arvix_search,
-    create_retriever_tool
-]
-
-def build_graph(provider: str = "openai"):
-    """Build the graph using OpenAI or Hugging Face"""
-    # Validate provider
-    if provider not in ["openai", "huggingface"]:
-        raise ValueError("Invalid provider. Choose 'openai' or 'huggingface'.")
-
-    # Initialize LLM
-    if provider == "openai":
-        from langchain_openai import ChatOpenAI
-        llm = ChatOpenAI(model="gpt-4o", temperature=0)
-    else:  # huggingface
-        from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
-        llm = ChatHuggingFace(
-            llm=HuggingFaceEndpoint(
-                endpoint_url="https://api-inference.huggingface.co/models/meta-llama/Llama-3.1-8B-Instruct",
-                temperature=0,
-            )
-        )
-
-    # Bind tools to LLM
-    llm_with_tools = llm.bind_tools(tools)
 
-    # Define nodes
-    def assistant(state: MessagesState):
-        """Assistant node - generates responses"""
-        messages = llm_with_tools.invoke(state["messages"])
-        # Generate response using LLM
-        # response = llm_with_tools.invoke(messages)
-        # Return new state with appended message
-        return {"messages": messages}
+all_rows = supabase.table("docs").select("content").execute().data
 
+qa_dict: dict[str, str] = {}
+for row in all_rows:
+    raw = row["content"]
+    if "Answer:" in raw:
+        parts = raw.split("Answer:", 1)
+        question_part = parts[0].strip()
+        answer_part = parts[1].strip()
+        if question_part.lower().startswith("question"):
+            question_part = question_part.split(":", 1)[1].strip()
+        qa_dict[question_part] = answer_part
+    else:
+        qa_dict[raw.strip()] = ""
 
-    def retriever(state: MessagesState):
-        """Retriever node - provides context from vector store"""
-        messages = state["messages"]
-        query = messages[-1].content
-        
-        # Retrieve similar documents
-        similar_docs = vector_store.similarity_search(query, k=1)
-        
-        if not similar_docs:
-            return {"messages": messages}
-        
-        context = similar_docs[0].page_content
-        context_msg = SystemMessage(content=f"Reference context:\n{context}")
-        
-        return {"messages": messages + [context_msg]}
 
-    # Build graph
-    builder = StateGraph(MessagesState)
-    
-    # Add nodes
-    # builder.add_node("retriever", retriever)
-    builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
-    
-    # Set up edges
-    builder.set_entry_point("assistant")
-    builder.set_finish_point("assistant")
-    
-    # Compile graph
-    return builder.compile()
-    
-    # def retriever(state: MessagesState):
-    #     """Retriever node - provides context from vector store"""
-    #     # Get current messages
-    #     messages = state["messages"]
-    #     # Last message is the user query
-    #     query = messages[-1].content
-        
-    #     # Retrieve similar documents
-    #     similar_docs = vector_store.similarity_search(query, k=1)
-        
-    #     if not similar_docs:
-    #         # Return original messages if no context found
-    #         return {"messages": messages}
-        
-    #     # Get context from first document
-    #     context = similar_docs[0].page_content
-    #     # Create system message with context
-    #     context_msg = SystemMessage(content=f"Reference context:\n{context}")
-        
-    #     # Append context to messages
-    #     return {"messages": messages + [context_msg]}
-
-    # # Build graph
-    # builder = StateGraph(MessagesState)
-    
-    # # Add nodes
-    # builder.add_node("retriever", retriever)
-    # builder.add_node("assistant", assistant)
-    # builder.add_node("tools", ToolNode(tools))
-    
-    # # Set up edges
-    # builder.set_entry_point("retriever")
-    # builder.add_edge("retriever", "assistant")
-    
-    # # Conditional tool usage
-    # builder.add_conditional_edges(
-    #     "assistant",
-    #     tools_condition,
-    #     {
-    #         # Continue to tools if needed
-    #         "continue": "tools",
-    #         # End conversation if no tools needed
-    #         "end": END
-    #     }
-    # )
-    
-    # # After tools execute, return to assistant for response generation
-    # builder.add_edge("tools", "assistant")
-    
-    # # builder.add_finish_point(END)  # Explicitly declare END as finish point
-    # return builder.compile()
-
-
-# def build_graph(provider: str = "openai"):
-#     """Build the graph using OpenAI or Hugging Face"""
-#     # Validate provider
-#     if provider not in ["openai", "huggingface"]:
-#         raise ValueError("Invalid provider. Choose 'openai' or 'huggingface'.")
-
-#     # Initialize LLM based on provider
-#     if provider == "openai":
-#         from langchain_openai import ChatOpenAI
-#         llm = ChatOpenAI(model="gpt-4o", temperature=0)
-#     else:  # huggingface
-#         from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
-#         llm = ChatHuggingFace(
-#             llm=HuggingFaceEndpoint(
-#                 endpoint_url="https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct",
-#                 temperature=0,
-#             )
-#         )
-
-#     # Bind tools to LLM
-#     llm_with_tools = llm.bind_tools(tools)
-
-#     # Define nodes
-#     def assistant(state: MessagesState):
-#         """Assistant node"""
-#         return {"messages": [llm_with_tools.invoke(state["messages"])]}
-    
-#     def retriever(state: MessagesState):
-#         """Retriever node - provides context from vector store"""
-#         query = state["messages"][-1].content
-#         similar_docs = vector_store.similarity_search(query, k=1)
-        
-#         if not similar_docs:
-#             return {"messages": [AIMessage(content="No relevant information found")]}
-            
-#         similar_doc = similar_docs[0]
-#         content = similar_doc.page_content
-        
-#         # Extract answer if formatted, otherwise use full content
-#         if "Final answer :" in content:
-#             answer = content.split("Final answer :")[-1].strip()
-#         else:
-#             answer = content.strip()
-
-#         return {"messages": [AIMessage(content=answer)]}
-
-#     # Build graph
-#     builder = StateGraph(MessagesState)
-    
-#     # Add nodes
-#     builder.add_node("retriever", retriever)
-#     builder.add_node("assistant", assistant)
-#     builder.add_node("tools", ToolNode(tools))
-    
-#     # Set up edges
-#     builder.set_entry_point("retriever")
-#     builder.add_edge("retriever", "assistant")
-#     builder.add_conditional_edges(
-#         "assistant",
-#         tools_condition,
-#         {"continue": "tools", "end": END}
-#     )
-#     builder.add_edge("tools", "assistant")
-    
-#     return builder.compile()
-# def build_graph(provider: str = "google"):
-#     """Build the graph"""
-#     # Load environment variables from .env file
-#     if provider == "google":
-#         # Google Gemini
-#         llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0)
-#     elif provider == "groq":
-#         # Groq https://console.groq.com/docs/models
-#         llm = ChatGroq(model="qwen-qwq-32b", temperature=0) # optional : qwen-qwq-32b gemma2-9b-it
-#     elif provider == "huggingface":
-#         # TODO: Add huggingface endpoint
-#         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' or 'huggingface'.")
-#     # Bind tools to LLM
-#     llm_with_tools = llm.bind_tools(tools)
-
-#     # Node
-#     def assistant(state: MessagesState):
-#         """Assistant node"""
-#         return {"messages": [llm_with_tools.invoke(state["messages"])]}
-    
-#     # def retriever(state: MessagesState):
-#       #  """Retriever node"""
-#        # similar_question = 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{similar_question[0].page_content}",
-#        # )
-#        # return {"messages": [sys_msg] + state["messages"] + [example_msg]}
-
-#     from langchain_core.messages import AIMessage
-
-#     def retriever(state: MessagesState):
-#         query = state["messages"][-1].content
-#         similar_doc = vector_store.similarity_search(query, k=1)[0]
-
-#         content = similar_doc.page_content
-#         if "Final answer :" in content:
-#             answer = content.split("Final answer :")[-1].strip()
-#         else:
-#             answer = content.strip()
-
-#         return {"messages": [AIMessage(content=answer)]}
-
-   # 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")
-
-    # builder = StateGraph(MessagesState)
-    # builder.add_node("retriever", retriever)
-
-    # # Retriever ist Start und Endpunkt
-    # builder.set_entry_point("retriever")
-    # builder.set_finish_point("retriever")
-
-    # # Compile graph
-    # return builder.compile()
-# def build_graph(provider: str = "openai"):
-#     """Build the graph using OpenAI or Hugging Face"""
-    
-#     if provider == "openai":
-#         # OpenAI ChatGPT (e.g., GPT-4 or GPT-3.5)
-#         from langchain.chat_models import ChatOpenAI
-#         llm = ChatOpenAI(model="gpt-4", temperature=0)
-
-#     elif provider == "huggingface":
-#         # Hugging Face endpoint
-#         from langchain.chat_models import ChatHuggingFace
-#         from langchain.llms import HuggingFaceEndpoint
-
-#         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 'openai' or 'huggingface'.")
-
-#     # Bind tools to LLM
-#     llm_with_tools = llm.bind_tools(tools)
-
-#     # return llm_with_tools
-
-#     # Node
-#     def assistant(state: MessagesState):
-#         """Assistant node"""
-#         return {"messages": [llm_with_tools.invoke(state["messages"])]}
-    
-#     # def retriever(state: MessagesState):
-#       #  """Retriever node"""
-#        # similar_question = 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{similar_question[0].page_content}",
-#        # )
-#        # return {"messages": [sys_msg] + state["messages"] + [example_msg]}
-
-#     from langchain_core.messages import AIMessage
-
-#     def retriever(state: MessagesState):
-#         query = state["messages"][-1].content
-#         similar_doc = vector_store.similarity_search(query, k=1)[0]
-#         if not similar_docs:
-#            return {"messages": [AIMessage(content="No relevant information found")]}
-#         similar_doc = similar_docs[0]
-
-#         content = similar_doc.page_content
-#         if "Final answer :" in content:
-#             answer = content.split("Final answer :")[-1].strip()
-#         else:
-#             answer = content.strip()
-
-#         return {"messages": [AIMessage(content=answer)]}
-
-#    # 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")
-
-#     builder = StateGraph(MessagesState)
-#     builder.add_node("retriever", retriever)
+@tool
+def find_answer(query: str) -> str:
+    """
+    If 'query' exactly matches a key in qa_dict, return qa_dict[query].
+    Otherwise, do an embedding search (k=1) in Supabase and return only the "Answer:" portion. 
+    """
 
-    
-#     builder.set_entry_point("retriever")
-#     builder.set_finish_point("retriever")
+    if query in qa_dict:
+        return qa_dict[query]
+    similar_docs = vector_store.similarity_search(query, k=1)
+    if not similar_docs:
+        return "Sorry, I couldn't find that question"
+    top_doc = similar_docs[0].page_content
+    if "Answer:" in top_doc:
+        return top_doc.split("Answer:", 1)[1].strip()
+    if "Final answer: " in top_doc:
+        return top_doc.split("Final answer :", 1)[1].strip()
+    return top_doc.strip()
 
-#     # Compile graph
-#     return builder.compile()
+tools = [find_answer]
 
+def build_graph(provider: str = "openai"):
+    """
+    Build a LangGraph where every HumanMessage is handled by find_answe(---),
+    and the returned AIMessage contains exactly the stored answer text.    
+    """
+    def retriever_node(state: MessageState):
+        user_query = state["messages"][-1].content
+        answer_text = find_answer(user_query)
+        return {"messages": state["messages"] + [AIMessage(content=answer_text)]}
+    builder = StateGraph(MessageState)
+    builder.add_node("retriever", retriever_node)
+    builder.set_entry_point("retriever")
+    builder.set_finish_point("retriever")
+    return builder.compile()