Upload agent.py

agent.py

@@ -1,5 +1,6 @@
 """LangGraph Agent"""
 import os
+import json
 from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import tools_condition
@@ -18,6 +19,20 @@ from supabase.client import Client, create_client
 
 load_dotenv()
 
+def safe_get_metadata(doc, key, default=""):
+    """Safely extract metadata from document, handling string and dict formats"""
+    try:
+        if isinstance(doc.metadata, str):
+            # Try to parse as JSON if it's a string
+            metadata = json.loads(doc.metadata)
+        elif isinstance(doc.metadata, dict):
+            metadata = doc.metadata
+        else:
+            return default
+        return metadata.get(key, default)
+    except (json.JSONDecodeError, AttributeError):
+        return default
+
 @tool
 def multiply(a: int, b: int) -> int:
     """Multiply two numbers.
@@ -75,13 +90,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="{safe_get_metadata(doc, "source")}" page="{safe_get_metadata(doc, "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"Error searching Wikipedia: {str(e)}"}
 
 @tool
 def web_search(query: str) -> str:
@@ -89,13 +107,25 @@ 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_tool = TavilySearchResults(max_results=3)
+        search_results = search_tool.invoke(query)
+        
+        # Handle the case where search_results might be a list of dicts or Document objects
+        if isinstance(search_results, list):
+            formatted_search_docs = "\n\n---\n\n".join(
+                [
+                    f'<Document source="{result.get("url", "")}" />\n{result.get("content", "")}\n</Document>'
+                    if isinstance(result, dict) else
+                    f'<Document source="{safe_get_metadata(result, "source")}" page="{safe_get_metadata(result, "page")}"/>\n{result.page_content}\n</Document>'
+                    for result in search_results
+                ])
+        else:
+            formatted_search_docs = str(search_results)
+            
+        return {"web_results": formatted_search_docs}
+    except Exception as e:
+        return {"web_results": f"Error searching web: {str(e)}"}
 
 @tool
 def arvix_search(query: str) -> str:
@@ -103,19 +133,23 @@ 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="{safe_get_metadata(doc, "source")}" page="{safe_get_metadata(doc, "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"Error searching Arxiv: {str(e)}"}
 
 # load the system prompt from the file
-with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    system_prompt = f.read()
+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."
 
 # System message
 sys_msg = SystemMessage(content=system_prompt)
@@ -125,20 +159,22 @@ embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-b
 supabase_url = "https://ajnakgegqblhwltzkzbz.supabase.co"
 supabase_key = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6ImFqbmFrZ2VncWJsaHdsdHpremJ6Iiwicm9sZSI6ImFub24iLCJpYXQiOjE3NDkyMDgxODgsImV4cCI6MjA2NDc4NDE4OH0.b9RPF-5otedg4yiaQu_uhOgYpXVXd9D_0oR-9cluUjo"
 
-supabase: Client = create_client(supabase_url, supabase_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(),
-    name="Question Search",
-    description="A tool to retrieve similar questions from a vector store.",
-)
-
-
+try:
+    supabase: Client = create_client(supabase_url, supabase_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(),
+        name="Question Search",
+        description="A tool to retrieve similar questions from a vector store.",
+    )
+except Exception as e:
+    print(f"Warning: Could not initialize vector store: {e}")
+    vector_store = None
 
 tools = [
     multiply,
@@ -169,39 +205,31 @@ def build_graph(provider: str = "groq"):
         """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]}
-
     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")
+        """Retriever node with error handling"""
+        try:
+            if vector_store is None:
+                return {"messages": [AIMessage(content="Vector store not available.")]}
+                
+            query = state["messages"][-1].content
+            similar_docs = vector_store.similarity_search(query, k=1)
+            
+            if not similar_docs:
+                return {"messages": [AIMessage(content="No similar documents found.")]}
+            
+            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:
+            return {"messages": [AIMessage(content=f"Error in retriever: {str(e)}")]}
 
     builder = StateGraph(MessagesState)
     builder.add_node("retriever", retriever)
@@ -211,4 +239,4 @@ def build_graph(provider: str = "groq"):
     builder.set_finish_point("retriever")
 
     # Compile graph
-    return builder.compile()
+    return builder.compile()