ert

agent.py

@@ -96,41 +96,38 @@ def wiki_search(query: str) -> str:
 #             for doc in search_docs
 #         ])
 #     return {"web_results": formatted_search_docs}
-
-
 @tool
 def web_search(query: str) -> str:
     """Search Tavily for a query and return maximum 3 results.
+    
     Args:
         query: The search query."""
     try:
+        # Initialize the tool
         tavily_tool = TavilySearchResults(max_results=3)
-        search_results = tavily_tool.invoke(input=query) # Correct parameter is 'input'
-
-        formatted_docs = []
-        if isinstance(search_results, list) and search_results:
-            if all(isinstance(doc, dict) for doc in search_results): # Check for list of dicts
-                for doc_dict in search_results:
-                    title = doc_dict.get("title", "N/A")
-                    source_url = doc_dict.get("url", "N/A")
-                    content = doc_dict.get("content", "")
-                    # Adjust formatting to be most useful for the LLM
-                    formatted_docs.append(
-                        f'<Document source="{source_url}" title="{title}">\n{content}\n</Document>'
-                    )
-                formatted_search_docs = "\n\n---\n\n".join(formatted_docs)
-            elif all(hasattr(doc, 'metadata') and hasattr(doc, 'page_content') for doc in search_results): # Original check for Document objects
-                formatted_search_docs = "\n\n---\n\n".join(
-                    [
-                        f'<Document source="{doc.metadata.get("source", "N/A")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
-                        for doc in search_results
-                    ])
-            else: # Fallback for other list types
-                print(f"Tavily search returned a list of unexpected item types: {type(search_results[0])}")
-                formatted_search_docs = str(search_results)
-        elif isinstance(search_results, str): # Single string result
+        # Invoke it correctly
+        search_results = tavily_tool.invoke(input=query) # <--- CORRECTED LINE
+
+        # The result of TavilySearchResults.invoke is usually a list of strings or a single string.
+        # Let's check its type and format accordingly.
+        # Typically, TavilySearchResults directly returns a list of Document objects
+        # or a list of dictionaries if you've configured it differently.
+        # For the default, it's often a list of strings or a single concatenated string.
+        # If it returns a list of Document objects (which is common for loaders/retrievers):
+        if isinstance(search_results, list) and all(hasattr(doc, 'metadata') and hasattr(doc, 'page_content') for doc in search_results):
+            formatted_search_docs = "\n\n---\n\n".join(
+                [
+                    f'<Document source="{doc.metadata.get("source", "N/A")}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
+                    for doc in search_results
+                ])
+        # If it returns a list of strings (less likely for Tavily in recent versions, but good to check)
+        elif isinstance(search_results, list) and all(isinstance(item, str) for item in search_results):
+            formatted_search_docs = "\n\n---\n\n".join(search_results)
+        # If it returns a single string
+        elif isinstance(search_results, str):
             formatted_search_docs = search_results
         else:
+            # Fallback or handle unexpected format
             print(f"Unexpected Tavily search result format: {type(search_results)}")
             formatted_search_docs = str(search_results)
 
@@ -251,14 +248,14 @@ def build_graph(provider: str = "huggingface"):
         example_msg = HumanMessage(
             content=f"Here I provide a similar question and answer for reference: \n\n{similar_question[0].page_content}",
         )
-        print(example_msg)
+        print("ex msgs"+[sys_msg] + state["messages"] + [example_msg])
         return {"messages": [sys_msg] + state["messages"] + [example_msg]}
 
     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(START, "assistant")
     builder.add_edge("retriever", "assistant")
     builder.add_conditional_edges(
         "assistant",
@@ -291,7 +288,7 @@ def build_graph(provider: str = "huggingface"):
     # --- END: Visualization code ---
 
     return compiled_graph # This should be the last line of the function
-
+    
 # test
 if __name__ == "__main__":
     question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"