Update agent.py

agent.py

@@ -1,5 +1,4 @@
-"""LangGraph Agent - Fixed Version"""
-
+"""LangGraph Agent"""
 import os
 from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
@@ -7,11 +6,7 @@ from langgraph.prebuilt import tools_condition
 from langgraph.prebuilt import ToolNode
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_groq import ChatGroq
-from langchain_huggingface import (
-    ChatHuggingFace,
-    HuggingFaceEndpoint,
-    HuggingFaceEmbeddings,
-)
+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
@@ -23,7 +18,6 @@ from supabase.client import Client, create_client
 
 load_dotenv()
 
-
 @tool
 def multiply(a: int, b: int) -> int:
     """Multiply two numbers.
@@ -33,33 +27,30 @@ def multiply(a: int, b: int) -> 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
@@ -68,116 +59,85 @@ def divide(a: int, b: int) -> int:
         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."""
-    try:
-        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}
-    except Exception as e:
-        return {"wiki_results": f"Wikipedia search failed: {str(e)}"}
-
+    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."""
-    try:
-        search_docs = TavilySearchResults(max_results=3).invoke(query=query)
-        formatted_search_docs = "\n\n---\n\n".join(
-            [
-                f'<Document source="{doc.get("url", "unknown")}" page="{doc.get("title", "")}"/>\n{doc.get("content", "")}\n</Document>'
-                for doc in search_docs
-            ]
-        )
-        return {"web_results": formatted_search_docs}
-    except Exception as e:
-        return {"web_results": f"Web search failed: {str(e)}"}
-
+    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}
 
 @tool
 def arvix_search(query: str) -> str:
     """Search Arxiv for a query and return maximum 3 result.
-
+    
     Args:
         query: The search query."""
-    try:
-        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}
-    except Exception as e:
-        return {"arvix_results": f"Arxiv search failed: {str(e)}"}
+    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
-try:
-    with open("system_prompt.txt", "r", encoding="utf-8") as f:
-        system_prompt = f.read()
-except FileNotFoundError:
-    system_prompt = "You are a helpful AI assistant. Answer questions accurately and concisely."
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
 
 # System message
 sys_msg = SystemMessage(content=system_prompt)
 
-# build a retriever with error handling
-def initialize_vector_store():
-    """Initialize vector store with proper error handling."""
-    try:
-        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",
-        )
-        return vector_store
-    except Exception as e:
-        print(f"Warning: Failed to initialize vector store: {e}")
-        return None
+# 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(
+    client=supabase,
+    embedding= embeddings,
+    table_name="documents2",
+    query_name="match_documents_2",
+)
+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.",
+)
 
-# Initialize vector store
-vector_store = initialize_vector_store()
 
-# Create retriever tool if vector store is available
-if vector_store:
-    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,
@@ -190,26 +150,20 @@ tools = [
     arvix_search,
 ]
 
-
 # Build graph function
 def build_graph(provider: str = "huggingface"):
     """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
+
+    if provider == "groq":
+        llm = ChatGroq(model="qwen-qwq-32b", temperature=0) # optional : qwen-qwq-32b gemma2-9b-it
     elif provider == "huggingface":
         llm = ChatHuggingFace(
-            llm=HuggingFaceEndpoint(repo_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
+            llm=HuggingFaceEndpoint(
+                repo_id = "Qwen/Qwen2.5-Coder-32B-Instruct"
+            ),
         )
     else:
         raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
-    
     # Bind tools to LLM
     llm_with_tools = llm.bind_tools(tools)
 
@@ -217,42 +171,14 @@ def build_graph(provider: str = "huggingface"):
     def assistant(state: MessagesState):
         """Assistant node"""
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
-
+    
     def retriever(state: MessagesState):
-        """Retriever node with proper error handling"""
-        try:
-            # Check if vector store is available
-            if vector_store is None:
-                print("Vector store not available, proceeding without retrieval")
-                return {"messages": [sys_msg] + state["messages"]}
-            
-            # Get the user question
-            user_question = state["messages"][-1].content if state["messages"] else ""
-            
-            if not user_question:
-                print("No user question found, proceeding without retrieval")
-                return {"messages": [sys_msg] + state["messages"]}
-            
-            # Perform similarity search
-            similar_questions = vector_store.similarity_search(user_question, k=1)
-            
-            # Check if we found any similar questions
-            if similar_questions and len(similar_questions) > 0:
-                # Extract the first similar question
-                similar_content = similar_questions[0].page_content
-                example_msg = HumanMessage(
-                    content=f"Here I provide a similar question and answer for reference: \n\n{similar_content}",
-                )
-                print("Found similar question for retrieval")
-                return {"messages": [sys_msg] + state["messages"] + [example_msg]}
-            else:
-                print("No similar questions found, proceeding without retrieval example")
-                return {"messages": [sys_msg] + state["messages"]}
-                
-        except Exception as e:
-            print(f"Error in retriever node: {e}")
-            # Fallback: proceed without retrieval
-            return {"messages": [sys_msg] + state["messages"]}
+        """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]}
 
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
@@ -269,7 +195,6 @@ def build_graph(provider: str = "huggingface"):
     # Compile graph
     return builder.compile()
 
-
 # test
 if __name__ == "__main__":
     question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"