rebased

agent.py

@@ -7,7 +7,7 @@ from langchain.chat_models import init_chat_model
 
 from langchain_community.document_loaders import WikipediaLoader, ArxivLoader, YoutubeLoader
 from langchain_community.tools import TavilySearchResults
-from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, AnyMessage
+from langchain_core.messages import HumanMessage
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langgraph.graph import add_messages, START, END, StateGraph
 from langchain_core.tools import tool
@@ -56,7 +56,7 @@ def get_graph(llm):
     from langchain_community.retrievers import TavilySearchAPIRetriever
 
     # Wikipedia retriever
-    wiki_retriever = WikipediaRetriever(load_max_docs =20)
+    wiki_retriever = WikipediaRetriever()
     # Tavily retriever
     tavily_retriever = TavilySearchAPIRetriever(k=3)
 
@@ -93,7 +93,7 @@ def get_graph(llm):
         return a - b
 
     @tool
-    def divide(a: int, b: int) -> int:
+    def divide(a: int, b: int) -> float:
         """Divide two numbers.
 
         Args:
@@ -125,7 +125,7 @@ def get_graph(llm):
         print("The query is: ", query)
         docs = wiki_retriever.invoke(query)
         serialized = "\n\n".join(
-            (f"\nContent:\n{doc.page_content}")
+            f"\nContent:\n{doc.page_content}"
             for doc in docs
         )
 
@@ -160,7 +160,7 @@ def get_graph(llm):
                 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}
+        return formatted_search_docs
 
     @tool
     def online_search(query: str):
@@ -172,7 +172,7 @@ def get_graph(llm):
         # docs = tavily_retriever.invoke(query)
         docs = TavilySearchResults(max_results=3).invoke({'query': query})
         serialized = "\n\n".join(
-            (f"\nContent:\n{doc.page_content}")
+            f"\nContent:\n{doc.page_content}"
             for doc in docs
         )
 
@@ -191,7 +191,7 @@ def get_graph(llm):
                 f'URL: {doc["url"]}\nTitle= {doc["title"]}\nContent: {doc["content"]}'
                 for doc in search_docs
             ])
-        return {"web_results": formatted_search_docs}
+        return formatted_search_docs
 
     @tool
     def arvix_search(query: str) -> str:
@@ -206,7 +206,7 @@ def get_graph(llm):
                 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}
+        return formatted_search_docs
 
     @tool
     def youtube_transcript(url: str) -> str:
@@ -220,10 +220,10 @@ def get_graph(llm):
         )
         docs = loader.load()
         transcript = "\n\n".join(
-            (f"\nContent:\n{doc.page_content}")
+            f"\nContent:\n{doc.page_content}"
             for doc in docs
         )
-        return {"youtube_transcript": transcript}
+        return transcript