Update agent.py

agent.py

@@ -1,5 +1,6 @@
-"""LangGraph Agent"""
+"""LangGraph Agent (patched for robustness)"""
 import os
+import traceback
 from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import tools_condition
@@ -13,14 +14,11 @@ from langchain_community.document_loaders import ArxivLoader
 from langchain_community.vectorstores import SupabaseVectorStore
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
-
 from supabase.client import Client, create_client
-# --- langchain create_retriever_tool fallback (paste near other imports) ---
-# NOTE: removed the unconditional import that caused ModuleNotFoundError.
-import traceback
 
+# --- Safe import + fallback for langchain.tools.retriever.create_retriever_tool ---
 try:
-    # Prefer the real helper if available
+    # Try to import the real helper (if the installed langchain provides it)
     from langchain.tools.retriever import create_retriever_tool  # type: ignore
     HAS_CREATE_RETRIEVER_TOOL = True
 except Exception:
@@ -29,6 +27,10 @@ except Exception:
     print(traceback.format_exc())
 
     class _SimpleRetrieverTool:
+        """
+        Minimal tool-like wrapper providing a `.run(query)` method.
+        Most templates call tool.run(query) — adapt if your code uses a different interface.
+        """
         def __init__(self, retriever, name="retriever", description=""):
             self.name = name
             self.description = description
@@ -69,6 +71,7 @@ except Exception:
         """
         return _SimpleRetrieverTool(retriever, name=name, description=description)
 
+
 load_dotenv()
 
 @tool
@@ -128,13 +131,16 @@ def wiki_search(query: str) -> str:
     
     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}
+    try:
+        search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
+        formatted_search_docs = "\n\n---\n\n".join(
+            [
+                f'<Document source="{doc.metadata.get("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_error": str(e)}
 
 @tool
 def web_search(query: str) -> str:
@@ -142,13 +148,16 @@ def web_search(query: str) -> str:
     
     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}
+    try:
+        search_docs = TavilySearchResults(max_results=3).invoke(query=query)
+        formatted_search_docs = "\n\n---\n\n".join(
+            [
+                f'<Document source="{doc.metadata.get("source","")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+                for doc in search_docs
+            ])
+        return {"web_results": formatted_search_docs}
+    except Exception as e:
+        return {"web_results_error": str(e)}
 
 @tool
 def arvix_search(query: str) -> str:
@@ -156,14 +165,16 @@ def arvix_search(query: str) -> str:
     
     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}
-
+    try:
+        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
+        formatted_search_docs = "\n\n---\n\n".join(
+            [
+                f'<Document source="{doc.metadata.get("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_error": str(e)}
 
 
 # load the system prompt from the file
@@ -173,24 +184,53 @@ with open("system_prompt.txt", "r", encoding="utf-8") as f:
 # System message
 sys_msg = SystemMessage(content=system_prompt)
 
-# 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="documents",
-    query_name="match_documents_langchain",
-)
-retriever_tool = create_retriever_tool(
-    retriever=vector_store.as_retriever(),
-    name="Question Search",
-    description="A tool to retrieve similar questions from a vector store.",
-)
-
+# --- Build a retriever (defensive: don't crash if heavy deps or credentials missing) ---
+retriever_tool = None
+vector_store = None
+embeddings = None
 
+# Try to create HuggingFaceEmbeddings and SupabaseVectorStore if dependencies and env are present.
+try:
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")  # dim=768
+except Exception as e:
+    print(f"⚠️  Could not initialize HuggingFaceEmbeddings: {e}")
+    embeddings = None
+
+SUPABASE_URL = os.environ.get("SUPABASE_URL")
+SUPABASE_SERVICE_KEY = os.environ.get("SUPABASE_SERVICE_KEY")
+
+if SUPABASE_URL and SUPABASE_SERVICE_KEY and embeddings is not None:
+    try:
+        supabase: Client = create_client(SUPABASE_URL, SUPABASE_SERVICE_KEY)
+        vector_store = SupabaseVectorStore(
+            client=supabase,
+            embedding=embeddings,
+            table_name="documents",
+            query_name="match_documents_langchain",
+        )
+    except Exception as e:
+        print(f"⚠️  Could not initialize SupabaseVectorStore: {e}")
+        vector_store = None
+else:
+    if not SUPABASE_URL or not SUPABASE_SERVICE_KEY:
+        print("⚠️  SUPABASE_URL or SUPABASE_SERVICE_KEY not set — skipping vector store initialization.")
+    elif embeddings is None:
+        print("⚠️  Embeddings not available — skipping vector store initialization.")
+    vector_store = None
+
+# Create a retriever tool only if vector_store exists
+if vector_store is not None:
+    try:
+        retriever_tool = create_retriever_tool(
+            retriever=vector_store.as_retriever(),
+            name="Question Search",
+            description="A tool to retrieve similar questions from a vector store.",
+        )
+    except Exception as e:
+        print(f"⚠️  Failed to create retriever tool from vector store: {e}")
+        retriever_tool = None
+else:
+    retriever_tool = None
 
 
 tools = [
@@ -204,6 +244,20 @@ tools = [
     arvix_search,
 ]
 
+# Add retriever_tool to tools if available and matches the callable interface
+if retriever_tool is not None:
+    try:
+        if hasattr(retriever_tool, "run"):
+            @tool
+            def retriever_wrapper(query: str) -> str:
+                return retriever_tool.run(query)
+            tools.append(retriever_wrapper)
+        else:
+            tools.append(retriever_tool)
+    except Exception as e:
+        print(f"⚠️  Could not append retriever tool to tools list: {e}")
+
+
 # Build graph function
 def build_graph(provider: str = "google"):
     """Build the graph"""
@@ -213,7 +267,7 @@ def build_graph(provider: str = "google"):
         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
+        llm = ChatGroq(model="qwen-qwq-32b", temperature=0)  # optional : qwen-qwq-32b gemma2-9b-it
     elif provider == "huggingface":
         # TODO: Add huggingface endpoint
         llm = ChatHuggingFace(
@@ -224,52 +278,53 @@ def build_graph(provider: str = "google"):
         )
     else:
         raise ValueError("Invalid provider. Choose 'google', 'groq' or 'huggingface'.")
-    # Bind tools to LLM
-    llm_with_tools = llm.bind_tools(tools)
 
-    # Node
+    # Bind tools to LLM
+    try:
+        llm_with_tools = llm.bind_tools(tools)
+    except Exception as e:
+        print(f"⚠️  Could not bind tools to LLM: {e}")
+        # fallback: keep LLM without tools
+        llm_with_tools = llm
+
+    # Node: assistant
     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]}
+        try:
+            return {"messages": [llm_with_tools.invoke(state["messages"])]}
+        except Exception as e:
+            print(f"⚠️  assistant node failed: {e}")
+            # return empty message so graph can continue
+            return {"messages": [HumanMessage(content="")]}
 
     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")
-
+        # If vector_store not available, return empty message so assistant proceeds normally
+        if vector_store is None:
+            return {"messages": [AIMessage(content="")]}
+
+        try:
+            similar_docs = vector_store.similarity_search(query, k=1)
+            if not similar_docs:
+                return {"messages": [AIMessage(content="")]}
+            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)]}
+        except Exception as e:
+            print(f"⚠️  retriever node failed: {e}")
+            return {"messages": [AIMessage(content="")]}
+
+    # Build the state graph: a simple retriever-only entry point (defensive)
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
 
-    # Retriever ist Start und Endpunkt
+    # Retriever is both the entry and finish point in this design
     builder.set_entry_point("retriever")
     builder.set_finish_point("retriever")