Update agent.py

agent.py

@@ -1,7 +1,17 @@
 # agent.py
 import os
+from supabase import create_client
+from sentence_transformers import SentenceTransformers
 from langgraph import Graph, LLM, tool #or other graph library
 
+# ----Supabase setup----
+SUPABASE_URL = os.getenv("SUPABASE_URL")
+SUPABASE_KEY = os.getenv("SUPABASE_SERVICE_KEY")
+EMBED_MODEL_ID = os.getenv("HF_EMBEDDING_MODEL")
+
+sb_client = create_client(SUPABASE_URL, SUPABASE_KEY)
+embedder = SentenceTransformers(EMBED_MODEL_ID)
+
 # 1) Define tools
 
 @tool
@@ -15,7 +25,28 @@ def calculator(expr: str) -> str:
 # @tool
 # def web_search(query:str) -> str:
 # ...
+@tool
+def retrieve_docs(query: str, k: int = 3) -> str:
+    """
+    Fetch tpo-k docs from Supabase vector store.
+    Returns the concatenated text.
+    """
+    # --- embed the query
+    q_emb = embedder.encode(query).tolist()
+
+    # --- query the embedding table
+    response = (
+        sb_client
+        .rpc("match_documents", {"query_embedding": q_emb, "match_count": k})
+        .execute()
+    )
+    rows = response.data
 
+    # ---- concatenate the content field
+    docs = [row["content"] for row in rows]
+    return "\n\n---\n\n".join(docs)
+    
+    
 # 2) Build your graph
 def build_graph(provider: str = "huggingface") -> Graph:
     # 2a) Instantiate your LLM endpoint