Fix

agent.py

@@ -2,13 +2,10 @@ import os
 import json
 from dotenv import load_dotenv
 
-# ---- Environment & Setup ----
 load_dotenv()
 os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
-
 hf_token = os.getenv("HUGGINGFACE_INFERENCE_TOKEN")
 
-# ---- Imports ----
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import tools_condition, ToolNode
 from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
@@ -20,7 +17,7 @@ from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain.schema import Document
 
-# ---- Tools ----
+# ---- Tool Definitions (with docstrings) ----
 
 @tool
 def multiply(a: int, b: int) -> int:
@@ -51,87 +48,58 @@ def modulus(a: int, b: int) -> int:
 
 @tool
 def wiki_search(query: str) -> str:
-    """Search Wikipedia for the given query and return formatted documents."""
+    """Search Wikipedia for the query and return text of up to 2 documents."""
     search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
     formatted = "\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
-        ]
+        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}
 
 @tool
 def web_search(query: str) -> str:
-    """Search the web using Tavily API for the given query."""
+    """Search the web for the query using Tavily and return up to 3 results."""
     search_docs = TavilySearchResults(max_results=3).invoke(query=query)
     formatted = "\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
-        ]
+        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}
 
 @tool
 def arvix_search(query: str) -> str:
-    """Search Arxiv for academic papers related to the query."""
+    """Search Arxiv for the query and return content from up to 3 papers."""
     search_docs = ArxivLoader(query=query, load_max_docs=3).load()
     formatted = "\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
-        ]
+        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}
 
-@tool
-def similar_question_search(query: str) -> str:
-    """Searches for questions similar to the input query using a vector database."""
-    matched_docs = vector_store.similarity_search(query, 3)
-    formatted = "\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 matched_docs
-        ]
-    )
-    return {"similar_questions": formatted}
-
-
-# ---- Embedding & Vector Store ----
-
+# Build vector store once
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-
-json_QA = []
-with open('metadata.jsonl', 'r') as jsonl_file:
-    for line in jsonl_file:
-        json_QA.append(json.loads(line))
-
+json_QA = [json.loads(line) for line in open("metadata.jsonl", "r")]
 documents = [
     Document(
         page_content=f"Question : {sample['Question']}\n\nFinal answer : {sample['Final answer']}",
         metadata={"source": sample["task_id"]}
-    )
-    for sample in json_QA
+    ) for sample in json_QA
 ]
-
 vector_store = Chroma.from_documents(
     documents=documents,
     embedding=embeddings,
     persist_directory="./chroma_db",
     collection_name="my_collection"
 )
-vector_store.persist()
 print("Documents inserted:", vector_store._collection.count())
 
 @tool
 def similar_question_search(query: str) -> str:
+    """Search for questions similar to the input query using the vector store."""
     matched_docs = vector_store.similarity_search(query, 3)
     formatted = "\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 matched_docs
-        ]
+        f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
+        for doc in matched_docs
     )
     return {"similar_questions": formatted}
 
@@ -143,17 +111,14 @@ Now, I will ask you a question. Report your thoughts, and finish your answer wit
 FINAL ANSWER: [YOUR FINAL ANSWER]. 
 YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings...
 """
-
 sys_msg = SystemMessage(content=system_prompt)
 
-# ---- Tool List ----
-
 tools = [
     multiply, add, subtract, divide, modulus,
     wiki_search, web_search, arvix_search, similar_question_search
 ]
 
-# ---- Graph Construction ----
+# ---- Graph Builder ----
 
 def build_graph(provider: str = "huggingface"):
     if provider == "huggingface":