Upload 2 files

agent.py

@@ -0,0 +1,233 @@
+import os
+import time
+import json
+from dotenv import load_dotenv
+from langgraph.graph import StateGraph, END
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.tools import DuckDuckGoSearchRun
+from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
+from langchain_core.messages import SystemMessage, AIMessage, HumanMessage
+from langchain_core.tools import tool
+from tenacity import retry, stop_after_attempt, wait_exponential
+from typing import TypedDict, Annotated, Sequence
+import operator
+
+# Load environment variables
+load_dotenv()
+google_api_key = os.getenv("GOOGLE_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+if not google_api_key:
+    raise ValueError("Missing GOOGLE_API_KEY environment variable")
+
+# --- Math Tools ---
+@tool
+def multiply(a: int, b: int) -> int:
+    """Multiply two integers."""
+    return a * b
+
+@tool
+def add(a: int, b: int) -> int:
+    """Add two integers."""
+    return a + b
+
+@tool
+def subtract(a: int, b: int) -> int:
+    """Subtract b from a."""
+    return a - b
+
+@tool
+def divide(a: int, b: int) -> float:
+    """Divide a by b, error on zero."""
+    if b == 0:
+        raise ValueError("Cannot divide by zero.")
+    return a / b
+
+@tool
+def modulus(a: int, b: int) -> int:
+    """Compute a mod b."""
+    return a % b
+
+# --- Browser Tools ---
+@tool
+def wiki_search(query: str) -> str:
+    """Search Wikipedia and return up to 3 relevant documents."""
+    try:
+        docs = WikipediaLoader(query=query, load_max_docs=3).load()
+        if not docs:
+            return "No Wikipedia results found."
+        
+        results = []
+        for doc in docs:
+            title = doc.metadata.get('title', 'Unknown Title')
+            content = doc.page_content[:2000]  # Limit content length
+            results.append(f"Title: {title}\nContent: {content}")
+        
+        return "\n\n---\n\n".join(results)
+    except Exception as e:
+        return f"Wikipedia search error: {str(e)}"
+
+@tool
+def arxiv_search(query: str) -> str:
+    """Search Arxiv and return up to 3 relevant papers."""
+    try:
+        docs = ArxivLoader(query=query, load_max_docs=3).load()
+        if not docs:
+            return "No arXiv papers found."
+        
+        results = []
+        for doc in docs:
+            title = doc.metadata.get('Title', 'Unknown Title')
+            authors = ", ".join(doc.metadata.get('Authors', []))
+            content = doc.page_content[:2000]  # Limit content length
+            results.append(f"Title: {title}\nAuthors: {authors}\nContent: {content}")
+        
+        return "\n\n---\n\n".join(results)
+    except Exception as e:
+        return f"arXiv search error: {str(e)}"
+
+@tool
+def web_search(query: str) -> str:
+    """Search the web using DuckDuckGo and return top results."""
+    try:
+        search = DuckDuckGoSearchRun()
+        result = search.run(query)
+        return f"Web search results for '{query}':\n{result[:2000]}"  # Limit content length
+    except Exception as e:
+        return f"Web search error: {str(e)}"
+
+# --- Load system prompt ---
+with open("system_prompt.txt", "r", encoding="utf-8") as f:
+    system_prompt = f.read()
+
+# --- Tool Setup ---
+tools = [
+    multiply,
+    add,
+    subtract,
+    divide,
+    modulus,
+    wiki_search,
+    arxiv_search,
+    web_search,
+]
+
+# --- Graph Builder ---
+def build_graph():
+    # Initialize model with Gemini 2.5 Flash
+    llm = ChatGoogleGenerativeAI(
+        model="gemini-2.5-flash",
+        temperature=0.3,
+        google_api_key=google_api_key,
+        max_retries=3
+    )
+    
+    # Bind tools to LLM
+    llm_with_tools = llm.bind_tools(tools)
+    
+    # 1. 定义状态结构
+    class AgentState(TypedDict):
+        messages: Annotated[Sequence, operator.add]
+    
+    # 2. 创建图
+    workflow = StateGraph(AgentState)
+    
+    # 3. 定义节点函数
+    def agent_node(state: AgentState):
+        """主代理节点"""
+        try:
+            # 添加请求间隔
+            time.sleep(1)
+            
+            # 带重试的调用
+            @retry(stop=stop_after_attempt(3),
+                   wait=wait_exponential(multiplier=1, min=4, max=10))
+            def invoke_with_retry():
+                return llm_with_tools.invoke(state["messages"])
+            
+            response = invoke_with_retry()
+            return {"messages": [response]}
+        
+        except Exception as e:
+            error_type = "UNKNOWN"
+            if "429" in str(e):
+                error_type = "QUOTA_EXCEEDED"
+            elif "400" in str(e):
+                error_type = "INVALID_REQUEST"
+                
+            error_msg = f"AGENT ERROR ({error_type}): {str(e)[:200]}"
+            return {"messages": [AIMessage(content=error_msg)]}
+    
+    def tool_node(state: AgentState):
+        """工具执行节点"""
+        last_msg = state["messages"][-1]
+        tool_calls = last_msg.additional_kwargs.get("tool_calls", [])
+        
+        responses = []
+        for call in tool_calls:
+            tool_name = call["function"]["name"]
+            tool_args = call["function"].get("arguments", {})
+            
+            # 查找工具
+            tool_func = next((t for t in tools if t.name == tool_name), None)
+            if not tool_func:
+                responses.append(f"Tool {tool_name} not available")
+                continue
+            
+            try:
+                # 解析参数
+                if isinstance(tool_args, str):
+                    tool_args = json.loads(tool_args)
+                
+                # 执行工具
+                result = tool_func.invoke(tool_args)
+                responses.append(f"{tool_name} result: {result[:1000]}")  # 限制结果长度
+            except Exception as e:
+                responses.append(f"{tool_name} error: {str(e)}")
+        
+        # 修复括号错误：确保正确关闭所有括号
+        tool_response_content = "\n".join(responses)
+        return {"messages": [AIMessage(content=tool_response_content)]}
+    
+    # 4. 添加节点到工作流
+    workflow.add_node("agent", agent_node)
+    workflow.add_node("tools", tool_node)
+    
+    # 5. 设置入口点
+    workflow.set_entry_point("agent")
+    
+    # 6. 定义条件边
+    def should_continue(state: AgentState):
+        last_msg = state["messages"][-1]
+        
+        # 错误情况直接结束
+        if "AGENT ERROR" in last_msg.content:
+            return "end"
+        
+        # 有工具调用则转到工具节点
+        if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
+            return "tools"
+        
+        # 包含最终答案则结束
+        if "FINAL ANSWER" in last_msg.content:
+            return "end"
+            
+        # 其他情况继续代理处理
+        return "agent"
+    
+    workflow.add_conditional_edges(
+        "agent",
+        should_continue,
+        {
+            "agent": "agent",
+            "tools": "tools",
+            "end": END
+        }
+    )
+    
+    # 7. 定义工具节点后的流向
+    workflow.add_edge("tools", "agent")
+    
+    # 8. 编译图
+    return workflow.compile()
+
+# 初始化代理图
+agent_graph = build_graph()

system_prompt.txt

@@ -0,0 +1,32 @@
+You are GAIA (General AI Assistant). Follow these rules:
+
+1. Always use available tools when needed
+2. Format responses as:
+   REASONING: <step-by-step logic>
+   FINAL ANSWER: <concise result>
+
+Available Tools:
+- Math tools: multiply, add, subtract, divide, modulus
+- Search tools: 
+  * wiki_search (Wikipedia) - for biographical and discography info
+  * arxiv_search (arXiv) - for scientific papers
+  * web_search (general web search) - for current info
+
+For music artist discography queries (like Mercedes Sosa):
+1. Use wiki_search("Artist Name discography") to find the discography section
+2. Identify studio albums specifically (exclude live albums, compilations, etc.)
+3. Check release years carefully - note that reissues or remasters don't count as new studio albums
+4. If Wikipedia is unclear, use web_search for additional sources
+5. Count only albums released within the specified year range
+
+Example for Mercedes Sosa question:
+USER: How many studio albums were published by Mercedes Sosa between 2000 and 2009?
+REASONING: 
+1. Use wiki_search("Mercedes Sosa discography")
+2. Parse discography section and filter for studio albums
+3. Identify albums released between 2000-2009:
+   - Acústico (2002, studio)
+   - Corazón Libre (2005, studio)
+   - Cantora (2009, studio - recorded in 2009)
+4. Final count: 3 studio albums
+FINAL ANSWER: 3