update

agents/root_agent.py

@@ -3,108 +3,43 @@ import asyncio
 import os
 import json
 import httpx
-import redis
 
 from .agent_manager_client import AgentManagerClient, AgentInfo, CreateAgentRequest
 from .mcp_agent_interface import MCPAgent # 导入 MCP Agent 接口
 
-# Redis 配置
-REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
-REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
-REDIS_DB = int(os.getenv("REDIS_DB", 0))
-AGENT_CHANNEL = "agent_discovery_channel"
-AGENT_KEY_PREFIX = "agent:"
-
 class RootAgent:
     def __init__(self):
         AGENT_MANAGER_BASE_URL = os.getenv("AGENT_MANAGER_BASE_URL", "http://localhost:7860")
         self.agent_manager_client = AgentManagerClient(base_url=AGENT_MANAGER_BASE_URL)
         self.active_agents: Dict[str, AgentInfo] = {} # 内存缓存活跃 Agent
-        self.redis_client: Optional[redis.Redis] = None
-        self.redis_pubsub: Optional[redis.client.PubSub] = None
-        self.discovery_task: Optional[asyncio.Task] = None
 
     async def initialize(self):
         """
-        RootAgent 启动时的初始化逻辑，连接 Redis 并启动 Agent 发现。
+        RootAgent 启动时的初始化逻辑，通过 Agent Manager 发现 Agent。
         """
-        self.redis_client = redis.Redis(host=REDIS_HOST, port=REDIS_PORT, db=REDIS_DB)
-        try:
-            self.redis_client.ping()
-            print("RootAgent connected to Redis successfully!")
-        except redis.exceptions.ConnectionError as e:
-            print(f"RootAgent could not connect to Redis: {e}")
-            # 可以在这里选择退出或以降级模式运行
-
-        self.redis_pubsub = self.redis_client.pubsub()
-        self.redis_pubsub.subscribe(AGENT_CHANNEL)
-        self.discovery_task = asyncio.create_task(self._agent_discovery_listener())
-
-        # 从 Redis 加载所有已知的 Agent
-        await self._load_existing_agents()
+        await self._refresh_active_agents()
         print("RootAgent initialized.")
 
     async def shutdown(self):
         """
         RootAgent 关闭时的清理逻辑。
         """
-        if self.discovery_task:
-            self.discovery_task.cancel()
-            try:
-                await self.discovery_task
-            except asyncio.CancelledError:
-                pass
-        if self.redis_pubsub:
-            self.redis_pubsub.unsubscribe(AGENT_CHANNEL)
-        if self.redis_client:
-            self.redis_client.close()
         print("RootAgent shutdown.")
 
-    async def _agent_discovery_listener(self):
+    async def _refresh_active_agents(self):
         """
-        监听 Redis 频道，实时更新活跃 Agent 列表。
+        从 Agent Manager 加载所有活跃的 Agent 到内存缓存。
         """
-        if not self.redis_pubsub:
-            return
-
-        while True:
-            message = self.redis_pubsub.get_message(ignore_subscribe_messages=True)
-            if message:
-                try:
-                    data = json.loads(message['data'].decode('utf-8'))
-                    agent_id = data.get("id")
-                    event_type = data.get("event_type")
-
-                    if agent_id:
-                        if event_type == "HEARTBEAT" or event_type == "ONLINE":
-                            # 从 Redis 获取最新的 Agent 信息
-                            agent_data = self.redis_client.hgetall(f"{AGENT_KEY_PREFIX}{agent_id}")
-                            if agent_data:
-                                agent_info = AgentInfo(**{k.decode('utf-8'): v.decode('utf-8') for k, v in agent_data.items()})
-                                self.active_agents[agent_id] = agent_info
-                                # print(f"RootAgent discovered/updated agent: {agent_id} ({agent_info.agent_type})")
-                            else:
-                                print(f"RootAgent received {event_type} for {agent_id}, but no data in Redis.")
-                        elif event_type == "OFFLINE":
-                            if agent_id in self.active_agents:
-                                del self.active_agents[agent_id]
-                                print(f"RootAgent removed offline agent: {agent_id}")
-                except json.JSONDecodeError:
-                    print(f"RootAgent received invalid JSON message from Redis: {message['data']}")
-                except Exception as e:
-                    print(f"RootAgent error processing Redis message: {e}")
-            await asyncio.sleep(0.1)
+        try:
+            agents = await self.agent_manager_client.list_agents()
+            self.active_agents = {agent.id: agent for agent in agents}
+            print(f"RootAgent refreshed active agents. Found {len(self.active_agents)} agents.")
+        except httpx.HTTPStatusError as e:
+            print(f"Failed to refresh agents from Agent Manager: {e.response.text}")
+            # 可以在这里选择抛出异常或以降级模式运行
+        except Exception as e:
+            print(f"An unexpected error occurred while refreshing agents: {e}")
 
-    async def _load_existing_agents(self):
-        """
-        从 Redis 加载所有已知的 Agent 到内存缓存。
-        """
-        for key in self.redis_client.keys(f"{AGENT_KEY_PREFIX}*"):
-            agent_data = self.redis_client.hgetall(key)
-            if agent_data:
-                agent_info = AgentInfo(**{k.decode('utf-8'): v.decode('utf-8') for k, v in agent_data.items()})
-                self.active_agents[agent_info.id] = agent_info
-                print(f"RootAgent loaded existing agent: {agent_info.id} ({agent_info.agent_type})")
 
     async def _get_or_create_agent(self, agent_type: str, image_name: str = "echo-agent:latest") -> AgentInfo:
         """
@@ -121,7 +56,7 @@ class RootAgent:
         create_request = CreateAgentRequest(
             agent_type=agent_type,
             image_name=image_name,
-            env_vars={"REDIS_HOST": REDIS_HOST, "REDIS_PORT": str(REDIS_PORT), "REDIS_DB": str(REDIS_DB)},
+            env_vars={}, # 移除 Redis 相关的环境变量，因为 RootAgent 不再直接使用 Redis
             config={"some_initial_config": "value"}
         )
         try:

app.py

@@ -1,30 +1,89 @@
 from fastapi import FastAPI, HTTPException, status
-from typing import Dict, Any, Optional
+from typing import Dict, Any, Optional, List
 import asyncio
 import os
+import requests
+from dotenv import load_dotenv
+import contextlib
+from pydantic import BaseModel
 
 from agents.root_agent import RootAgent
 
+# 定义任务子项的模型
+class Subtask(BaseModel):
+    agent_type: str
+    payload: Dict[str, Any]
+
+# 定义 process_task endpoint 的请求体模型
+class ProcessTaskRequest(BaseModel):
+    subtasks: List[Subtask]
+    global_config: Optional[Dict[str, Any]] = None # 可选的全局配置
+
 app = FastAPI()
 root_agent_instance: Optional[RootAgent] = None
 
-@app.on_event("startup")
-async def startup_event():
-    global root_agent_instance
+# Upstash Redis 相关的全局变量
+UPSTASH_REDIS_REST_URL: Optional[str] = None
+UPSTASH_REDIS_REST_TOKEN: Optional[str] = None
+
+@contextlib.asynccontextmanager
+async def lifespan(app: FastAPI):
+    global root_agent_instance, UPSTASH_REDIS_REST_URL, UPSTASH_REDIS_REST_TOKEN
+
+    # 加载 .env 文件中的环境变量
+    load_dotenv()
+    UPSTASH_REDIS_REST_URL = os.getenv("UPSTASH_REDIS_REST_URL")
+    UPSTASH_REDIS_REST_TOKEN = os.getenv("UPSTASH_REDIS_REST_TOKEN")
+
+    if not UPSTASH_REDIS_REST_URL or not UPSTASH_REDIS_REST_TOKEN:
+        print("错误：UPSTASH_REDIS_REST_URL 或 UPSTASH_REDIS_REST_TOKEN 未设置。请检查 .env 文件。")
+        # 可以在这里选择抛出异常或继续，取决于应用程序对 Redis 的依赖程度
+        # 为了演示，我们选择继续，但会在 Redis 操作时检查这些变量
+    else:
+        print("Upstash Redis 环境变量已加载。")
+        # 可以在这里添加一个简单的 Redis 连接测试
+        headers = {
+            "Authorization": f"Bearer {UPSTASH_REDIS_REST_TOKEN}",
+            "Content-Type": "application/json"
+        }
+        try:
+            # 尝试执行一个简单的 PING 命令来测试连接
+            ping_url = f"{UPSTASH_REDIS_REST_URL}/ping"
+            response = requests.post(ping_url, headers=headers)
+            response.raise_for_status()
+            ping_result = response.json()
+            if ping_result and ping_result.get('result') == "PONG":
+                print("Upstash Redis 连接测试成功！")
+            else:
+                print(f"Upstash Redis 连接测试失败: {ping_result}")
+        except Exception as e:
+            print(f"Upstash Redis 连接测试发生错误: {e}")
+
+    # 初始化 RootAgent
     root_agent_instance = RootAgent()
-    await root_agent_instance.initialize()
+    try:
+        await root_agent_instance.initialize()
+        print("RootAgent 已初始化。")
+    except Exception as e:
+        print(f"RootAgent 初始化失败: {e}")
+        # 可以在这里选择抛出异常或以降级模式运行
+        raise # 重新抛出异常以阻止应用程序启动，直到问题解决
 
-@app.on_event("shutdown")
-async def shutdown_event():
+    yield # 应用程序在此处启动
+
+    # 应用程序关闭时执行的清理工作
     if root_agent_instance:
         await root_agent_instance.shutdown()
+        print("RootAgent 已关闭。")
+
+app = FastAPI(lifespan=lifespan)
 
 @app.get("/")
 async def read_root():
     return {"message": "Worker service with RootAgent is running!"}
 
 @app.post("/process_task")
-async def process_task_endpoint(task_data: Dict[str, Any]):
+async def process_task_endpoint(request: ProcessTaskRequest):
     """
     接收外部任务请求，并交由 RootAgent 处理。
     """
@@ -32,6 +91,8 @@ async def process_task_endpoint(task_data: Dict[str, Any]):
         raise HTTPException(status_code=500, detail="RootAgent not initialized.")
     
     try:
+        # 将 Pydantic 模型转换为 RootAgent 期望的 Dict 格式
+        task_data = request.dict()
         result = await root_agent_instance.process_task(task_data)
         return result
     except Exception as e:

requirements.txt

@@ -3,3 +3,4 @@ uvicorn[standard]
 qwen-agent
 python-dotenv
 python-dateutil
+requests

upstash_redis_test.py

@@ -0,0 +1,109 @@
+import os
+import requests
+from dotenv import load_dotenv
+
+# 加载 .env 文件中的环境变量
+load_dotenv()
+
+UPSTASH_REDIS_REST_URL = os.getenv("UPSTASH_REDIS_REST_URL")
+UPSTASH_REDIS_REST_TOKEN = os.getenv("UPSTASH_REDIS_REST_TOKEN")
+
+if not UPSTASH_REDIS_REST_URL or not UPSTASH_REDIS_REST_TOKEN:
+    print("错误：UPSTASH_REDIS_REST_URL 或 UPSTASH_REDIS_REST_TOKEN 未设置。请检查 .env 文件。")
+    exit(1)
+
+headers = {
+    "Authorization": f"Bearer {UPSTASH_REDIS_REST_TOKEN}",
+    "Content-Type": "application/json"
+}
+
+def execute_redis_command(command, *args):
+    url = ""
+    send_json_payload = False
+    payload = []
+
+    if command.upper() == "SET":
+        if len(args) != 2:
+            print("SET 命令需要两个参数：key 和 value。")
+            return None
+        key, value = args
+        url = f"{UPSTASH_REDIS_REST_URL}/set/{key}/{value}"
+    elif command.upper() == "GET":
+        if len(args) != 1:
+            print("GET 命令需要一个参数：key。")
+            return None
+        key = args[0]
+        url = f"{UPSTASH_REDIS_REST_URL}/get/{key}"
+    elif command.upper() == "DEL":
+        if len(args) != 1:
+            print("DEL 命令需要一个参数：key。")
+            return None
+        key = args[0]
+        url = f"{UPSTASH_REDIS_REST_URL}/del/{key}"
+    else:
+        # 对于其他命令，使用通用格式，并发送 JSON 数组作为 payload
+        url = f"{UPSTASH_REDIS_REST_URL}/{command}"
+        payload = [arg for arg in args]
+        send_json_payload = True
+
+    try:
+        if send_json_payload:
+            response = requests.post(url, headers=headers, json=payload)
+        else:
+            response = requests.post(url, headers=headers) # 不发送 json=payload
+        response.raise_for_status()  # 如果请求失败，抛出 HTTPError
+        return response.json()
+    except requests.exceptions.HTTPError as http_err:
+        print(f"HTTP 错误发生: {http_err}")
+        print(f"响应内容: {response.text}")
+    except requests.exceptions.ConnectionError as conn_err:
+        print(f"连接错误发生: {conn_err}")
+    except requests.exceptions.Timeout as timeout_err:
+        print(f"请求超时: {timeout_err}")
+    except requests.exceptions.RequestException as req_err:
+        print(f"发生未知错误: {req_err}")
+    return None
+
+def test_redis_connection():
+    print("正在测试 Upstash Redis 连接...")
+
+    # SET 命令
+    key = "mykey"
+    value = "hello_upstash"
+    print(f"正在设置键 '{key}' 为值 '{value}'...")
+    set_result = execute_redis_command("SET", key, value)
+    if set_result:
+        print(f"SET 结果: {set_result}")
+        if set_result and set_result.get('result') == "OK":
+            print("SET 命令成功。")
+        else:
+            print("SET 命令失败。")
+            return # 如果 SET 失败，才返回
+
+    # GET 命令
+    print(f"正在获取键 '{key}' 的值...")
+    get_result = execute_redis_command("GET", key)
+    if get_result is not None:
+        print(f"GET 结果: {get_result}")
+        if get_result and get_result.get('result') == value:
+            print("GET 命令成功，值匹配。")
+        else:
+            print("GET 命令失败，值不匹配或结果格式不正确。")
+    else:
+        print("GET 命令失败。")
+        return
+
+    # DEL 命令 (可选，用于清理)
+    print(f"正在删除键 '{key}'...")
+    del_result = execute_redis_command("DEL", key)
+    if del_result is not None:
+        print(f"DEL 结果: {del_result}")
+        if del_result and del_result.get('result') == 1:
+            print("DEL 命令成功。")
+        else:
+            print("DEL 命令失败或结果格式不正确。")
+    else:
+        print("DEL 命令失败。")
+
+if __name__ == "__main__":
+    test_redis_connection()