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

dp/agent/cli/templates/device/tescan_device.py.template

@@ -0,0 +1,158 @@
+"""
+Tescan device implementation.
+
+This module contains the implementation of a Tescan device with
+specific parameter and result types.
+"""
+import time
+from typing import Dict, TypedDict
+from dp.agent.device.device.types import SuccessResult, BaseParams, ErrorResult
+from dp.agent.device.device.device import Device, action
+
+class TakePictureParams(BaseParams):
+    """Parameters for take_picture action."""
+    horizontal_width: str
+
+class GetStagePositionParams(BaseParams):
+    """Parameters for get_stage_position action."""
+    pass
+
+class MoveStageParams(BaseParams):
+    """Parameters for move_stage action."""
+    x: float
+    y: float
+    z: float
+
+class StagePosition(TypedDict):
+    """Stage position data."""
+    x: float
+    y: float
+    z: float
+
+class PictureData(TypedDict):
+    """Picture data."""
+    image_id: str
+
+
+class SleepParams(BaseParams):
+    """Parameters for sleep action."""
+    seconds: int
+
+class StagePositionResult(SuccessResult):
+    """Result for get_stage_position action."""
+    data: StagePosition
+
+class PictureResult(SuccessResult):
+    """Result for take_picture action."""
+    data: PictureData
+
+class MoveStageResult(SuccessResult):
+    """Result for move_stage action."""
+    data: Dict[str, float]
+
+class TescanDevice(Device):
+    """Tescan device implementation."""
+    device_name = "tescan"
+
+    @action("start_electron_beam")
+    def start_electron_beam(self, params: BaseParams):
+        print("start electron beam")
+    
+    @action("start_ion_beam")
+    def start_ion_beam(self, params: BaseParams):
+        print("start ion beam")
+    
+    @action("toggle_electron_scanning")
+    def toggle_electron_scanning(self, params: BaseParams):
+        """start scanning, this will cost 5 seconds to finish. call sleep(5) after toggle this.
+        """
+        print("toggle electron scanning")
+    
+    @action("toggle_ion_scanning")
+    def toggle_ion_scanning(self, params: BaseParams):
+        print("toggle ion scanning")
+
+    @action("sleep")
+    def sleep(self, params: SleepParams):
+        time.sleep(params.get("seconds", 0))
+
+    @action("take_picture")
+    def take_picture(self, params: TakePictureParams) -> PictureResult:
+        """Take a picture with the microscope.
+        
+        Args:
+            params: Parameters for the action
+                - horizontal_width: Horizontal width of the image
+                
+        Returns:
+            Result of the action
+        """
+        time.sleep(1)
+        
+        hw = params.get("horizontal_width", "default")
+        
+        return PictureResult(
+            message=f"Picture taken with {self.device_name} (width: {hw})",
+            data={"image_id": "mock_image_123"}
+        )
+    
+    @action("get_stage_position")
+    def get_stage_position(self, params: GetStagePositionParams) -> StagePositionResult:
+        """Get the stage position.
+        
+        Args:
+            params: Parameters for the action (none required)
+                
+        Returns:
+            Result of the action with x, y, z coordinates
+        """
+        time.sleep(0.5)
+        
+        return StagePositionResult(
+            message=f"Stage position retrieved for {self.device_name}",
+            data=StagePosition(
+                x=10.5,
+                y=20.3,
+                z=5.1
+            )
+        )
+    
+    @action("move_stage")
+    def move_stage(self, params: MoveStageParams) -> MoveStageResult:
+        """Move the stage to a new position.
+        
+        Args:
+            params: Parameters for the action
+                - x: X coordinate
+                - y: Y coordinate
+                - z: Z coordinate
+                
+        Returns:
+            Result of the action
+        """
+        time.sleep(2)
+        
+        if not all(k in params for k in ["x", "y", "z"]):
+            # Use a proper MoveStageResult with error status instead of ErrorResult
+            return MoveStageResult(
+                message=f"Missing required parameters: x, y, z",
+                data={}
+            )
+        
+        try:
+            new_position = {
+                "x": float(params["x"]),
+                "y": float(params["y"]),
+                "z": float(params["z"])
+            }
+            
+            return MoveStageResult(
+                message=f"Stage moved for {self.device_name}",
+                data=new_position
+            )
+        except (KeyError, TypeError, ValueError):
+            # Use a proper MoveStageResult with error status instead of ErrorResult
+            return MoveStageResult(
+                message=f"Invalid parameters: x, y, z must be numeric values",
+                data={}
+            )

dp/agent/cli/templates/main.py.template

@@ -0,0 +1,67 @@
+import sys
+import signal
+from pathlib import Path
+import os
+
+# 从SDK包中导入所需模块
+from dp.agent.cloud.mcp import mcp
+from dp.agent.cloud.mqtt import get_mqtt_cloud_instance
+from dp.agent.device.mqtt_device_twin import DeviceTwin
+# 从本地导入 TescanDevice
+
+def run_cloud():
+    """Run in cloud environment"""
+    mqtt_instance = get_mqtt_cloud_instance()
+    from dp.agent.device.device.device import register_mcp_tools
+    from device.tescan_device import TescanDevice
+    register_mcp_tools(mcp, TescanDevice())
+    def signal_handler(sig, frame):
+        """Handle SIGINT signal to gracefully shutdown."""
+        print("Shutting down...")
+        mqtt_instance.stop()
+        sys.exit(0)
+    
+    # Register signal handler for graceful shutdown
+    signal.signal(signal.SIGINT, signal_handler)
+    
+    # Start MCP server
+    print("Starting MCP server...")
+    mcp.run(transport="sse")
+def run_device():
+    """Run in lab environment"""
+    # Initialize device and twin
+    from device.tescan_device import TescanDevice
+    tescan_device = TescanDevice()
+    device_twin = DeviceTwin(tescan_device)
+    
+    # Run device twin
+    try:
+        device_twin.run()
+    except KeyboardInterrupt:
+        print("\nShutting down lab environment...")
+        sys.exit(0)
+
+def run_calculation():
+    """Run in calculation environment"""
+    # 在此实现你自己的计算环境逻辑
+    print("Running calculation task...")
+    # 模拟任务处理逻辑
+    from calculation import simple
+    simple.run()
+    
+if __name__ == "__main__":
+    if len(sys.argv) != 2:
+        print("Usage: python main.py [device|cloud|calculation]")
+        sys.exit(1)
+        
+    mode = sys.argv[1]
+    if mode == "device":
+        run_device()
+    elif mode == "cloud":
+        run_cloud()
+    elif mode == "calculation":  
+        run_calculation()
+    else:
+        print(f"Unknown mode: {mode}")
+        print("Usage: python main.py [device|cloud|calculation]")
+        sys.exit(1) 

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

@@ -0,0 +1 @@
+# UI 模板包

dp/agent/cli/templates/ui/api/__init__.py

@@ -0,0 +1 @@
+# API !W

dp/agent/cli/templates/ui/api/config.py

@@ -0,0 +1,32 @@
+# Configuration API
+from fastapi import Request
+from fastapi.responses import JSONResponse
+
+from server.utils import get_ak_info_from_request
+from config.agent_config import agentconfig
+
+
+async def get_config(request: Request):
+    """Get frontend configuration"""
+    access_key, _ = get_ak_info_from_request(request.headers)
+    return JSONResponse(content={
+        "agent": agentconfig.config.get("agent", {}),
+        "ui": agentconfig.get_ui_config(),
+        "files": agentconfig.get_files_config(),
+        "websocket": agentconfig.get_websocket_config(),
+        "user_type": "registered" if access_key else "temporary"
+    })
+
+
+async def status():
+    """API status"""
+    return {
+        "message": f"{agentconfig.config.get('agent', {}).get('name', 'Agent')} WebSocket 服务器正在运行",
+        "mode": "session",
+        "endpoints": {
+            "websocket": "/ws",
+            "files": "/api/files",
+            "file_tree": "/api/files/tree",
+            "config": "/api/config"
+        }
+    }

dp/agent/cli/templates/ui/api/constants.py

@@ -0,0 +1,61 @@
+"""
+Constants for API modules
+"""
+
+# File upload limits
+MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
+MAX_FILES_PER_UPLOAD = 10
+
+# Allowed file extensions for upload
+ALLOWED_EXTENSIONS = {
+    'txt', 'pdf', 'csv', 'json', 'xml', 
+    'png', 'jpg', 'jpeg', 'gif', 'svg', 
+    'py', 'js', 'ts', 'java', 'cpp', 'c',
+    'md', 'rst', 'yaml', 'yml', 'log',
+    'html', 'htm', 'css', 'scss',
+    'sh', 'bash', 'sql', 'toml'
+}
+
+# WebSocket message types
+WS_MESSAGE_TYPES = {
+    'MESSAGE': 'message',
+    'CREATE_SESSION': 'create_session',
+    'SWITCH_SESSION': 'switch_session',
+    'GET_SESSIONS': 'get_sessions',
+    'DELETE_SESSION': 'delete_session',
+    'SET_PROJECT_ID': 'set_project_id',
+    'ERROR': 'error',
+    'PROJECT_ID_SET': 'project_id_set',
+    'SESSION_MESSAGES': 'session_messages'
+}
+
+# Session states
+SESSION_STATES = {
+    'ACTIVE': 'active',
+    'INACTIVE': 'inactive',
+    'DELETED': 'deleted'
+}
+
+# API endpoints
+API_ENDPOINTS = {
+    'WEBSOCKET': '/ws',
+    'FILES': '/api/files',
+    'FILE_TREE': '/api/files/tree',
+    'CONFIG': '/api/config',
+    'PROJECTS': '/api/projects',
+    'SESSIONS': '/api/sessions',
+    'DEBUG': '/api/debug'
+}
+
+# File operations
+FILE_OPS = {
+    'DOWNLOAD': 'download',
+    'UPLOAD': 'upload',
+    'DELETE': 'delete',
+    'LIST': 'list'
+}
+
+# Default values
+DEFAULT_WORKSPACE_NAME = "Workspace"
+DEFAULT_SESSION_LIMIT = 50
+DEFAULT_MESSAGE_LIMIT = 100

dp/agent/cli/templates/ui/api/debug.py

@@ -0,0 +1,257 @@
+"""
+Debug API endpoints for troubleshooting
+"""
+import os
+from fastapi import HTTPException
+from fastapi.responses import JSONResponse
+from config.agent_config import agentconfig
+from api.websocket import manager
+
+async def get_runner_status():
+    """
+    Get status of all runners
+    """
+    try:
+        runner_status = {}
+        
+        # Get status of all runners
+        for runner_key, runner in manager.runners.items():
+            runner_status[runner_key] = {
+                "exists": True,
+                "type": type(runner).__name__,
+                "created": True
+            }
+        
+        # Get all runner errors
+        runner_errors = getattr(manager, '_runner_errors', {})
+        for error_key, error_msg in runner_errors.items():
+            runner_status[error_key] = {
+                "exists": False,
+                "error": error_msg,
+                "created": False
+            }
+        
+        return JSONResponse({
+            "success": True,
+            "runners": runner_status,
+            "total_runners": len(manager.runners),
+            "total_errors": len(runner_errors),
+            "active_connections": len(manager.active_connections)
+        })
+    except Exception as e:
+        return JSONResponse({
+            "success": False,
+            "error": str(e)
+        }, status_code=500)
+
+async def get_config_status():
+    """
+    Get configuration status
+    """
+    try:
+        config = agentconfig.config
+        
+        # Get agent configuration
+        agent_config = config.get("agent", {})
+        
+        # Test if agent can be loaded
+        test_result = {
+            "can_load": False,
+            "error": None,
+            "agent_type": None
+        }
+        
+        try:
+            # Try to load agent (without parameters)
+            test_agent = agentconfig.get_agent()
+            test_result["can_load"] = True
+            test_result["agent_type"] = type(test_agent).__name__
+        except Exception as e:
+            test_result["error"] = str(e)
+        
+        return JSONResponse({
+            "success": True,
+            "config": {
+                "agent_name": agent_config.get("name"),
+                "agent_module": agent_config.get("module"),
+                "root_agent": agent_config.get("rootAgent"),
+                "config_path": str(agentconfig.config_path),
+                "config_exists": agentconfig.config_path.exists()
+            },
+            "test_result": test_result,
+            "environment": {
+                "user_working_dir": os.environ.get('USER_WORKING_DIR', os.getcwd()),
+                "bohr_project_id": os.environ.get('BOHR_PROJECT_ID'),
+                "agent_config_path": os.environ.get('AGENT_CONFIG_PATH')
+            }
+        })
+    except Exception as e:
+        return JSONResponse({
+            "success": False,
+            "error": str(e)
+        }, status_code=500)
+
+async def get_session_status():
+    """
+    Get session status
+    """
+    try:
+        session_info = {}
+        
+        # Get session information for each user
+        for user_id, session_service in manager.session_services.items():
+            try:
+                response = await session_service.list_sessions(
+                    app_name=manager.app_name,
+                    user_id=user_id
+                )
+                sessions = response.sessions if hasattr(response, 'sessions') else []
+                session_info[user_id] = {
+                    "session_count": len(sessions),
+                    "sessions": [
+                        {
+                            "id": s.id,
+                            "state": s.state if hasattr(s, 'state') else {}
+                        } for s in sessions[:3]  # Only show first 3
+                    ]
+                }
+            except Exception as e:
+                session_info[user_id] = {
+                    "error": str(e)
+                }
+        
+        return JSONResponse({
+            "success": True,
+            "users": len(manager.session_services),
+            "session_info": session_info,
+            "active_connections": [
+                {
+                    "user": ctx.get_user_identifier(),
+                    "project_id": ctx.project_id,
+                    "current_session": ctx.current_session_id,
+                    "is_registered": ctx.is_registered_user()
+                }
+                for ctx in manager.active_connections.values()
+            ]
+        })
+    except Exception as e:
+        return JSONResponse({
+            "success": False,
+            "error": str(e)
+        }, status_code=500)
+
+async def test_agent_creation():
+    """
+    Test agent creation process
+    """
+    try:
+        steps = []
+        
+        # Step 1: Load configuration
+        steps.append({
+            "step": "Load configuration",
+            "success": True,
+            "details": {
+                "config_path": str(agentconfig.config_path),
+                "agent_module": agentconfig.config.get("agent", {}).get("module")
+            }
+        })
+        
+        # Step 2: Create agent (without parameters)
+        try:
+            agent = agentconfig.get_agent()
+            steps.append({
+                "step": "Create agent (no parameters)",
+                "success": True,
+                "details": {
+                    "agent_type": type(agent).__name__,
+                    "has_run_method": hasattr(agent, 'run') or hasattr(agent, '__call__')
+                }
+            })
+        except Exception as e:
+            steps.append({
+                "step": "Create agent (no parameters)",
+                "success": False,
+                "error": str(e)
+            })
+            
+        # Step 3: Create agent (with project_id)
+        try:
+            project_id = os.environ.get('BOHR_PROJECT_ID')
+            if project_id:
+                agent_with_project = agentconfig.get_agent(
+                    ak="",
+                    app_key="",
+                    project_id=int(project_id)
+                )
+                steps.append({
+                    "step": f"Create agent (project_id={project_id})",
+                    "success": True,
+                    "details": {
+                        "agent_type": type(agent_with_project).__name__
+                    }
+                })
+        except Exception as e:
+            steps.append({
+                "step": f"创建 Agent (project_id={project_id})",
+                "success": False,
+                "error": str(e)
+            })
+        
+        # Step 4: Check SessionService
+        try:
+            from google.adk.sessions import InMemorySessionService
+            test_service = InMemorySessionService()
+            steps.append({
+                "step": "Create SessionService",
+                "success": True,
+                "details": {
+                    "service_type": type(test_service).__name__
+                }
+            })
+        except Exception as e:
+            steps.append({
+                "step": "Create SessionService",
+                "success": False,
+                "error": str(e)
+            })
+        
+        # Step 5: Create Runner
+        try:
+            from google.adk import Runner
+            if 'agent' in locals() and 'test_service' in locals():
+                test_runner = Runner(
+                    agent=agent,
+                    session_service=test_service,
+                    app_name="TestApp"
+                )
+                steps.append({
+                    "step": "Create Runner",
+                    "success": True,
+                    "details": {
+                        "runner_type": type(test_runner).__name__
+                    }
+                })
+        except Exception as e:
+            steps.append({
+                "step": "创建 Runner", 
+                "success": False,
+                "error": str(e)
+            })
+        
+        all_success = all(step.get("success", False) for step in steps)
+        
+        return JSONResponse({
+            "success": all_success,
+            "steps": steps,
+            "summary": {
+                "total_steps": len(steps),
+                "successful": sum(1 for s in steps if s.get("success", False)),
+                "failed": sum(1 for s in steps if not s.get("success", False))
+            }
+        })
+    except Exception as e:
+        return JSONResponse({
+            "success": False,
+            "error": str(e)
+        }, status_code=500)