Update app.py

app.py

@@ -7,6 +7,7 @@ import redis
 import time
 import os
 from typing import Dict, List, Optional
+from redis.connection import ConnectionPool  # 引入连接池
 
 # ---------------------- 1. 初始化配置 ----------------------
 app = FastAPI(title="Distributed Agent Controller", version="1.0")
@@ -14,36 +15,48 @@ app = FastAPI(title="Distributed Agent Controller", version="1.0")
 # 跨域配置（适配Hugging Face Space环境）
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # Space环境下允许跨域调用
+    allow_origins=["*"],  # 生产环境替换为具体域名
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Redis连接（必须使用公网Redis，Space内无本地Redis）
-def get_redis_client():
+# ---------------------- 关键修复：Redis连接池 + 依赖注入 ----------------------
+def create_redis_pool() -> ConnectionPool:
+    """创建Redis连接池（替代全局客户端，避免析构问题）"""
     try:
-        return redis.Redis(
+        pool = ConnectionPool(
             host=os.getenv("REDIS_HOST"),
             port=int(os.getenv("REDIS_PORT", 6379)),
             password=os.getenv("REDIS_PASSWORD"),
             decode_responses=True,
             socket_timeout=5,
-            retry_on_timeout=True
+            retry_on_timeout=True,
+            max_connections=10  # 限制连接数，避免资源泄漏
         )
+        # 验证连接池可用性
+        with redis.Redis(connection_pool=pool) as client:
+            client.ping()
+        return pool
     except Exception as e:
-        raise RuntimeError(f"Redis connection failed: {str(e)}")
+        raise RuntimeError(f"Redis connection pool failed: {str(e)}")
 
-redis_client = get_redis_client()
+# 初始化连接池（全局仅创建一次）
+redis_pool = create_redis_pool()
 
-# API密钥验证
+def get_redis_client() -> redis.Redis:
+    """依赖注入：每次请求创建独立客户端，请求结束自动关闭"""
+    with redis.Redis(connection_pool=redis_pool) as client:
+        yield client  # 用yield确保客户端使用后自动释放
+
+# ---------------------- 2. API密钥验证 ----------------------
 def verify_api_key(api_key: Optional[str] = None):
     valid_key = os.getenv("CONTROLLER_API_KEY", "default-key")
     if api_key != valid_key:
         raise HTTPException(status_code=401, detail="Invalid API key")
     return api_key
 
-# 处理节点配置（从环境变量读取，支持动态扩展）
+# ---------------------- 3. 处理节点配置 ----------------------
 PROCESSING_NODES = [
     {
         "url": os.getenv("NODE_1_URL", "https://your-username-node1.hf.space"),
@@ -61,7 +74,7 @@ PROCESSING_NODES = [
     }
 ]
 
-# ---------------------- 2. 数据模型 ----------------------
+# ---------------------- 4. 数据模型 ----------------------
 class CreateTaskRequest(BaseModel):
     user_id: str
     task_name: str
@@ -74,7 +87,7 @@ class SendMessageRequest(BaseModel):
     max_new_tokens: int = 512
     temperature: float = 0.7
 
-# ---------------------- 3. 工具函数 ----------------------
+# ---------------------- 5. 工具函数 ----------------------
 async def check_node_health(node: Dict) -> bool:
     """检查节点健康状态，5分钟缓存结果"""
     current_time = time.time()
@@ -93,35 +106,38 @@ async def check_node_health(node: Dict) -> bool:
             return True
         node["health"] = "dead"
         return False
-    except:
+    except Exception as e:
         node["health"] = "dead"
         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
-    # 按当前负载排序，选择负载最低的
     return min(healthy_nodes, key=lambda x: x["load"])
 
 def task_key(user_id: str, task_id: str) -> str:
     return f"user:{user_id}:task:{task_id}"
 
-# ---------------------- 4. 核心接口 ----------------------
+# ---------------------- 6. 核心接口（修复Redis调用逻辑） ----------------------
 @app.post("/task/create")
-async def create_task(req: CreateTaskRequest, api_key: str = Depends(verify_api_key)):
+async def create_task(
+    req: CreateTaskRequest,
+    api_key: str = Depends(verify_api_key),
+    redis_client: redis.Redis = Depends(get_redis_client)  # 依赖注入Redis客户端
+):
     """创建新任务并生成初始响应"""
     task_id = f"task_{uuid.uuid4().hex[:8]}"
     task_data = {
         "task_id": task_id,
         "name": req.task_name,
-        "created": int(time.time()),
-        "updated": int(time.time()),
-        "history": [{"role": "user", "content": req.initial_prompt}]
+        "created": str(int(time.time())),  # 存储为字符串，避免Redis哈希类型转换问题
+        "updated": str(int(time.time())),
+        "history": '[{"role": "user", "content": "%s"}]' % req.initial_prompt.replace('"', '\\"')  # 转义双引号，避免JSON错误
     }
 
-    # 保存任务到Redis
+    # 保存任务到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)
 
@@ -132,26 +148,30 @@ async def create_task(req: CreateTaskRequest, api_key: str = Depends(verify_api_
 
     # 调用节点生成代码
     try:
-        node["load"] += 1  # 增加负载计数
+        node["load"] += 1
         async with httpx.AsyncClient(timeout=60) as client:
             resp = await client.post(
                 f"{node['url']}/generate/code",
                 json={"prompt": req.initial_prompt, "max_new_tokens": 512},
                 params={"api_key": node["api_key"]}
             )
-        node["load"] -= 1  # 减少负载计数
+        node["load"] -= 1
         resp.raise_for_status()
     except Exception as e:
         node["load"] -= 1
         raise HTTPException(status_code=500, detail=f"Node error: {str(e)}")
 
-    # 更新任务历史
+    # 更新任务历史（修复eval风险，用json.loads替代）
+    import json  # 建议在文件顶部导入，此处为了突出修复点
     result = resp.json()
-    history = eval(task_data["history"])  # 从字符串恢复列表
-    history.append({"role": "assistant", "content": result["code"]})
+    history = json.loads(task_data["history"])  # 用json.loads替代eval，更安全
+    history.append({"role": "assistant", "content": result["code"].replace('"', '\\"')})
     redis_client.hset(
         task_key(req.user_id, task_id),
-        mapping={"history": str(history), "updated": int(time.time())}
+        mapping={
+            "history": json.dumps(history),  # 用json.dumps存储，避免格式错误
+            "updated": str(int(time.time()))
+        }
     )
 
     return {
@@ -161,16 +181,21 @@ async def create_task(req: CreateTaskRequest, api_key: str = Depends(verify_api_
     }
 
 @app.post("/task/message")
-async def send_message(req: SendMessageRequest, api_key: str = Depends(verify_api_key)):
+async def send_message(
+    req: SendMessageRequest,
+    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 = eval(task_data["history"])
-    history.append({"role": "user", "content": req.message})
+    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])
@@ -201,17 +226,25 @@ async def send_message(req: SendMessageRequest, api_key: str = Depends(verify_ap
 
     # 更新历史
     result = resp.json()
-    history.append({"role": "assistant", "content": result["code"]})
+    history.append({"role": "assistant", "content": result["code"].replace('"', '\\"')})
     redis_client.hset(
         key,
-        mapping={"history": str(history), "updated": int(time.time())}
+        mapping={
+            "history": json.dumps(history),
+            "updated": str(int(time.time()))
+        }
     )
 
     return {"status": "success", "response": result["code"]}
 
 @app.get("/user/tasks/{user_id}")
-async def list_tasks(user_id: str, api_key: str = Depends(verify_api_key)):
+async def list_tasks(
+    user_id: str,
+    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 = []
     for tid in task_ids:
@@ -222,7 +255,7 @@ async def list_tasks(user_id: str, api_key: str = Depends(verify_api_key)):
                 "name": data["name"],
                 "created": time.ctime(int(data["created"])),
                 "updated": time.ctime(int(data["updated"])),
-                "history_length": len(eval(data["history"]))
+                "history_length": len(json.loads(data["history"]))
             })
     return {"status": "success", "tasks": tasks}
 
@@ -231,11 +264,12 @@ 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"],
             "health": node["health"],
             "load": node["load"],
             "last_check": time.ctime(node["last_check"]) if node["last_check"] else "Never"
         })
-    return {"status": "success", "nodes": status}
-    
+    return {"status": "success", "nodes": status}