Update app.py

app.py

@@ -7,45 +7,58 @@ import redis
 import time
 import os
 from typing import Dict, List, Optional
-from redis.connection import ConnectionPool  # 引入连接池
+from redis.connection import ConnectionPool, SSLConnection  # 引入SSL连接类
 
 # ---------------------- 1. 初始化配置 ----------------------
 app = FastAPI(title="Distributed Agent Controller", version="1.0")
 
-# 跨域配置（适配Hugging Face Space环境）
+# 跨域配置
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # 生产环境替换为具体域名
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# ---------------------- 关键修复：Redis连接池 + 依赖注入 ----------------------
+# ---------------------- 修复Redis连接池（兼容低版本redis-py） ----------------------
 def create_redis_pool() -> ConnectionPool:
     try:
+        # 从环境变量获取配置（必须在Space中正确设置）
+        redis_host = os.getenv("REDIS_HOST")
+        redis_port = int(os.getenv("REDIS_PORT", 6380))
+        redis_password = os.getenv("REDIS_PASSWORD")
+        
+        if not all([redis_host, redis_password]):
+            raise ValueError("REDIS_HOST和REDIS_PASSWORD必须配置")
+
+        # 关键修复：使用SSLConnection类，替代ssl=True参数（兼容低版本redis-py）
         pool = ConnectionPool(
-            host=os.getenv("REDIS_HOST"),  # 填 Upstash 的 Host（xxx.upstash.io）
-            port=int(os.getenv("REDIS_PORT", 6380)),  # 固定填 6380（加密端口）
-            password=os.getenv("REDIS_PASSWORD"),  # 填 Upstash 的实例密码
+            host=redis_host,
+            port=redis_port,
+            password=redis_password,
             decode_responses=True,
-            ssl=True,  # 必须开启 SSL，Upstash 强制加密连接
+            connection_class=SSLConnection,  # 强制使用SSL加密连接
             socket_timeout=5,
-            retry_on_timeout=True
+            retry_on_timeout=True,
+            max_connections=10
         )
+        
+        # 验证连接
         with redis.Redis(connection_pool=pool) as client:
-            client.ping()  # 验证连接
+            client.ping()  # 成功会返回True
+        print("Redis连接池初始化成功")
         return pool
     except Exception as e:
-        raise RuntimeError(f"Redis connection failed: {str(e)}")
+        raise RuntimeError(f"Redis连接池初始化失败：{str(e)}")
 
-# 初始化连接池（全局仅创建一次）
+# 初始化连接池
 redis_pool = create_redis_pool()
 
+# 依赖注入：获取Redis客户端
 def get_redis_client() -> redis.Redis:
-    """依赖注入：每次请求创建独立客户端，请求结束自动关闭"""
     with redis.Redis(connection_pool=redis_pool) as client:
-        yield client  # 用yield确保客户端使用后自动释放
+        yield client
 
 # ---------------------- 2. API密钥验证 ----------------------
 def verify_api_key(api_key: Optional[str] = None):
@@ -61,7 +74,7 @@ PROCESSING_NODES = [
         "api_key": os.getenv("NODE_API_KEY", "node-key"),
         "health": "unknown",
         "last_check": 0,
-        "load": 0  # 节点负载计数
+        "load": 0
     },
     {
         "url": os.getenv("NODE_2_URL", "https://your-username-node2.hf.space"),
@@ -87,7 +100,6 @@ class SendMessageRequest(BaseModel):
 
 # ---------------------- 5. 工具函数 ----------------------
 async def check_node_health(node: Dict) -> bool:
-    """检查节点健康状态，5分钟缓存结果"""
     current_time = time.time()
     if current_time - node["last_check"] < 300:
         return node["health"] == "alive"
@@ -109,7 +121,6 @@ async def check_node_health(node: Dict) -> bool:
         return False
 
 def get_least_loaded_node() -> Optional[Dict]:
-    """选择负载最低的健康节点"""
     healthy_nodes = [n for n in PROCESSING_NODES if check_node_health(n)]
     if not healthy_nodes:
         return None
@@ -118,33 +129,29 @@ def get_least_loaded_node() -> Optional[Dict]:
 def task_key(user_id: str, task_id: str) -> str:
     return f"user:{user_id}:task:{task_id}"
 
-# ---------------------- 6. 核心接口（修复Redis调用逻辑） ----------------------
+# ---------------------- 6. 核心接口 ----------------------
 @app.post("/task/create")
 async def create_task(
     req: CreateTaskRequest,
     api_key: str = Depends(verify_api_key),
-    redis_client: redis.Redis = Depends(get_redis_client)  # 依赖注入Redis客户端
+    redis_client: redis.Redis = Depends(get_redis_client)
 ):
-    """创建新任务并生成初始响应"""
     task_id = f"task_{uuid.uuid4().hex[:8]}"
     task_data = {
         "task_id": task_id,
         "name": req.task_name,
-        "created": str(int(time.time())),  # 存储为字符串，避免Redis哈希类型转换问题
+        "created": str(int(time.time())),
         "updated": str(int(time.time())),
-        "history": '[{"role": "user", "content": "%s"}]' % req.initial_prompt.replace('"', '\\"')  # 转义��引号，避免JSON错误
+        "history": '[{"role": "user", "content": "%s"}]' % req.initial_prompt.replace('"', '\\"')
     }
 
-    # 保存任务到Redis（使用注入的客户端）
     redis_client.hset(task_key(req.user_id, task_id), mapping=task_data)
     redis_client.sadd(f"user:{req.user_id}:tasks", task_id)
 
-    # 分配节点处理
     node = get_least_loaded_node()
     if not node:
         raise HTTPException(status_code=503, detail="No available nodes")
 
-    # 调用节点生成代码
     try:
         node["load"] += 1
         async with httpx.AsyncClient(timeout=60) as client:
@@ -159,15 +166,14 @@ async def create_task(
         node["load"] -= 1
         raise HTTPException(status_code=500, detail=f"Node error: {str(e)}")
 
-    # 更新任务历史（修复eval风险，用json.loads替代）
-    import json  # 建议在文件顶部导入，此处为了突出修复点
+    import json
     result = resp.json()
-    history = json.loads(task_data["history"])  # 用json.loads替代eval，更安全
+    history = json.loads(task_data["history"])
     history.append({"role": "assistant", "content": result["code"].replace('"', '\\"')})
     redis_client.hset(
         task_key(req.user_id, task_id),
         mapping={
-            "history": json.dumps(history),  # 用json.dumps存储，避免格式错误
+            "history": json.dumps(history),
             "updated": str(int(time.time()))
         }
     )
@@ -184,26 +190,21 @@ async def send_message(
     api_key: str = Depends(verify_api_key),
     redis_client: redis.Redis = Depends(get_redis_client)
 ):
-    """向现有任务发送消息"""
     import json
     key = task_key(req.user_id, req.task_id)
     if not redis_client.exists(key):
         raise HTTPException(status_code=404, detail="Task not found")
 
-    # 获取历史对话（用json.loads替代eval）
     task_data = redis_client.hgetall(key)
     history = json.loads(task_data["history"])
     history.append({"role": "user", "content": req.message.replace('"', '\\"')})
 
-    # 构建上下文
     context = "\n".join([f"{m['role']}: {m['content']}" for m in history])
 
-    # 分配节点处理
     node = get_least_loaded_node()
     if not node:
         raise HTTPException(status_code=503, detail="No available nodes")
 
-    # 调用节点生成响应
     try:
         node["load"] += 1
         async with httpx.AsyncClient(timeout=60) as client:
@@ -222,7 +223,6 @@ async def send_message(
         node["load"] -= 1
         raise HTTPException(status_code=500, detail=f"Node error: {str(e)}")
 
-    # 更新历史
     result = resp.json()
     history.append({"role": "assistant", "content": result["code"].replace('"', '\\"')})
     redis_client.hset(
@@ -241,7 +241,6 @@ async def list_tasks(
     api_key: str = Depends(verify_api_key),
     redis_client: redis.Redis = Depends(get_redis_client)
 ):
-    """获取用户所有任务"""
     import json
     task_ids = redis_client.smembers(f"user:{user_id}:tasks")
     tasks = []
@@ -259,10 +258,8 @@ async def list_tasks(
 
 @app.get("/nodes/health")
 async def node_health(api_key: str = Depends(verify_api_key)):
-    """查看节点健康状态"""
     status = []
     for node in PROCESSING_NODES:
-        # 主动触发健康检查（更新最新状态）
         await check_node_health(node)
         status.append({
             "url": node["url"],