bohr-agent-sdk 0.1.105__tar.gz → 0.1.107__tar.gz

{bohr_agent_sdk-0.1.105 → bohr_agent_sdk-0.1.107}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.105
+Version: 0.1.107
 Summary: SDK for scientific agents
 Home-page: https://github.com/dptech-corp/bohr-agent-sdk/
 Author: DP Technology
@@ -30,7 +30,7 @@ Requires-Dist: paho-mqtt>=2.1.0
 Requires-Dist: redis>=6.2.0
 Requires-Dist: twine>=6.1.0
 Requires-Dist: build>=1.2.2.post1
-Requires-Dist: cloudpickle==2.2.0
+Requires-Dist: cloudpickle
 Requires-Dist: watchdog>=6.0.0
 Requires-Dist: fastapi>=0.116.0
 Requires-Dist: bohrium-open-sdk
@@ -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.105 → bohr_agent_sdk-0.1.107}/README.md

@@ -200,6 +200,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.105 → bohr_agent_sdk-0.1.107}/pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "redis>=6.2.0",
     "twine>=6.1.0",
     "build>=1.2.2.post1",
-    "cloudpickle==2.2.0",
+    "cloudpickle",
     "watchdog>=6.0.0",
     "fastapi>=0.116.0",
     "bohrium-open-sdk"
@@ -131,4 +131,4 @@ requires = [
     "setuptools>=61.0.0",
     "wheel>=0.37.0",
 ]
-build-backend = "setuptools.build_meta"
+build-backend = "setuptools.build_meta"

{bohr_agent_sdk-0.1.105 → bohr_agent_sdk-0.1.107}/setup.py

@@ -9,7 +9,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="bohr-agent-sdk",
-    version="0.1.105",
+    version="0.1.107",
     description="SDK for science agent and mcp tools",
     long_description=long_description,
     long_description_content_type="text/markdown",
@@ -35,6 +35,7 @@ setup(
         "google-genai",          # Google AI SDK
         "google-generativeai",   # Google Generative AI
         "aiofiles",
+        "bohrium-open-sdk==0.1.5",
     ],
     # extras_require={
     #     "dev": [
@@ -75,4 +76,4 @@ setup(
         "Source": "https://github.com/dptech-corp/bohr-agent-sdk/",
         # "Documentation": "https://bohr-agent-sdk.readthedocs.io/",
     },
-)
+)

{bohr_agent_sdk-0.1.105 → bohr_agent_sdk-0.1.107}/src/bohr_agent_sdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.105
+Version: 0.1.107
 Summary: SDK for scientific agents
 Home-page: https://github.com/dptech-corp/bohr-agent-sdk/
 Author: DP Technology
@@ -30,7 +30,7 @@ Requires-Dist: paho-mqtt>=2.1.0
 Requires-Dist: redis>=6.2.0
 Requires-Dist: twine>=6.1.0
 Requires-Dist: build>=1.2.2.post1
-Requires-Dist: cloudpickle==2.2.0
+Requires-Dist: cloudpickle
 Requires-Dist: watchdog>=6.0.0
 Requires-Dist: fastapi>=0.116.0
 Requires-Dist: bohrium-open-sdk
@@ -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.105 → bohr_agent_sdk-0.1.107}/src/bohr_agent_sdk.egg-info/SOURCES.txt

@@ -59,6 +59,8 @@ src/dp/agent/cli/templates/ui/server/models.py
 src/dp/agent/cli/templates/ui/server/session_manager.py
 src/dp/agent/cli/templates/ui/server/user_files.py
 src/dp/agent/cli/templates/ui/server/utils.py
+src/dp/agent/client/__init__.py
+src/dp/agent/client/mcp_client.py
 src/dp/agent/cloud/__init__.py
 src/dp/agent/cloud/main.py
 src/dp/agent/cloud/mcp.py

{bohr_agent_sdk-0.1.105 → bohr_agent_sdk-0.1.107}/src/bohr_agent_sdk.egg-info/requires.txt

@@ -10,7 +10,7 @@ paho-mqtt>=2.1.0
 redis>=6.2.0
 twine>=6.1.0
 build>=1.2.2.post1
-cloudpickle==2.2.0
+cloudpickle
 watchdog>=6.0.0
 fastapi>=0.116.0
 bohrium-open-sdk

{bohr_agent_sdk-0.1.105 → bohr_agent_sdk-0.1.107}/src/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:

{bohr_agent_sdk-0.1.105 → bohr_agent_sdk-0.1.107}/src/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

{bohr_agent_sdk-0.1.105 → bohr_agent_sdk-0.1.107}/src/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):
         """

bohr_agent_sdk-0.1.107/src/dp/agent/client/__init__.py

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

bohr_agent_sdk-0.1.107/src/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()

{bohr_agent_sdk-0.1.105 → bohr_agent_sdk-0.1.107}/src/dp/agent/server/calculation_mcp_server.py

@@ -138,14 +138,16 @@ 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):