Update agent.py

agent.py

@@ -1,210 +1,107 @@
 import os
 from dotenv import load_dotenv
 from langgraph.graph import START, StateGraph, MessagesState
-from langgraph.prebuilt import tools_condition
-from langgraph.prebuilt import ToolNode
+from langgraph.prebuilt import tools_condition, ToolNode
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_community.document_loaders import WikipediaLoader
-# Removed: from langchain_community.document_loaders import ArxivLoader
-from langchain_community.tools import DuckDuckGoSearchRun # Added DuckDuckGo import
+from langchain_community.tools import DuckDuckGoSearchRun
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
-# from langchain_openai import ChatOpenAI
 from langchain_deepseek import ChatDeepSeek
 
-# load_dotenv() # 假设你在 app.py 或其他地方加载了 .env
-# Ensure API keys are set
+load_dotenv()
+
 DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
-TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") # 需要在 Space Secrets 中添加 TAVILY_API_KEY
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 
 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:
     print("Warning: TAVILY_API_KEY not found. Tavily search tool may not work.")
 
-
-# --- Removed math tools (multiply, add, subtract, divide, modulo) ---
-
-
-# Keep Wikipedia search
 @tool
 def wiki_search(query: str) -> str:
-    "Using Wikipedia, search for a query and return up to 2 relevant results."
     try:
-        search_docs = WikipediaLoader(query=query, load_max_docs=2, doc_content_chars_max=2000).load() # Limit content length
+        search_docs = WikipediaLoader(query=query, load_max_docs=2, doc_content_chars_max=2000).load()
         if not search_docs:
-             return "Wikipedia search found no relevant pages."
+            return "Wikipedia search found no relevant pages."
         formatted_search_docs = "\n\n---\n\n".join(
             [
                 f'<Document source="Wikipedia - {doc.metadata.get("source", "")}" page="{doc.metadata.get("page", "")}">\n{doc.page_content}\n</Document>'
                 for doc in search_docs
             ])
-        return formatted_search_docs # Return string directly
+        return formatted_search_docs
     except Exception as e:
         return f"An error occurred during Wikipedia search: {e}"
 
-
-# --- Removed Arxiv search (arvix_search) ---
-
-
-# *** 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."
+        return "Tavily search is not available because TAVILY_API_KEY is not set."
     try:
-        tavily = TavilySearchResults(max_results=5) # Get up to 5 results
+        tavily = TavilySearchResults(max_results=5)
         results = tavily.invoke(query)
         if not results:
-             return "Web search (Tavily) found no relevant results."
-        # Format Tavily results
+            return "Web search (Tavily) found no relevant results."
         formatted_results = "\n\n---\n\n".join([
             f'<SearchResult source="{r["source"]}">\nTitle: {r["title"]}\nContent: {r["content"]}\n</SearchResult>'
             for r in results
         ])
-        return formatted_results # Return string directly
+        return formatted_results
     except Exception as 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
+duckduckgo_search_tool_instance = DuckDuckGoSearchRun()
 
 @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 "DuckDuckGo search found no relevant results."
         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
 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.")
-
+    print("Warning: system_prompt.txt not found. Using a default system message.")
+    sys_msg = SystemMessage(content="You are a helpful AI assistant.")
 
-# Updated tools list: Removed math and Arxiv, Added DuckDuckGo
 tools = [
     wiki_search,
-    web_search, # Tavily search
-    duckduckgo_search, # DuckDuckGo search
+    web_search,
+    duckduckgo_search,
 ]
 
-
 def build_graph():
     llm = ChatDeepSeek(
         model="deepseek-chat",
-        temperature=0, # Keep low for factual answers
+        temperature=0,
         max_tokens=None,
         timeout=None,
         max_retries=2,
         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
-        # Include the system message at the beginning of the conversation
+        print("---Calling Assistant---")
         messages_for_llm = [sys_msg] + state["messages"]
         result = llm_with_tools.invoke(messages_for_llm)
-        print(f"---Assistant Response: {result}") # Added print for debugging
+        print(f"---Assistant Response: {result}")
         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")
-
-    # The tools_condition checks if the last message from "assistant" is a tool call.
-    # If yes, it transitions to "tools".
-    # If no, the graph implicitly ends. This is how the agent stops.
-    builder.add_conditional_edges(
-        "assistant",
-        tools_condition,
-        # If tool_condition is false (no tool calls detected), the default is None,
-        # which implicitly ends the graph execution for that path.
-        # We don't need to explicitly define other paths here for a simple graph.
-    )
-
-    # After a tool is executed, the result is added to the state, and the control
-    # goes back to the assistant to process the tool result and decide the next step.
+    builder.add_conditional_edges("assistant", tools_condition)
     builder.add_edge("tools", "assistant")
 
-    # You can optionally increase the recursion limit if your graph is expected to be complex,
-    # but it's better to fix the LLM's logic via the prompt first.
-    # return builder.compile(recursion_limit=50) # Example of increasing limit
-    return builder.compile()
-
-
-if __name__ == "__main__":
-    # Example Usage (for local testing)
-    # 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 (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)
-    ]
-
-
-    graph = build_graph()
-
-    # Optional: Draw graph
-    # try:
-    #     png_data = graph.get_graph().draw_mermaid_png()
-    #     with open("graph.png", "wb") as f:
-    #         f.write(png_data)
-    #     print("Graph visualization saved to graph.png")
-    # except Exception as e:
-    #      print(f"Could not draw graph: {e}. Make sure 'pygraphviz' and graphviz system libraries are installed.")
-
-
-    print("\n--- Running single question tests ---")
-    for i, question in enumerate(questions_for_testing):
-        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"]:
-                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)
+    return builder.compile()