Upload agent.py

agent.py

@@ -10,10 +10,27 @@ from langchain_core.messages import HumanMessage, SystemMessage
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import ToolNode, tools_condition
 from IPython.display import Image, display
+from langchain_core.messages import AIMessage
+from langchain_community.vectorstores import SupabaseVectorStore
+from supabase.client import Client, create_client
+import os
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
 
 load_dotenv('../config.env')
 
 llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash")
+embedding_model = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004")
+
+supabase_url = os.environ.get("SUPABASE_URL")
+supabase_key = os.environ.get("SUPABASE_SERVICE_KEY")
+supabase: Client = create_client(supabase_url, supabase_key)
+
+vector_store = SupabaseVectorStore(
+    client=supabase,
+    embedding= embedding_model,
+    table_name="documents",
+    query_name="match_documents_langchain",
+)
 
 # load the system prompt from the file
 with open('system_prompt.txt', 'r') as f:
@@ -159,24 +176,47 @@ def build_graph():
 
         return {"messages": [llm_with_tools.invoke([sys_msg] + state["messages"])]}
 
+    def retriever(state: AgentState):
+        query = state["messages"][-1].content
+        results = vector_store.similarity_search(query, k=1)
+
+        if not results:
+            # If no documents are found, provide a fallback response.
+            answer = "I couldn't find anything relevant in the knowledge base. Please try rephrasing your question."
+        else:
+            similar_doc = results[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)]}
+
 
 
     # Graph
     builder = StateGraph(AgentState)
 
     # Define nodes: these do the work
-    builder.add_node("assistant", assistant)
-    builder.add_node("tools", ToolNode(tools))
+    # builder.add_node("assistant", assistant)
+    # builder.add_node("tools", ToolNode(tools))
+
+    # # Define edges: these determine how the control flow moves
+    # builder.add_edge(START, "assistant")
+    # builder.add_conditional_edges(
+    #     "assistant",
+    #     # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
+    #     # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
+    #     tools_condition,
+    # )
+    # builder.add_edge("tools", "assistant")
+    builder.add_node("retriever", retriever)
 
     # Define edges: these determine how the control flow moves
-    builder.add_edge(START, "assistant")
-    builder.add_conditional_edges(
-        "assistant",
-        # If the latest message (result) from assistant is a tool call -> tools_condition routes to tools
-        # If the latest message (result) from assistant is a not a tool call -> tools_condition routes to END
-        tools_condition,
-    )
-    builder.add_edge("tools", "assistant")
+    builder.add_edge(START, "retriever")
+    builder.set_finish_point("retriever")
+
     react_graph = builder.compile()
 
     # Show