updated

agent.py

@@ -144,7 +144,7 @@ def get_graph(llm):
                 f'<Document source="{doc.metadata["source"]}" page="{doc.metadata["title"]} 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 wiki_table_search(url: str, title: str) -> str:
@@ -157,7 +157,7 @@ def get_graph(llm):
         search_docs = WikipediaTableLoader(url=url, title=title).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>'
+                f'<Document source="{doc.metadata["source"]}" title="{doc.metadata["title"]}" table_index={doc.metadata["table_index"]}/>\n{doc.page_content}\n</Document>'
                 for doc in search_docs
             ])
         return formatted_search_docs