bohr-agent-sdk 0.1.101__py3-none-any.whl → 0.1.103__py3-none-any.whl

dp/agent/cli/templates/ui/server/user_files.py

@@ -0,0 +1,85 @@
+"""
+User file manager
+"""
+from pathlib import Path
+
+
+class UserFileManager:
+    """Manage user-specific file directories"""
+    
+    def __init__(self, base_dir: str, sessions_dir: str = ".agent_sessions"):
+        self.base_dir = Path(base_dir)
+        # Support custom sessions directory, can be absolute or relative path
+        sessions_path = Path(sessions_dir)
+        if sessions_path.is_absolute():
+            self.sessions_dir = sessions_path
+        else:
+            self.sessions_dir = self.base_dir / sessions_dir
+        self.user_sessions_dir = self.sessions_dir / "user_sessions"
+        self.temp_sessions_dir = self.sessions_dir / "temp_sessions"
+        
+        # Ensure directories exist
+        self.user_sessions_dir.mkdir(parents=True, exist_ok=True)
+        self.temp_sessions_dir.mkdir(parents=True, exist_ok=True)
+    
+    def _get_user_dir(self, user_id: str) -> str:
+        """Get user directory name"""
+        # Use user_id directly as directory name
+        return user_id
+    
+    def get_user_files_dir(self, user_id: str) -> Path:
+        """Get user's file directory
+        
+        Args:
+            user_id: User's unique identifier (Bohrium user_id or temporary user ID)
+            
+        Returns:
+            User's file directory path
+        """
+        if not user_id:
+            # If no user_id, use default directory
+            user_files_dir = self.temp_sessions_dir / "default" / "files"
+        elif user_id.startswith("user_"):
+            # Temporary user (generated ID starts with user_)
+            user_files_dir = self.temp_sessions_dir / user_id / "files"
+        else:
+            # Registered user (Bohrium user_id)
+            user_dir_name = self._get_user_dir(user_id)
+            user_files_dir = self.user_sessions_dir / user_dir_name / "files"
+        
+        # Ensure directory exists
+        user_files_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Create default output subdirectory
+        output_dir = user_files_dir / "output"
+        output_dir.mkdir(exist_ok=True)
+        
+        return user_files_dir
+    
+    def cleanup_temp_files(self, max_age_days: int = 7):
+        """Clean up expired temporary user files
+        
+        Args:
+            max_age_days: Maximum days to keep files
+        """
+        import time
+        import shutil
+        
+        current_time = time.time()
+        max_age_seconds = max_age_days * 24 * 60 * 60
+        
+        # Only clean temporary sessions directory
+        if not self.temp_sessions_dir.exists():
+            return
+        
+        for session_dir in self.temp_sessions_dir.iterdir():
+            if not session_dir.is_dir():
+                continue
+            
+            # Check directory modification time
+            dir_mtime = session_dir.stat().st_mtime
+            if current_time - dir_mtime > max_age_seconds:
+                try:
+                    shutil.rmtree(session_dir)
+                except Exception:
+                    pass

dp/agent/cli/templates/ui/server/utils.py

@@ -0,0 +1,50 @@
+"""
+Utility functions
+"""
+import os
+from typing import Tuple
+from http.cookies import SimpleCookie
+import socket
+
+def get_ak_info_from_request(headers) -> Tuple[str, str]:
+    """Extract AK info from request headers
+    
+    Priority:
+    1. Get from Cookie (production environment)
+    2. Get from environment variables (development debugging)
+    3. Return empty string to allow user custom input (commented out restriction)
+    """
+    # First try to get from cookie
+    cookie_header = headers.get("cookie", "")
+    if cookie_header:
+        simple_cookie = SimpleCookie()
+        simple_cookie.load(cookie_header)
+        
+        access_key = ""
+        app_key = ""
+        
+        if "appAccessKey" in simple_cookie:
+            access_key = simple_cookie["appAccessKey"].value
+        if "clientName" in simple_cookie:
+            app_key = simple_cookie["clientName"].value
+            
+        # If got valid values from cookie, return directly
+        if access_key or app_key:
+            return access_key, app_key
+    
+    # If not in cookie, try to get from environment variables (for development debugging)
+    access_key = os.environ.get("BOHR_ACCESS_KEY", "")
+    app_key = os.environ.get("BOHR_APP_KEY", "")
+    
+    return access_key, app_key
+
+
+def check_port_available(port: int) -> bool:
+    """Check if port is available"""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    try:
+        sock.bind(('', port))
+        sock.close()
+        return True
+    except OSError:
+        return False

dp/agent/cli/templates/ui/test_download.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""测试文件下载功能"""
+
+import os
+import sys
+import time
+import tempfile
+import requests
+from pathlib import Path
+
+def test_download_api():
+    """测试下载 API"""
+    
+    # API 基础地址
+    base_url = "http://localhost:8001"
+    
+    print("=" * 60)
+    print("文件下载功能测试")
+    print("=" * 60)
+    
+    # 1. 创建测试文件
+    test_dir = Path(tempfile.gettempdir()) / "download_test"
+    test_dir.mkdir(exist_ok=True)
+    
+    # 创建测试文件
+    test_file = test_dir / "test.txt"
+    test_file.write_text("This is a test file for download functionality.\n测试文件下载功能。")
+    
+    test_json = test_dir / "data.json"
+    test_json.write_text('{"name": "test", "value": 123}')
+    
+    # 创建子目录和文件
+    sub_dir = test_dir / "subdir"
+    sub_dir.mkdir(exist_ok=True)
+    (sub_dir / "file1.txt").write_text("File 1 content")
+    (sub_dir / "file2.txt").write_text("File 2 content")
+    
+    print(f"✅ 测试文件已创建在: {test_dir}")
+    print()
+    
+    # 2. 测试单文件下载
+    print("测试单文件下载...")
+    try:
+        # 模拟文件下载请求
+        file_path = str(test_file)
+        download_url = f"{base_url}/api/download/file{file_path}"
+        
+        print(f"  下载 URL: {download_url}")
+        print(f"  预期结果: 文件应该可以正常下载")
+        print()
+        
+        # 注意：实际测试需要运行服务器并通过浏览器或 curl 测试
+        print("  💡 请在浏览器中打开以下链接测试下载:")
+        print(f"     {download_url}")
+        print()
+        
+    except Exception as e:
+        print(f"  ❌ 错误: {e}")
+        print()
+    
+    # 3. 测试文件夹下载
+    print("测试文件夹下载...")
+    try:
+        # 模拟文件夹下载请求
+        folder_path = str(sub_dir)
+        download_url = f"{base_url}/api/download/folder{folder_path}"
+        
+        print(f"  下载 URL: {download_url}")
+        print(f"  预期结果: 文件夹应打包为 zip 下载")
+        print()
+        
+        print("  💡 请在浏览器中打开以下链接测试下载:")
+        print(f"     {download_url}")
+        print()
+        
+    except Exception as e:
+        print(f"  ❌ 错误: {e}")
+        print()
+    
+    # 4. 测试说明
+    print("=" * 60)
+    print("测试说明:")
+    print("1. 确保 Agent UI 服务正在运行 (端口 8001)")
+    print("2. 在浏览器中打开文件浏览器")
+    print("3. 测试以下功能:")
+    print("   - 点击文件旁的下载图标下载单个文件")
+    print("   - 点击文件夹旁的下载图标下载整个文件夹（zip格式）")
+    print("   - 在文件预览界面点击下载按钮")
+    print("4. 验证下载的文件内容是否正确")
+    print("=" * 60)
+    
+    # 清理测试文件（可选）
+    # import shutil
+    # shutil.rmtree(test_dir)
+    # print(f"\n✅ 测试文件已清理")
+
+if __name__ == "__main__":
+    test_download_api()

dp/agent/cli/templates/ui/ui_utils.py

@@ -0,0 +1,260 @@
+import os
+import json
+import subprocess
+import signal
+import sys
+import time
+from pathlib import Path
+from typing import Dict, Any, Optional, List
+import click
+
+
+class UIConfigManager:
+    """管理 UI 配置的工具类"""
+    
+    DEFAULT_CONFIG = {
+        "agent": {
+            "module": "agent",  # 用户必须提供具体的模块路径
+            "rootAgent": "root_agent",
+            "name": "DP Agent Assistant",
+            "description": "AI Assistant powered by DP Agent SDK",
+            "welcomeMessage": "欢迎使用 DP Agent Assistant！我可以帮助您进行科学计算、数据分析等任务。",
+        },
+        "ui": {
+            "title": "DP Agent Assistant"
+        },
+        "server": {
+            "port": int(os.environ.get('AGENT_SERVER_PORT', '50002')),
+            "host": ["*"]  # 默认允许所有主机访问
+        }
+    }
+    
+    def __init__(self, config_path: Optional[str] = None):
+        self.config_path = Path(config_path) if config_path else Path.cwd() / "agent-config.json"
+        self.config = self._load_config()
+    
+    def _load_config(self) -> Dict[str, Any]:
+        """加载配置文件，如果不存在则使用默认配置"""
+        if self.config_path.exists():
+            with open(self.config_path, 'r') as f:
+                user_config = json.load(f)
+                # 深度合并用户配置和默认配置
+                return self._deep_merge(self.DEFAULT_CONFIG.copy(), user_config)
+        return self.DEFAULT_CONFIG.copy()
+    
+    def _deep_merge(self, base: Dict, update: Dict) -> Dict:
+        """深度合并两个字典"""
+        for key, value in update.items():
+            if key in base and isinstance(base[key], dict) and isinstance(value, dict):
+                base[key] = self._deep_merge(base[key], value)
+            else:
+                base[key] = value
+        return base
+    
+    def save_config(self, config_path: Optional[Path] = None):
+        """保存配置到文件"""
+        save_path = config_path or self.config_path
+        with open(save_path, 'w') as f:
+            json.dump(self.config, f, indent=2)
+    
+    def update_from_cli(self, **kwargs):
+        """从命令行参数更新配置"""
+        if kwargs.get('agent'):
+            module, _, variable = kwargs['agent'].partition(':')
+            self.config['agent']['module'] = module
+            if variable:
+                self.config['agent']['rootAgent'] = variable
+        
+        if kwargs.get('port'):
+            self.config['server']['port'] = kwargs['port']
+        
+
+
+
+class UIProcessManager:
+    """管理 UI 相关进程的工具类"""
+    
+    def __init__(self, ui_dir: Path, config: Dict[str, Any]):
+        self.ui_dir = ui_dir
+        self.config = config
+        self.processes: List[subprocess.Popen] = []
+        self._setup_signal_handlers()
+    
+    def _setup_signal_handlers(self):
+        """设置信号处理器以优雅关闭进程"""
+        signal.signal(signal.SIGINT, self._signal_handler)
+        signal.signal(signal.SIGTERM, self._signal_handler)
+    
+    def _signal_handler(self, signum, frame):
+        """处理终止信号"""
+        # Prevent multiple signal handling
+        signal.signal(signal.SIGINT, signal.SIG_IGN)
+        signal.signal(signal.SIGTERM, signal.SIG_IGN)
+        self.cleanup()
+        # Don't exit here, let the main process handle it
+    
+    def start_websocket_server(self):
+        """启动 WebSocket 服务器"""
+        # 统一使用 server.port
+        server_port = self.config.get('server', {}).get('port', int(os.environ.get('AGENT_SERVER_PORT', '50002')))
+        
+        websocket_script = self.ui_dir / "websocket-server.py"
+        if not websocket_script.exists():
+            raise FileNotFoundError(f"找不到 websocket-server.py: {websocket_script}")
+        
+        # 设置环境变量
+        env = os.environ.copy()
+        env['AGENT_CONFIG_PATH'] = str(self.ui_dir / "config" / "agent-config.temp.json")
+        env['USER_WORKING_DIR'] = str(Path.cwd())  # 传递用户工作目录
+        env['UI_TEMPLATE_DIR'] = str(self.ui_dir)  # 传递UI模板目录
+        # Ensure PYTHONPATH includes the user's working directory
+        if 'PYTHONPATH' in env:
+            env['PYTHONPATH'] = f"{str(Path.cwd())}:{env['PYTHONPATH']}"
+        else:
+            env['PYTHONPATH'] = str(Path.cwd())
+        
+        # 静默启动，将输出重定向到日志文件
+        # 使用 'w' 模式清空旧日志
+        log_file = open(Path.cwd() / "websocket.log", "w")
+        process = subprocess.Popen(
+            [sys.executable, str(websocket_script)],
+            cwd=str(Path.cwd()),  # 在用户工作目录运行
+            env=env,
+            stdout=log_file,
+            stderr=subprocess.STDOUT
+        )
+        self.processes.append(process)
+        
+        # 等待服务器启动
+        time.sleep(2)
+        
+        if process.poll() is not None:
+            raise RuntimeError("WebSocket 服务器启动失败")
+        
+        click.echo(f"🚀 WebSocket 服务器已启动（端口 {server_port}）")
+        click.echo("📝 查看日志: websocket.log")
+        
+        return process
+    
+    def start_frontend_server(self, dev_mode: bool = True):
+        """启动前端服务器"""
+        server_port = self.config.get('server', {}).get('port', int(os.environ.get('AGENT_SERVER_PORT', '50002')))
+        
+        ui_path = self.ui_dir / "frontend"
+        if not ui_path.exists():
+            raise FileNotFoundError(f"找不到 UI 目录: {ui_path}")
+        
+        # 检查是否有构建好的静态文件
+        dist_path = ui_path / "ui-static"
+        if dist_path.exists() and not dev_mode:
+            # 生产模式：静态文件由 WebSocket 服务器提供
+            click.echo(f"✨ Agent UI 已启动: http://{os.environ.get('AGENT_HOST', 'localhost')}:{server_port}")
+            return None
+        
+        if not dev_mode and not dist_path.exists():
+            click.echo("警告: 未找到构建的静态文件，将使用开发模式")
+            dev_mode = True
+        
+        # 检查是否已安装依赖
+        node_modules = ui_path / "node_modules"
+        if not node_modules.exists():
+            click.echo("检测到未安装前端依赖，正在安装...")
+            subprocess.run(["npm", "install"], cwd=str(ui_path), check=True)
+        
+        # 设置环境变量
+        env = os.environ.copy()
+        # 开发模式下，前端开发服务器端口
+        frontend_dev_port = int(os.environ.get('FRONTEND_DEV_PORT', '3000'))
+        env['FRONTEND_PORT'] = str(frontend_dev_port)
+        # 告诉前端后端服务器在哪个端口
+        env['VITE_WS_PORT'] = str(server_port)
+        
+        # 启动命令
+        if dev_mode:
+            cmd = ["npm", "run", "dev"]
+            click.echo(f"启动前端开发服务器...")
+        else:
+            cmd = ["npm", "run", "build"]
+            click.echo("构建前端生产版本...")
+        
+        # 启动前端
+        log_file_path = Path.cwd() / "frontend.log"
+        with open(log_file_path, "a") as log_file:
+            process = subprocess.Popen(
+                cmd,
+                cwd=str(ui_path),
+                env=env,
+                stdout=log_file,
+                stderr=subprocess.STDOUT
+            )
+            self.processes.append(process)
+            
+            # 等待服务器启动
+            time.sleep(3)
+            
+            # 检查进程状态
+            if process.poll() is not None:
+                # 读取错误日志
+                with open(log_file_path, "r") as f:
+                    error_log = f.read()
+                    if "EADDRINUSE" in error_log:
+                        raise RuntimeError(f"端口 {frontend_dev_port} 已被占用")
+                    else:
+                        raise RuntimeError(f"前端服务器启动失败，请查看 frontend.log 了解详情")
+        
+        if dev_mode:
+            click.echo(f"\n✨ 前端开发服务器: http://localhost:{frontend_dev_port}")
+            click.echo(f"📡 后端服务器: http://localhost:{server_port}\n")
+        
+        return process
+    
+    
+    def wait_for_processes(self):
+        """等待所有进程结束"""
+        try:
+            for process in self.processes:
+                if process:  # 处理可能的 None
+                    process.wait()
+        except KeyboardInterrupt:
+            # Don't handle here, let it bubble up to main
+            raise
+    
+    def cleanup(self):
+        """清理所有进程"""
+        if not self.processes:
+            return
+        
+        click.echo("\n🛑 正在停止所有进程...")
+        
+        # First attempt to terminate all processes gracefully
+        for process in self.processes:
+            if process and process.poll() is None:
+                try:
+                    process.terminate()
+                except:
+                    pass
+        
+        # Give processes time to terminate gracefully
+        time.sleep(1)
+        
+        # Force kill any remaining processes
+        for process in self.processes:
+            if process and process.poll() is None:
+                try:
+                    if sys.platform == "win32":
+                        # Windows specific kill
+                        subprocess.run(["taskkill", "/F", "/PID", str(process.pid)], capture_output=True)
+                    else:
+                        # Unix-like systems
+                        process.kill()
+                    process.wait(timeout=1)
+                except:
+                    pass
+        
+        self.processes.clear()
+        
+        # 等待端口释放
+        time.sleep(0.5)
+        click.echo("✅ 所有进程已停止")
+
+

dp/agent/cli/templates/ui/websocket-server.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""
+Agent WebSocket 服务器 - 主入口文件
+使用 Session 运行 rootagent，并通过 WebSocket 与前端通信
+"""
+
+import os
+import sys
+import warnings
+
+# 忽略 paramiko 的加密算法弃用警告
+warnings.filterwarnings("ignore", category=DeprecationWarning, module="paramiko")
+
+# Add user working directory to Python path first
+user_working_dir = os.environ.get('USER_WORKING_DIR')
+if user_working_dir and user_working_dir not in sys.path:
+    sys.path.insert(0, user_working_dir)
+
+# Add UI template directory to Python path for config imports
+ui_template_dir = os.environ.get('UI_TEMPLATE_DIR')
+if ui_template_dir and ui_template_dir not in sys.path:
+    sys.path.insert(0, ui_template_dir)
+
+import uvicorn
+from server.app import create_app
+from server.utils import check_port_available
+from config.agent_config import agentconfig
+
+
+if __name__ == "__main__":
+    print("🚀 启动 Agent WebSocket 服务器...")
+    
+    # 统一使用 server 配置
+    server_config = agentconfig.config.get('server', {})
+    port = server_config.get('port', 8000)
+    # host 数组中的第一个作为显示用
+    hosts = server_config.get('host', ['localhost'])
+    display_host = hosts[0] if isinstance(hosts, list) else hosts
+    
+    
+    # 创建应用
+    app = create_app()
+    
+    print("📡 使用 Session 模式运行 rootagent")
+    print(f"🌐 服务器地址: http://{display_host}:{port}")
+    print(f"🔌 WebSocket 端点: ws://{display_host}:{port}/ws")
+    print("🛑 使用 Ctrl+C 优雅关闭服务器")
+    
+    # uvicorn 始终监听 0.0.0.0 以支持所有配置的主机
+    uvicorn.run(
+        app, 
+        host="0.0.0.0", 
+        port=port,
+        log_level="info",  # 使用 info 级别，过滤掉 warning
+        access_log=False,  # 禁用访问日志，减少噪音
+        # 添加自定义的日志配置
+        log_config={
+            "version": 1,
+            "disable_existing_loggers": False,
+            "formatters": {
+                "default": {
+                    "fmt": "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+                    "datefmt": "%Y-%m-%d %H:%M:%S"
+                }
+            },
+            "handlers": {
+                "default": {
+                    "formatter": "default",
+                    "class": "logging.StreamHandler",
+                    "stream": "ext://sys.stdout"
+                }
+            },
+            "root": {
+                "level": "INFO",
+                "handlers": ["default"]
+            },
+            "loggers": {
+                "uvicorn.error": {
+                    "level": "ERROR"
+                },
+                "uvicorn.access": {
+                    "handlers": [],
+                    "propagate": False
+                }
+            }
+        }
+    )

dp/agent/server/storage/http_storage.py

@@ -51,4 +51,4 @@ class HTTPStorage(BaseStorage):
 
 
 class HTTPSStorage(HTTPStorage):
-    scheme = "https"
+    scheme = "https"