Update agent.py

agent.py

@@ -1,4 +1,3 @@
-"""LangGraph Agent"""
 import os
 from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
@@ -121,33 +120,21 @@ with open("system_prompt.txt", "r", encoding="utf-8") as f:
 sys_msg = SystemMessage(content=system_prompt)
 
 # build a retriever
-try:
-    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="documents",
-        query_name="match_documents_langchain",
-    )
-    
-    # Test the connection
-    test_results = vector_store.similarity_search("test query", k=1)
-    print(f"Vector store initialized successfully. Test returned {len(test_results)} results.")
-    
-except Exception as e:
-    print(f"Warning: Vector store initialization failed: {e}")
-    vector_store = None
-
+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="documents",
+    query_name="match_documents_langchain",
+)
 create_retriever_tool = create_retriever_tool(
-    retriever=vector_store.as_retriever() if vector_store else None,
+    retriever=vector_store.as_retriever(),
     name="Question Search",
     description="A tool to retrieve similar questions from a vector store.",
-) if vector_store else None
+)
 
 
 
@@ -163,20 +150,25 @@ tools = [
 ]
 
 # Build graph function
-def build_graph(provider: str = "huggingface"):
-    """Build the graph with improved error handling"""
-
-    if provider == "groq":
-        llm = ChatGroq(
-            model="qwen-qwq-32b", temperature=0
-        )  # optional : qwen-qwq-32b gemma2-9b-it
+def build_graph(provider: str = "groq"):
+    """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(repo_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
+            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)
 
@@ -184,76 +176,28 @@ def build_graph(provider: str = "huggingface"):
     def assistant(state: MessagesState):
         """Assistant node"""
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
+    
+    from langchain_core.messages import AIMessage
 
     def retriever(state: MessagesState):
-        """Retriever node with 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"]}
-                
-            similar_question = vector_store.similarity_search(state["messages"][0].content)
-            
-            # Check if we have results before accessing them
-            if similar_question and len(similar_question) > 0:
-                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]}
-            else:
-                # No similar questions found, proceed without reference
-                print("No similar questions found in vector store")
-                return {"messages": [sys_msg] + state["messages"]}
-                
-        except Exception as e:
-            print(f"Error in retriever: {e}")
-            # Fallback: continue without retrieval
-            return {"messages": [sys_msg] + state["messages"]}
+        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")
 
-    # Compile graph
-    return builder.compile()
-    
-def retriever(state: MessagesState):
-    """Retriever node with error handling"""
-    try:
-        similar_question = vector_store.similarity_search(state["messages"][0].content)
-        
-        # Check if we have results before accessing them
-        if similar_question and len(similar_question) > 0:
-            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]}
-        else:
-            # No similar questions found, proceed without reference
-            print("No similar questions found in vector store")
-            return {"messages": [sys_msg] + state["messages"]}
-            
-    except Exception as e:
-        print(f"Error in retriever: {e}")
-        # Fallback: continue without retrieval
-        return {"messages": [sys_msg] + state["messages"]}
+    # Retriever ist Start und Endpunkt
+    builder.set_entry_point("retriever")
+    builder.set_finish_point("retriever")
 
-# 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?"
-    # Build the graph
-    graph = build_graph(provider="groq")
-    # Run the graph
-    messages = [HumanMessage(content=question)]
-    messages = graph.invoke({"messages": messages})
-    for m in messages["messages"]:
-        m.pretty_print()
+    # Compile graph
+    return builder.compile()