Feat/gmy/support bohrium ak

agents/DPA_Agent/agent/agent.py

@@ -1,19 +1,13 @@
-import asyncio
-import os
+import nest_asyncio
 import sys
-from pathlib import Path
-from typing import Any, Dict
 
-import nest_asyncio
 from dotenv import load_dotenv
+
 from dp.agent.adapter.adk import CalculationMCPToolset
-from google.adk import Agent
-from google.adk.agents import LlmAgent, SequentialAgent
+from google.adk.agents import LlmAgent
 from google.adk.models.lite_llm import LiteLlm
-from google.adk.runners import Runner
-from google.adk.sessions import InMemorySessionService
 from google.adk.tools.mcp_tool.mcp_session_manager import SseServerParams
-from google.genai import types
+from pathlib import Path
 
 # Add the project root to Python path
 project_root = Path(__file__).resolve().parent.parent
@@ -23,91 +17,84 @@
 load_dotenv()
 nest_asyncio.apply()
 
-# Global Configuration
-BOHRIUM_EXECUTOR_CALC = {
-    "type": "dispatcher",
-    "machine": {
-        "batch_type": "Bohrium",
-        "context_type": "Bohrium",
-        "remote_profile": {
-            "email": os.getenv("BOHRIUM_EMAIL"),
-            "password": os.getenv("BOHRIUM_PASSWORD"),
-            "program_id": int(os.getenv("BOHRIUM_PROJECT_ID")),
-            "input_data": {
-                "image_name": "registry.dp.tech/dptech/dp/native/prod-19853/dpa-mcp:dev-0704",
+def create_root_agent(ak: str, app_key: str, project_id: int) -> LlmAgent:
+    """动态创建 agent，使用传入的 access_key 和 app_key"""
+    BOHRIUM_EXECUTOR_CALC = {
+        "type": "dispatcher",
+        "machine": {
+            "batch_type": "OpenAPI",
+            "context_type": "OpenAPI",
+            "remote_profile": {
+                "access_key": ak,
+                "app_key": app_key,  # 使用传入的 app_key
+                "project_id": project_id,
+                "image_address": "registry.dp.tech/dptech/dp/native/prod-19853/dpa-mcp:dev-0704",
                 "job_type": "container",
                 "platform": "ali",
-                "scass_type": "1 * NVIDIA V100_32g"
+                "machine_type": "1 * NVIDIA V100_32g"
             }
         }
     }
-}
-BOHRIUM_EXECUTOR_TE = {
-    "type": "dispatcher",
-    "machine": {
-        "batch_type": "Bohrium",
-        "context_type": "Bohrium",
-        "remote_profile": {
-            "email": os.getenv("BOHRIUM_EMAIL"),
-            "password": os.getenv("BOHRIUM_PASSWORD"),
-            "program_id": int(os.getenv("BOHRIUM_PROJECT_ID")),
-            "input_data": {
-                "image_name": "registry.dp.tech/dptech/dp/native/prod-435364/agents:0.1.0",
+
+    BOHRIUM_EXECUTOR_TE = {
+        "type": "dispatcher",
+        "machine": {
+            "batch_type": "OpenAPI",
+            "context_type": "OpenAPI",
+            "remote_profile": {
+                "access_key": ak,
+                "app_key": app_key,  # 使用传入的 app_key
+                "project_id": project_id,
+                "image_address": "registry.dp.tech/dptech/dp/native/prod-435364/agents:0.1.0",
                 "job_type": "container",
                 "platform": "ali",
-                "scass_type": "1 * NVIDIA V100_32g"
+                "machine_type": "1 * NVIDIA V100_32g"
             }
         }
     }
-}
-LOCAL_EXECUTOR = {
-    "type": "local"
-}
-HTTPS_STORAGE = {
-  "type": "https",
-  "plugin": {
-        "type": "bohrium",
-        "username": os.getenv("BOHRIUM_EMAIL"),
-        "password": os.getenv("BOHRIUM_PASSWORD"),
-        "project_id": int(os.getenv("BOHRIUM_PROJECT_ID"))
+
+    HTTPS_STORAGE = {
+        "type": "https",
+        "plugin": {
+            "type": "bohrium",
+            "access_key": ak,
+            "app_key": app_key,  # 使用传入的 app_key
+            "project_id": project_id
+        }
     }
-}
 
-print('-----', HTTPS_STORAGE)
+    # 其余代码保持不变...
 
+    mcp_tools_dpa = CalculationMCPToolset(
+        connection_params=SseServerParams(url="https://dpa-uuid1750659890.app-space.dplink.cc/sse?token=a0d87a7abf8d47cb92403018fd4a9468"),
+        storage=HTTPS_STORAGE,
+        executor=BOHRIUM_EXECUTOR_CALC,
+        executor_map={
+            "build_bulk_structure": None,
+            "build_molecule_structure": None,
+            "build_surface_slab": None,
+            "build_surface_adsorbate": None
+        }
+    )
+    mcp_tools_thermoelectric = CalculationMCPToolset(
+        connection_params=SseServerParams(url="https://thermoelectricmcp000-uuid1750905361.app-space.dplink.cc/sse?token=1c1f2140a5504ebcb680f6a7fa2c03db"),
+        storage=HTTPS_STORAGE,
+        executor=BOHRIUM_EXECUTOR_TE
+    )
 
-mcp_tools_dpa = CalculationMCPToolset(
-    connection_params=SseServerParams(url="https://dpa-uuid1750659890.app-space.dplink.cc/sse?token=a0d87a7abf8d47cb92403018fd4a9468"),
-    # connection_params=SseServerParams(url="http://pfmx1355864.bohrium.tech:50002/sse"),
-    storage=HTTPS_STORAGE,
-    executor=None,
-    executor_map={
-        "build_structure": None
-    }
-)
-mcp_tools_thermoelectric = CalculationMCPToolset(
-    connection_params=SseServerParams(url="https://thermoelectricmcp000-uuid1750905361.app-space.dplink.cc/sse?token=1c1f2140a5504ebcb680f6a7fa2c03db"),
-    storage=HTTPS_STORAGE,
-    executor=BOHRIUM_EXECUTOR_TE
-)
-# mcp_tools_superconductor = CalculationMCPToolset(
-#     connection_params=SseServerParams(url="https://superconductor-ambient-010-uuid1750845273.app-space.dplink.cc/sse?token=57578d394b564682943a723697f992b1"),
-#     storage=HTTPS_STORAGE,
-#     executor=None
-# )
-root_agent = LlmAgent(
-    model=LiteLlm(model="deepseek/deepseek-chat"),
-    name="dpa_agent",
-    description="An agent specialized in computational research using Deep Potential",
-    instruction=(
-        "You are an expert in materials science and computational chemistry. "
-        "Help users perform Deep Potential calculations including structure optimization, molecular dynamics and property calculations. "
-        "Use default parameters if the users do not mention, but let users confirm them before submission. "
-        "Always verify the input parameters to users and provide clear explanations of results."
-    ),
-    tools=[
-        mcp_tools_dpa, 
-        mcp_tools_thermoelectric, 
-        # mcp_tools_superconductor
-    ],
-)
+    return LlmAgent(
+        model=LiteLlm(model="deepseek/deepseek-chat"),
+        name="dpa_agent",
+        description="An agent specialized in computational research using Deep Potential",
+        instruction=(
+            "You are an expert in materials science and computational chemistry. "
+            "Help users perform Deep Potential calculations including structure optimization, molecular dynamics and property calculations. "
+            "Use default parameters if the users do not mention, but let users confirm them before submission. "
+            "In multi-step workflows involving file outputs, always use the URI of the previous step's file as the input for the next tool. "
+            "Always verify the input parameters to users and provide clear explanations of results."
+        ),
+        tools=[
+            mcp_tools_dpa,
+            mcp_tools_thermoelectric,
+        ],
+    )

agents/DPA_Agent/config/agent_config.py

@@ -53,15 +53,10 @@ def _get_default_config(self) -> Dict[str, Any]:
 
     def get_agent(self):
         """Dynamically import and return the configured agent"""
-        agentconfig = self.config.get("agent", {})
-        module_name = agentconfig.get("module", "agent.subagent")
-        agentname = agentconfig.get("rootAgent", "rootagent")
-
-        try:
-            module = importlib.import_module(module_name)
-            return getattr(module, agentname)
-        except (ImportError, AttributeError) as e:
-            raise ImportError(f"Failed to load agent {agentname} from {module_name}: {e}")
+        raise RuntimeError(
+            "Agent 必须在运行时动态创建（需要 access_key），"
+            "请使用 agent.create_agent(ak, project_id) 而不是 config.get_agent()"
+        )
 
     def get_ui_config(self) -> Dict[str, Any]:
         """Get UI-specific configuration"""

agents/DPA_Agent/websocket-server.py

@@ -20,6 +20,7 @@
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse, PlainTextResponse
+from http.cookies import SimpleCookie
 from starlette.middleware.base import BaseHTTPMiddleware
 import uvicorn
 
@@ -29,9 +30,8 @@
 
 # Import configuration
 from config.agent_config import agentconfig
+from agent.agent import create_root_agent
 
-# Get agent from configuration
-rootagent = agentconfig.get_agent()
 
 # 配置日志
 logging.basicConfig(level=logging.INFO)
@@ -112,7 +112,7 @@ async def dispatch(self, request: Request, call_next):
 
 class ConnectionContext:
     """每个WebSocket连接的独立上下文"""
-    def __init__(self, websocket: WebSocket):
+    def __init__(self, websocket: WebSocket, app_key: str, access_key: str):
         self.websocket = websocket
         self.sessions: Dict[str, Session] = {}
         self.runners: Dict[str, Runner] = {}
@@ -124,6 +124,8 @@ def __init__(self, websocket: WebSocket):
         }
         # 为每个连接生成唯一的user_id
         self.user_id = f"user_{uuid.uuid4().hex[:8]}"
+        self.access_key = access_key 
+        self.app_key = app_key
 
 class SessionManager:
     def __init__(self):
@@ -157,6 +159,13 @@ def handle_init_error(future):
     async def _init_session_runner(self, context: ConnectionContext, session_id: str):
         """异步初始化会话的runner"""
         try:
+            project_id = int(os.getenv("BOHRIUM_PROJECT_ID", 27666))  # 可改为从配置读
+            # ✅ 使用 context 中的 access_key 和 app_key 动态创建 agent
+            user_agent = create_root_agent(
+                ak=context.access_key, 
+                app_key=context.app_key,
+                project_id=project_id
+            )
             session_service = InMemorySessionService()
             await session_service.create_session(
                 app_name=self.app_name,
@@ -165,19 +174,15 @@ async def _init_session_runner(self, context: ConnectionContext, session_id: str
             )
 
             runner = Runner(
-                agent=rootagent,
+                agent=user_agent,  # ✅ 使用动态创建的 agent
                 session_service=session_service,
                 app_name=self.app_name
             )
-
             context.session_services[session_id] = session_service
             context.runners[session_id] = runner
-
             logger.info(f"Runner 初始化完成: {session_id}")
-
         except Exception as e:
             logger.error(f"初始化Runner失败: {e}")
-            # 清理失败的会话
             if session_id in context.sessions:
                 del context.sessions[session_id]
             if session_id in context.session_services:
@@ -213,20 +218,18 @@ async def switch_session(self, context: ConnectionContext, session_id: str) -> b
             return True
         return False
 
-    async def connect_client(self, websocket: WebSocket):
+    async def connect_client(self, websocket: WebSocket, access_key: str, app_key: str):
         """连接新客户端"""
         await websocket.accept()
-        
+
         # 为新连接创建独立的上下文
-        context = ConnectionContext(websocket)
+        context = ConnectionContext(websocket, app_key, access_key)
         self.active_connections[websocket] = context
-
-        logger.info(f"新用户连接: {context.user_id}")
-
+        logger.info(f"新用户连接: {context.user_id}, AK: {access_key[:4]}...{access_key[-4:]}, AppKey: {app_key}")
+
         # 创建默认会话
         session = await self.create_session(context)
         context.current_session_id = session.id
-
         # 发送初始会话信息
         await self.send_sessions_list(context)
 
@@ -506,11 +509,59 @@ async def process_message(self, context: ConnectionContext, message: str):
 # 创建全局管理器
 manager = SessionManager()
 
+def get_ak_info_from_headers(headers: list) -> tuple:
+    """从请求头中获取access key和app key"""
+    cookie_header = None
+    for k, v in headers:
+        if k.lower() == b"cookie":  # 使用小写比较，确保兼容性
+            cookie_header = v.decode("utf-8")
+            break
+
+    access_key = None
+    app_key = None
+
+    if cookie_header:
+        try:
+            cookie = SimpleCookie()
+            cookie.load(cookie_header)
+            if "appAccessKey" in cookie:
+                access_key = cookie["appAccessKey"].value
+                logger.info(f"从Cookie获取Access Key: {access_key[:4]}...{access_key[-4:]}")
+            if "clientName" in cookie:
+                app_key = cookie["clientName"].value
+                logger.info(f"从Cookie获取App Key: {app_key}")
+        except Exception as e:
+            logger.warning(f"解析Cookie失败: {e}")
+
+    # 如果没有从Cookie获取到，尝试从Authorization头获取
+    if not access_key:
+        for k, v in headers:
+            if k.lower() == b"authorization":
+                auth_header = v.decode("utf-8")
+                if auth_header.startswith("Bearer "):
+                    access_key = auth_header[7:]
+                    logger.info(f"从Authorization头获取Access Key: {access_key[:4]}...{access_key[-4:]}")
+
+    if not app_key:
+        app_key = ""
+        logger.info("未在Cookie中找到clientName，使用默认App Key: mat-master")
+
+    if not access_key:
+        logger.warning("未找到Access Key")
+        return "", app_key
+
+    return access_key, app_key
+
 @app.websocket("/ws")
 async def websocket_endpoint(websocket: WebSocket):
     """WebSocket 端点"""
-    await manager.connect_client(websocket)
+    access_key, app_key = get_ak_info_from_headers(websocket.scope.get("headers", []))
+    if not access_key:
+        logger.warning("WebSocket 连接缺少 appAccessKey")
+        await websocket.close(code=1008, reason="Missing appAccessKey")
+        return
 
+    await manager.connect_client(websocket, access_key, app_key)
     # 获取该连接的上下文
     context = manager.active_connections.get(websocket)
     if not context: