Update agent.py

agent.py

@@ -3,9 +3,10 @@ from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import tools_condition
 from langgraph.prebuilt import ToolNode
-from langchain_community.tools.tavily_search import TavilySearchResults # 已经导入了
+from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
-from langchain_community.document_loaders import ArxivLoader
+# Removed: from langchain_community.document_loaders import ArxivLoader
+from langchain_community.tools import DuckDuckGoSearchRun # Added DuckDuckGo import
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 # from langchain_openai import ChatOpenAI
@@ -18,47 +19,16 @@ TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") # 需要在 Space Secrets 中添加
 
 if not DEEPSEEK_API_KEY:
     raise ValueError("DEEPSEEK_API_KEY not found in environment variables.")
+# Tavily is still included, so its key is needed if you want to use it.
+# If you ONLY want DuckDuckGo, you could remove Tavily and this check.
 if not TAVILY_API_KEY:
-    # Tavily is critical for most questions, raise error if not set
-    raise ValueError("TAVILY_API_KEY not found in environment variables. Please add it to your Space Secrets.")
+    print("Warning: TAVILY_API_KEY not found. Tavily search tool may not work.")
 
 
-@tool
-def multiply(a: int, b: int) -> int:
-    """Multiplies two numbers."""
-    return a * b
-
-
-@tool
-def add(a: int, b: int) -> int:
-    """Adds two numbers."""
-    return a + b
+# --- Removed math tools (multiply, add, subtract, divide, modulo) ---
 
 
-@tool
-def subtract(a: int, b: int) -> int:
-    """Subtracts two numbers."""
-    return a - b
-
-
-@tool
-def divide(a: int, b: int) -> int:
-    """Divides two numbers."""
-    # Added check for division by zero
-    if b == 0:
-        return "Error: Division by zero."
-    return a / b
-
-
-@tool
-def modulo(a: int, b: int) -> int:
-    """Returns the remainder of two numbers."""
-    # Added check for modulo by zero
-    if b == 0:
-        return "Error: Modulo by zero."
-    return a % b
-
-# Keep Wikipedia and Arxiv, but the general search will be more used
+# Keep Wikipedia search
 @tool
 def wiki_search(query: str) -> str:
     "Using Wikipedia, search for a query and return up to 2 relevant results."
@@ -76,32 +46,20 @@ def wiki_search(query: str) -> str:
         return f"An error occurred during Wikipedia search: {e}"
 
 
-@tool
-def arvix_search(query: str) -> str:
-    """Search Arxiv for scientific papers by query and return maximum 3 results (title and summary)."""
-    try:
-        search_docs = ArxivLoader(query=query, load_max_docs=3).load()
-        if not search_docs:
-             return "Arxiv search found no relevant papers."
-        # Format results to be more concise
-        formatted_search_docs = "\n\n---\n\n".join(
-            [
-                f'<Document source="Arxiv - {doc.metadata.get("source", "")}">\nTitle: {doc.metadata.get("Title", "N/A")}\nSummary: {doc.page_content[:1000]}...\n</Document>' # Limit summary length
-                for doc in search_docs
-            ])
-        return formatted_search_docs # Return string directly
-    except Exception as e:
-        return f"An error occurred during Arxiv search: {e}"
+# --- Removed Arxiv search (arvix_search) ---
+
 
-# *** ADD TAVILY WEB SEARCH TOOL ***
+# *** ADD TAVILY WEB SEARCH TOOL *** (Kept as requested implicitly by keeping web_search)
 @tool
 def web_search(query: str) -> str:
     """Search the web for a query using Tavily and return relevant snippets."""
+    if not TAVILY_API_KEY:
+         return "Tavily search is not available because TAVILY_API_KEY is not set."
     try:
         tavily = TavilySearchResults(max_results=5) # Get up to 5 results
         results = tavily.invoke(query)
         if not results:
-             return "Web search found no relevant results."
+             return "Web search (Tavily) found no relevant results."
         # Format Tavily results
         formatted_results = "\n\n---\n\n".join([
             f'<SearchResult source="{r["source"]}">\nTitle: {r["title"]}\nContent: {r["content"]}\n</SearchResult>'
@@ -109,24 +67,41 @@ def web_search(query: str) -> str:
         ])
         return formatted_results # Return string directly
     except Exception as e:
-        return f"An error occurred during web search: {e}"
+        return f"An error occurred during web search (Tavily): {e}"
+
+# *** ADD DUCKDUCKGO WEB SEARCH TOOL ***
+duckduckgo_search_tool_instance = DuckDuckGoSearchRun() # Instantiate the DuckDuckGo tool
+
+@tool
+def duckduckgo_search(query: str) -> str:
+    """Search the web for a query using DuckDuckGo."""
+    try:
+        # The DuckDuckGoSearchRun tool directly returns a formatted string result
+        results = duckduckgo_search_tool_instance.run(query)
+        if not results:
+             return "DuckDuckGo search found no relevant results."
+        # DuckDuckGoSearchRun often returns results as a string ready to be used
+        return results
+    except Exception as e:
+        return f"An error occurred during DuckDuckGo search: {e}"
 
 
 # load the system prompt from the file
 # Ensure this file exists and has the content from Step 2
-with open("system_prompt.txt", "r", encoding="utf-8") as f:
-    system_prompt = f.read()
-sys_msg = SystemMessage(content=system_prompt)
+try:
+    with open("system_prompt.txt", "r", encoding="utf-8") as f:
+        system_prompt = f.read()
+    sys_msg = SystemMessage(content=system_prompt)
+except FileNotFoundError:
+     print("Warning: system_prompt.txt not found. Using a default system message.")
+     sys_msg = SystemMessage(content="You are a helpful AI assistant. You can use tools to find information.")
+
 
+# Updated tools list: Removed math and Arxiv, Added DuckDuckGo
 tools = [
-    multiply,
-    add,
-    subtract,
-    divide,
-    modulo,
     wiki_search,
-    arvix_search,
-    web_search, # *** ADDED TAVILY WEB SEARCH ***
+    web_search, # Tavily search
+    duckduckgo_search, # DuckDuckGo search
 ]
 
 
@@ -140,17 +115,21 @@ def build_graph():
         api_key=DEEPSEEK_API_KEY,
         base_url="https://api.deepseek.com"
     )
+    # Bind the updated tools list to the LLM
     llm_with_tools = llm.bind_tools(tools)
 
     def assistant(state: MessagesState):
         """Assistant node: invoke LLM with tools."""
         print("---Calling Assistant---") # Added print for debugging
-        result = llm_with_tools.invoke(state["messages"])
+        # Include the system message at the beginning of the conversation
+        messages_for_llm = [sys_msg] + state["messages"]
+        result = llm_with_tools.invoke(messages_for_llm)
         print(f"---Assistant Response: {result}") # Added print for debugging
         return {"messages": [result]}
 
     builder = StateGraph(MessagesState)
     builder.add_node("assistant", assistant)
+    # The ToolNode needs the list of functions, not just the names
     builder.add_node("tools", ToolNode(tools))
 
     builder.add_edge(START, "assistant")
@@ -178,20 +157,24 @@ def build_graph():
 
 if __name__ == "__main__":
     # Example Usage (for local testing)
-    # To run this part, make sure you have DEEPSEEK_API_KEY and TAVILY_API_KEY
-    # set in your environment or a .env file loaded beforehand.
+    # To run this part, make sure you have DEEPSEEK_API_KEY set.
+    # TAVILY_API_KEY is needed for the web_search tool. DuckDuckGo usually works without a key.
     # If running locally, you'd typically use `load_dotenv()` here or in app.py
 
+    print("Note: Ensure DEEPSEEK_API_KEY is set.")
+    print("Note: TAVILY_API_KEY is required for the 'web_search' tool (Tavily). DuckDuckGo usually works without a key.")
+
     # Test questions covering different tool needs
+    # Removed purely math questions and Arxiv questions.
+    # Added questions that might benefit from multiple search tools.
     questions_for_testing = [
-        "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)?", # Web Search
-        "In the video https://www.youtube.com/watch?v=L1vXCYZAYYM, what is the highest number of bird species seen?", # Requires video analysis (will likely fail with current tools)
-        ".rewsna eht sa \"tfel\" drow eht fo etisoppo eht etirw ,ecnetnes siht dnatsrednu uoy fI", # Text manipulation (no tool needed)
-        "What is 12345 * 6789?", # Calculator
-        "Who nominated the only Featured Article on English Wikipedia about a dinosaur that was promoted in November 2023?", # Web Search/Wikipedia
-        "What country had the least number of athletes at the 1928 Summer Olympics?", # Web Search
-        "Review the chess position provided in the image. It is black's turn. Provide the correct next move from this position: [Describe the position or mention image input which is not supported]", # Requires image analysis (will likely fail)
-        # Add more questions from your evaluation set to test
+        "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)?", # Web Search (Tavily or DDG)
+        "Who nominated the only Featured Article on English Wikipedia about a dinosaur that was promoted in November 2023? Use Wikipedia first if possible.", # Wikipedia or Web Search
+        "What country had the least number of athletes at the 1928 Summer Olympics? Find this information using web search.", # Web Search (Tavily or DDG)
+        "Tell me about the Voyager 1 probe. Use Wikipedia.", # Wikipedia
+        "What is the current population of Tokyo?", # Web Search (Tavily or DDG)
+        "Give me a brief overview of the concept of 'LangGraph'.", # Web Search (Tavily or DDG)
+        ".rewsna eht sa \"tfel\" drow ehT etirw ,ecnetnes siht dnatsrednu uoy fI", # Text manipulation (no tool needed)
     ]
 
 
@@ -204,7 +187,7 @@ if __name__ == "__main__":
     #         f.write(png_data)
     #     print("Graph visualization saved to graph.png")
     # except Exception as e:
-    #      print(f"Could not draw graph: {e}")
+    #      print(f"Could not draw graph: {e}. Make sure 'pygraphviz' and graphviz system libraries are installed.")
 
 
     print("\n--- Running single question tests ---")
@@ -212,11 +195,16 @@ if __name__ == "__main__":
         print(f"\n--- Testing Question {i+1}: {question}")
         try:
             # LangGraph returns the final state after execution completes or hits recursion limit
+            # Need to start with the system message and the first human message
+            # The assistant node prepends the system message internally now.
             final_state = graph.invoke({"messages": [HumanMessage(content=question)]})
             print("\n--- Final State Messages ---")
+            # Print messages more readably
             for m in final_state["messages"]:
-                 m.pretty_print()
+                print(f"{m.__class__.__name__}: {m.content}")
             print("-" * 30)
         except Exception as e:
              print(f"--- Error running graph for this question: {e}")
+             import traceback
+             traceback.print_exc() # Print full traceback for debugging
              print("-" * 30)