bohr-agent-sdk 0.1.106__py3-none-any.whl → 0.1.109__py3-none-any.whl

{bohr_agent_sdk-0.1.106.dist-info → bohr_agent_sdk-0.1.109.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.106
+Version: 0.1.109
 Summary: SDK for scientific agents
 Home-page: https://github.com/dptech-corp/bohr-agent-sdk/
 Author: DP Technology
@@ -259,6 +259,33 @@ BOHRIUM_PASSWORD=your_password
 
 Note: The `dp-agent fetch config` command automatically downloads configuration files and replaces dynamic variables (such as MQTT_DEVICE_ID). For security reasons, this feature is only available in internal network environments.
 
+## 🔒 Authentication Configuration
+
+For private deployments or development environment debugging, you need to configure the following environment variables:
+
+- `BOHR_ACCESS_KEY`: Requires a real Access Key obtained from [Bohrium User Settings](https://www.bohrium.com/settings/user)
+- `BOHR_APP_KEY`: Can be set to any value for development
+
+### Linux/macOS:
+```bash
+export BOHR_ACCESS_KEY=your_real_ak_from_bohrium_settings
+export BOHR_APP_KEY=any_value_for_dev
+```
+
+### Windows (Command Prompt):
+```cmd
+set BOHR_ACCESS_KEY=your_real_ak_from_bohrium_settings
+set BOHR_APP_KEY=any_value_for_dev
+```
+
+### Windows (PowerShell):
+```powershell
+$env:BOHR_ACCESS_KEY="your_real_ak_from_bohrium_settings"
+$env:BOHR_APP_KEY="any_value_for_dev"
+```
+
+For agents deployed on Bohrium APP, authentication parameters will be automatically obtained from cookies.
+
 ## 🎯 Application Scenarios
 
 - **Materials Science Computing**: Molecular dynamics simulation, first-principles calculations

{bohr_agent_sdk-0.1.106.dist-info → bohr_agent_sdk-0.1.109.dist-info}/RECORD

@@ -22,13 +22,13 @@ dp/agent/cli/templates/ui/api/__init__.py,sha256=eq1djuBJYvWlGTgxXvT9A386Eu21FNC
 dp/agent/cli/templates/ui/api/config.py,sha256=ZhHApcXDrmO1JjI-Ok09ETYB-BVis390ZrMRt140DhQ,1018
 dp/agent/cli/templates/ui/api/constants.py,sha256=mHY81zf4ksfKEOsSiVsWD_A_KcQ6wveLvAImCay_kVM,1391
 dp/agent/cli/templates/ui/api/debug.py,sha256=URC1k429s84Ue7ElcK7jFh73QeK04IlRyXEjVgLXnbY,8490
-dp/agent/cli/templates/ui/api/files.py,sha256=DDqU8-u5yQmHz2sd9vuEURpdgTkQVlBjAyv7dhWz-mA,16469
+dp/agent/cli/templates/ui/api/files.py,sha256=9jQNUOdCVeCCd0PfB95NqKXQX_vHzcHIIZBkAMEsVwo,16026
 dp/agent/cli/templates/ui/api/files_upload.py,sha256=SOPy7GdQQG4grj5DvuJt3p1_HHJYdk6EuehqtKMAcqs,4411
 dp/agent/cli/templates/ui/api/files_user.py,sha256=mEFCpZChHzxNUxYUM7A3VqyMbgM3cR0M9FdT93wXJR8,1718
 dp/agent/cli/templates/ui/api/messages.py,sha256=umIkjfVA5rlXIonWQP7JRLqpfjDdvOc046gBFidMk5A,4992
 dp/agent/cli/templates/ui/api/projects.py,sha256=KOMRAPrfsbBKNBPotKHEFIKyBypQpox4lSO3ocnjA_M,5008
 dp/agent/cli/templates/ui/api/sessions.py,sha256=Mi_Sq4mfW163oZ2K0lRFiSgE7MpvVArmmSItDL2fBFo,2994
-dp/agent/cli/templates/ui/api/utils.py,sha256=6oTf1P105GEGcwDct-TB3v2pt2iQezTmywfniwrlOqM,4728
+dp/agent/cli/templates/ui/api/utils.py,sha256=FLbFJpBQ5O9F4BoLUWv5PvGQG8IAOMWI6QM7albVKTo,4825
 dp/agent/cli/templates/ui/api/websocket.py,sha256=kPc9ZArxHSpvEZHl8So5vAH7rWaKNcfTk2O5_61wYUg,9013
 dp/agent/cli/templates/ui/config/__init__.py,sha256=F7JR_fnXNnZDHeqRCdEiWtVMy03_qFr8cOilIQ6XLUk,14
 dp/agent/cli/templates/ui/config/agent_config.py,sha256=WFpukdBuZ0t-QQKnhE7NQ6tYPl98FP6fmweRpNbOZQM,11255
@@ -43,13 +43,15 @@ dp/agent/cli/templates/ui/frontend/ui-static/assets/index-DfN2raU9.css,sha256=Ks
 dp/agent/cli/templates/ui/scripts/build_ui.py,sha256=zoK6agl1dt-InaYQyXSABxjtJj24gtCbm0KKVv0RTSA,1661
 dp/agent/cli/templates/ui/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dp/agent/cli/templates/ui/server/app.py,sha256=ZX7ZA_qJyA-Ur7BODICzSn_S4FwGnRWcQP_qq2MZJhw,3803
-dp/agent/cli/templates/ui/server/connection.py,sha256=huet_WWxGjo_J9ozj-G8tEYSiEYTsgxOTgjEEJ6CubU,6457
+dp/agent/cli/templates/ui/server/connection.py,sha256=DVlhx49gMjKVPxOpi7spAZiHZA5illXkyC3MhczXsQk,6614
 dp/agent/cli/templates/ui/server/file_watcher.py,sha256=t0e7IMgf4E1y0CxeboCqsVeZPBo1CGSw_qoO8wG6Wy4,2923
 dp/agent/cli/templates/ui/server/middleware.py,sha256=bNR5aFqKjxv2einQ0e26aq6tiUW4SMses3IHkrmHDCk,1474
 dp/agent/cli/templates/ui/server/models.py,sha256=obfMkFFB9qZZXDdwoC25u3tBlUC6Rz7sjeyuh7a_UWk,1680
 dp/agent/cli/templates/ui/server/session_manager.py,sha256=ZbNHGCFvswa-LKWn6c6RMHWDdfeOvY5L6A8CfrsF0OE,46536
 dp/agent/cli/templates/ui/server/user_files.py,sha256=khkiyY2UOOysHqO6JgCPUDqtrInp83G1M62i3Lj-0aY,2995
 dp/agent/cli/templates/ui/server/utils.py,sha256=f4NfwFBq_RdZyFn_KCW6ZThYW8TvQyVruK7PJZ-DA80,1530
+dp/agent/client/__init__.py,sha256=yu7HYZwAkD7g5dL9JttLkGmspBcyOf-6OoCjci4oPDA,59
+dp/agent/client/mcp_client.py,sha256=xp8GSM5effjjp34yGp5-HwldalyqT-CWb6PxlqkclCA,6313
 dp/agent/cloud/__init__.py,sha256=e16ymCZX2f-S8DyGB5jSK8gnQqVObRIsvtLXLALIKxQ,441
 dp/agent/cloud/main.py,sha256=5QIEjpZ1RxWnR8wyLf-vlgz1bn9oOnxCYn158LBaLN4,727
 dp/agent/cloud/mcp.py,sha256=tsAwC3doVMLYr6Oh8PxVqF-qCygYkDZJTIhoF_h8eGQ,4537
@@ -60,7 +62,7 @@ dp/agent/device/device/__init__.py,sha256=w7_1S16S1vWUq0RGl0GFgjq2vFkc5oNvy8cQTn
 dp/agent/device/device/device.py,sha256=9ZRIJth-4qMO-i-u_b_cO3d6a4eTbTQjPaxFsV_zEkc,9643
 dp/agent/device/device/types.py,sha256=JuxB-hjf1CjjvfBxCLwRAXVFlYS-nPEdiJpBWLFVCzo,1924
 dp/agent/server/__init__.py,sha256=rckaYd8pbYyB4ENEhgjXKeGMXjdnrgcJpdM1gu5u1Wc,508
-dp/agent/server/calculation_mcp_server.py,sha256=iRFOdgTxySMGk7ZaSseNssEp-A7zT5cW1Ym2_MIKnG4,12602
+dp/agent/server/calculation_mcp_server.py,sha256=hsTxuguyqgq_4HSeJYS7VnW-jTvqf9DI0mpvRj0wK3w,13566
 dp/agent/server/preprocessor.py,sha256=XUWu7QOwo_sIDMYS2b1OTrM33EXEVH_73vk-ju1Ok8A,1264
 dp/agent/server/utils.py,sha256=ui3lca9EagcGqmYf8BKLsPARIzXxJ3jgN98yuEO3OSQ,1668
 dp/agent/server/executor/__init__.py,sha256=s95M5qKQk39Yi9qaVJZhk_nfj54quSf7EDghR3OCFUA,248
@@ -73,8 +75,8 @@ dp/agent/server/storage/bohrium_storage.py,sha256=EsKX4dWWvZTn2TEhZv4zsvihfDK0mm
 dp/agent/server/storage/http_storage.py,sha256=KiySq7g9-iJr12XQCKKyJLn8wJoDnSRpQAR5_qPJ1ZU,1471
 dp/agent/server/storage/local_storage.py,sha256=t1wfjByjXew9ws3PuUxWxmZQ0-Wt1a6t4wmj3fW62GI,1352
 dp/agent/server/storage/oss_storage.py,sha256=pgjmi7Gir3Y5wkMDCvU4fvSls15fXT7Ax-h9MYHFPK0,3359
-bohr_agent_sdk-0.1.106.dist-info/METADATA,sha256=gAYz_SBcnkqaV314ppLvyi9O8dekysp24LlvNhUGfdU,10219
-bohr_agent_sdk-0.1.106.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.106.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.106.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.106.dist-info/RECORD,,
+bohr_agent_sdk-0.1.109.dist-info/METADATA,sha256=B9SUVukVXm_P3E8ZWP41ApkenlCwobucnf9DSN_F2Ds,11070
+bohr_agent_sdk-0.1.109.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.109.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.109.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.109.dist-info/RECORD,,

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

@@ -36,22 +36,9 @@ def get_user_identifier(access_key: str = None, app_key: str = None, session_id:
         if cached_user_id:
             return cached_user_id
     
-    # If no cache, call OpenSDK
-    if access_key and app_key:
-        try:
-            # Use OpenSDK to get user info
-            client = OpenSDK(
-                access_key=access_key,
-                app_key=app_key
-            )
-            user_info = client.user.get_info()
-            if user_info and user_info.get('code') == 0:
-                data = user_info.get('data', {})
-                bohrium_user_id = data.get('user_id')
-                if bohrium_user_id:
-                    return bohrium_user_id
-        except Exception as e:
-            pass
+    # If no cache, return access_key directly (consistent with WebSocket logic)
+    if access_key:
+        return access_key
     
     # If has session_id, use it
     if session_id:

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

@@ -141,21 +141,22 @@ def safe_filename(filename: str) -> str:
 async def check_project_id_required(context_manager, user_identifier: str) -> bool:
     """
     Check if project_id is required and set
-    
+
     Args:
         context_manager: WebSocket manager instance
-        user_identifier: User identifier
-        
+        user_identifier: User identifier (can be access_key or bohrium_user_id)
+
     Returns:
         True if project_id is set, False otherwise
     """
     # First check environment variable
     if os.environ.get('BOHR_PROJECT_ID'):
         return True
-        
+
     # Check user's connection context
     for context in context_manager.active_connections.values():
-        if context.get_user_identifier() == user_identifier:
+        # Match by bohrium_user_id or access_key
+        if context.get_user_identifier() == user_identifier or context.access_key == user_identifier:
             return bool(context.project_id)
-            
+
     return False

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

@@ -86,21 +86,25 @@ class ConnectionContext:
     def get_user_identifier(self) -> str:
         """
         Get unique user identifier
-        Prefer bohrium_user_id, fallback to generated user_id
-        
+        Use access_key if available (for registered users), otherwise use generated user_id
+
         Returns:
             User identifier
         """
-        return self.bohrium_user_id if self.bohrium_user_id else self.user_id
+        # 优先使用 access_key，保持与 HTTP API 的一致性
+        if self.access_key:
+            return self.access_key
+        # 临时用户使用生成的 user_id
+        return self.user_id
         
     def is_registered_user(self) -> bool:
         """
         Check if user is registered
-        
+
         Returns:
-            True if has bohrium_user_id
+            True if has access_key (registered Bohrium user)
         """
-        return bool(self.bohrium_user_id)
+        return bool(self.access_key)
         
     def set_project_id(self, project_id: int):
         """

dp/agent/client/__init__.py

@@ -0,0 +1,3 @@
+from .mcp_client import MCPClient
+
+__all__ = ["MCPClient"]

dp/agent/client/mcp_client.py

@@ -0,0 +1,162 @@
+import asyncio
+import json
+import os
+from contextlib import AsyncExitStack
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+from mcp.client.sse import sse_client
+from mcp.client.streamable_http import streamablehttp_client
+
+from ..server.utils import get_logger
+logger = get_logger(__name__)
+
+
+class MCPClient:
+    def __init__(self, server, query_interval=4):
+        self.server = server
+        self.session = None
+        self.exit_stack = AsyncExitStack()
+        self.query_interval = query_interval
+
+    def _is_session_disconnected(self, session: ClientSession) -> bool:
+        """Checks if a session is disconnected or closed.
+
+        Args:
+            session: The ClientSession to check.
+
+        Returns:
+            True if the session is disconnected, False otherwise.
+        """
+        return session._read_stream._closed or session._write_stream._closed
+
+    async def get_session(self):
+        if self.session is not None:
+            # Check if the existing session is still connected
+            if not self._is_session_disconnected(self.session):
+                # Session is still good, return it
+                return self.session
+            else:
+                # Session is disconnected, clean it up
+                logger.info(f'Cleaning up disconnected session: {self.server}')
+            try:
+                await self.exit_stack.aclose()
+            except Exception as e:
+                logger.warning('Error during disconnected session cleanup: %s',
+                               e)
+            finally:
+                self.session = None
+
+        is_python = self.server.endswith('.py')
+        is_js = self.server.endswith('.js')
+        is_sse = self.server.startswith('http') and "/sse" in self.server
+        is_streamablehttp = self.server.startswith('http') \
+            and "/sse" not in self.server
+        if not (is_python or is_js or is_sse or is_streamablehttp):
+            raise ValueError(
+                "Server script must be a .py or .js file or a http link")
+
+        if is_sse:
+            try:
+                logger.info(f"SSE: {self.server}")
+                streams = await self.exit_stack.enter_async_context(sse_client(
+                    self.server))
+                self.session = await self.exit_stack.enter_async_context(
+                    ClientSession(*streams))
+                await self.session.initialize()
+                return self.session
+            except Exception as e:
+                logger.error(
+                    f"An error occurred while connecting to the server: {e}")
+                raise
+        if is_streamablehttp:
+            try:
+                logger.info(f"StreamableHTTP: {self.server}")
+                transports = await self.exit_stack.enter_async_context(
+                    streamablehttp_client(self.server))
+                self.session = await self.exit_stack.enter_async_context(
+                    ClientSession(*transports[:2]))
+                await self.session.initialize()
+                return self.session
+            except Exception as e:
+                logger.error(
+                    f"An error occurred while connecting to the server: {e}")
+                raise
+
+        command = "python" if is_python else "node"
+        server_params = StdioServerParameters(
+            command=command,
+            args=[self.server],
+            env=os.environ.copy(),
+        )
+
+        try:
+            stdio_transport = await self.exit_stack.enter_async_context(
+                stdio_client(server_params))
+            stdio, write = stdio_transport
+            self.session = await self.exit_stack.enter_async_context(
+                ClientSession(stdio, write))
+            await self.session.initialize()
+            return self.session
+        except Exception as e:
+            logger.error(
+                f"An error occurred while connecting to the server: {e}")
+            raise
+
+    async def connect_to_server(self):
+        await self.get_session()
+        response = await self.session.list_tools()
+        tools = []
+        for tool in response.tools:
+            if tool.name.startswith("submit_") or tool.name in [
+                    "query_job_status", "get_job_results", "terminate_job"]:
+                continue
+            tools.append(tool)
+        logger.info(
+            f"Connected to server with tools:{[tool.name for tool in tools]}")
+        return tools
+
+    async def call_tool(self, tool_name: str, arguments: dict,
+                        async_mode=False):
+        await self.get_session()
+        if not async_mode:
+            result = await self.session.call_tool(tool_name, arguments)
+            return result
+
+        executor = arguments.get("executor")
+        storage = arguments.get("storage")
+        res = await self.session.call_tool("submit_" + tool_name, arguments)
+        if res.isError:
+            logger.error("Failed to submit %s: %s" % (
+                tool_name, res.content[0].text))
+            return res
+        job_id = json.loads(res.content[0].text)["job_id"]
+        job_info = res.content[0].job_info
+        logger.info("Job submitted (ID: %s)" % job_id)
+        if job_info.get("extra_info"):
+            logger.info(job_info["extra_info"])
+
+        while True:
+            res = await self.session.call_tool("query_job_status", {
+                "job_id": job_id, "executor": executor})
+            if res.isError:
+                logger.error(res.content[0].text)
+            else:
+                status = res.content[0].text
+                logger.info("Job %s status is %s" % (job_id, status))
+                if status != "Running":
+                    break
+            await asyncio.sleep(self.query_interval)
+
+        res = await self.session.call_tool("get_job_results", {
+            "job_id": job_id, "executor": executor, "storage": storage})
+        if res.isError:
+            logger.error("Job %s failed: %s" % (job_id, res.content[0].text))
+        else:
+            logger.info("Job %s result is %s" % (job_id, res.content))
+        res.content[0].job_info = {**job_info,
+                                   **getattr(res.content[0], "job_info", {})}
+        return res
+
+    async def cleanup(self):
+        await self.exit_stack.aclose()

dp/agent/server/calculation_mcp_server.py

@@ -1,4 +1,5 @@
 import inspect
+import json
 import os
 from collections.abc import Callable
 from contextlib import contextmanager
@@ -64,6 +65,20 @@ def set_directory(workdir: str):
         os.chdir(cwd)
 
 
+def load_executor(executor):
+    if not executor and os.path.exists("executor.json"):
+        with open("executor.json", "r") as f:
+            executor = json.load(f)
+    return executor
+
+
+def load_storage(storage):
+    if not storage and os.path.exists("storage.json"):
+        with open("storage.json", "r") as f:
+            storage = json.load(f)
+    return storage
+
+
 def query_job_status(job_id: str, executor: Optional[dict] = None
                      ) -> Literal["Running", "Succeeded", "Failed"]:
     """
@@ -75,6 +90,7 @@ def query_job_status(job_id: str, executor: Optional[dict] = None
     """
     trace_id, exec_id = job_id.split("/")
     with set_directory(trace_id):
+        executor = load_executor(executor)
         _, executor = init_executor(executor)
         status = executor.query_status(exec_id)
         logger.info("Job %s status is %s" % (job_id, status))
@@ -89,6 +105,7 @@ def terminate_job(job_id: str, executor: Optional[dict] = None):
     """
     trace_id, exec_id = job_id.split("/")
     with set_directory(trace_id):
+        executor = load_executor(executor)
         _, executor = init_executor(executor)
         executor.terminate(exec_id)
         logger.info("Job %s is terminated" % job_id)
@@ -138,17 +155,21 @@ def handle_output_artifacts(results, exec_id, storage):
     return results, output_artifacts
 
 
+# MCP does not regard Any as serializable in Python 3.12
+# use Optional[Any] to work around
 def get_job_results(job_id: str, executor: Optional[dict] = None,
-                    storage: Optional[dict] = None) -> Any:
+                    storage: Optional[dict] = None) -> Optional[Any]:
     """
     Get results of a calculation job
     Args:
         job_id (str): The ID of the calculation job
     Returns:
-        results (dict): results of the calculation job
+        results (Any): results of the calculation job
     """
     trace_id, exec_id = job_id.split("/")
     with set_directory(trace_id):
+        executor = load_executor(executor)
+        storage = load_storage(storage)
         _, executor = init_executor(executor)
         results = executor.get_results(exec_id)
         results, output_artifacts = handle_output_artifacts(
@@ -247,6 +268,12 @@ class CalculationMCPServer:
                     if preprocess_func is not None:
                         executor, storage, kwargs = preprocess_func(
                             executor, storage, kwargs)
+                    if executor:
+                        with open("executor.json", "w") as f:
+                            json.dump(executor, f, indent=4)
+                    if storage:
+                        with open("storage.json", "w") as f:
+                            json.dump(storage, f, indent=4)
                     kwargs, input_artifacts = handle_input_artifacts(
                         fn, kwargs, storage)
                     executor_type, executor = init_executor(executor)