bohr-agent-sdk 0.1.18__tar.gz → 0.1.20__tar.gz

{bohr_agent_sdk-0.1.18 → bohr_agent_sdk-0.1.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.18
+Version: 0.1.20
 Summary: SDK for scientific agents
 Home-page: https://github.com/dptech-corp/bohr-agent-sdk/
 Author: DP Technology

{bohr_agent_sdk-0.1.18 → bohr_agent_sdk-0.1.20}/setup.py

@@ -9,7 +9,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="bohr-agent-sdk",
-    version="0.1.18",
+    version="0.1.20",
     description="SDK for science agent and mcp tools",
     long_description=long_description,
     long_description_content_type="text/markdown",

{bohr_agent_sdk-0.1.18 → bohr_agent_sdk-0.1.20}/src/bohr_agent_sdk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.18
+Version: 0.1.20
 Summary: SDK for scientific agents
 Home-page: https://github.com/dptech-corp/bohr-agent-sdk/
 Author: DP Technology

{bohr_agent_sdk-0.1.18 → bohr_agent_sdk-0.1.20}/src/bohr_agent_sdk.egg-info/SOURCES.txt

@@ -10,6 +10,7 @@ src/bohr_agent_sdk.egg-info/top_level.txt
 src/dp/__init__.py
 src/dp/agent/__init__.py
 src/dp/agent/adapter/adk/__init__.py
+src/dp/agent/adapter/adk/storage_artifact_service.py
 src/dp/agent/adapter/adk/utils.py
 src/dp/agent/adapter/adk/client/__init__.py
 src/dp/agent/adapter/adk/client/calculation_mcp_tool.py

{bohr_agent_sdk-0.1.18 → bohr_agent_sdk-0.1.20}/src/dp/agent/adapter/adk/__init__.py

@@ -3,6 +3,7 @@ from .client import (
     CalculationMCPToolset,
     BackgroundJobWatcher,
 )
+from .storage_artifact_service import StorageArtifactService
 from .utils import (
     search_error_in_memory_handler,
     update_session_handler,
@@ -11,4 +12,5 @@ from .utils import (
 
 __all__ = ["CalculationMCPTool", "CalculationMCPToolset",
            "update_session_handler", "search_error_in_memory_handler",
-           "BackgroundJobWatcher", "extract_job_info"]
+           "BackgroundJobWatcher", "extract_job_info",
+           "StorageArtifactService"]

bohr_agent_sdk-0.1.20/src/dp/agent/adapter/adk/storage_artifact_service.py

@@ -0,0 +1,184 @@
+
+import logging
+import mimetypes
+import os
+import tempfile
+from typing import Optional
+
+from google.adk.artifacts import BaseArtifactService
+from google.genai import types
+from typing_extensions import override
+
+from ...server.storage import BaseStorage
+
+logger = logging.getLogger(__name__)
+
+
+class StorageArtifactService(BaseArtifactService):
+    """An artifact service implementation using storage plugin."""
+    def __init__(self, storage: BaseStorage):
+        self.storage = storage
+
+    def _file_has_user_namespace(self, filename: str) -> bool:
+        """Checks if the filename has a user namespace.
+
+        Args:
+            filename: The filename to check.
+
+        Returns:
+            True if the filename has a user namespace (starts with "user:"),
+            False otherwise.
+        """
+        return filename.startswith("user:")
+
+    def _get_key(
+        self,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        version: int,
+    ) -> str:
+        """Constructs the key.
+
+        Args:
+            app_name: The name of the application.
+            user_id: The ID of the user.
+            session_id: The ID of the session.
+            filename: The name of the artifact file.
+            version: The version of the artifact.
+
+        Returns:
+            The constructed key.
+        """
+        if self._file_has_user_namespace(filename):
+            return f"{app_name}/{user_id}/user/{filename}/{version}"
+        return f"{app_name}/{user_id}/{session_id}/{filename}/{version}"
+
+    @override
+    async def save_artifact(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        artifact: types.Part,
+    ) -> int:
+        versions = await self.list_versions(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+        )
+        version = 0 if not versions else max(versions) + 1
+
+        key = self._get_key(
+            app_name, user_id, session_id, filename, version
+        )
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, filename)
+            with open(path, "wb") as f:
+                f.write(artifact.inline_data.data)
+            self.storage._upload(key, path)
+
+        return version
+
+    @override
+    async def load_artifact(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        version: Optional[int] = None,
+    ) -> Optional[types.Part]:
+        if version is None:
+            versions = await self.list_versions(
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename,
+            )
+            if not versions:
+                return None
+            version = max(versions)
+
+        key = self._get_key(
+            app_name, user_id, session_id, filename, version
+        )
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = os.path.join(tmpdir, filename)
+            self.storage._download(key, path)
+            with open(path, "rb") as f:
+                artifact_bytes = f.read()
+
+        mime_type, _ = mimetypes.guess_type(filename)
+        artifact = types.Part.from_bytes(
+            data=artifact_bytes, mime_type=mime_type
+        )
+        return artifact
+
+    @override
+    async def list_artifact_keys(
+        self, *, app_name: str, user_id: str, session_id: str
+    ) -> list[str]:
+        filenames = set()
+
+        session_prefix = f"{app_name}/{user_id}/{session_id}/"
+        keys = self.storage.list(session_prefix)
+        for key in keys:
+            _, _, _, filename, _ = key.split("/")[-5:]
+            filenames.add(filename)
+
+        user_namespace_prefix = f"{app_name}/{user_id}/user/"
+        user_namespace_keys = self.storage.list(user_namespace_prefix)
+        for key in user_namespace_keys:
+            _, _, _, filename, _ = key.split("/")[-5:]
+            filenames.add(filename)
+
+        return sorted(list(filenames))
+
+    @override
+    async def delete_artifact(
+        self, *, app_name: str, user_id: str, session_id: str, filename: str
+    ) -> None:
+        raise NotImplementedError()
+
+    @override
+    async def list_versions(
+        self, *, app_name: str, user_id: str, session_id: str, filename: str
+    ) -> list[int]:
+        prefix = self._get_key(app_name, user_id, session_id, filename, "")
+        keys = self.storage.list(prefix)
+        versions = []
+        for key in keys:
+            _, _, _, _, version = key.split("/")[-5:]
+            versions.append(int(version))
+        return versions
+
+    async def get_permanent_read_url(
+        self,
+        *,
+        app_name: str,
+        user_id: str,
+        session_id: str,
+        filename: str,
+        version: Optional[int] = None,
+    ) -> str:
+        if version is None:
+            versions = await self.list_versions(
+                app_name=app_name,
+                user_id=user_id,
+                session_id=session_id,
+                filename=filename,
+            )
+            if not versions:
+                return None
+            version = max(versions)
+
+        key = self._get_key(
+            app_name, user_id, session_id, filename, version
+        )
+        return self.storage.get_http_url(key)

{bohr_agent_sdk-0.1.18 → bohr_agent_sdk-0.1.20}/src/dp/agent/server/executor/dispatcher_executor.py

@@ -12,12 +12,15 @@ from dpdispatcher import Machine, Resources, Task, Submission
 
 from .base_executor import BaseExecutor
 from .... import __path__
+
 config = {
     "username": os.environ.get("BOHRIUM_USERNAME", ""),
     "password": os.environ.get("BOHRIUM_PASSWORD", ""),
     "project_id": os.environ.get("BOHRIUM_PROJECT_ID", ""),
     "access_key": os.environ.get("BOHRIUM_ACCESS_KEY", ""),
     "app_key": os.environ.get("BOHRIUM_APP_KEY", "agent"),
+    "bohrium_url": os.environ.get("BOHRIUM_BOHRIUM_URL",
+                                  "https://bohrium.dp.tech"),
 }
 logger = logging.getLogger(__name__)
 
@@ -26,17 +29,17 @@ def get_source_code(fn):
     source_lines, start_line = inspect.getsourcelines(fn)
     source_file = inspect.getsourcefile(fn)
     with open(source_file, "r", encoding="utf-8") as fd:
-        pre_lines = fd.readlines()[:start_line-1]
+        pre_lines = fd.readlines()[:start_line - 1]
     return "".join(pre_lines + source_lines) + "\n"
 
 
 class DispatcherExecutor(BaseExecutor):
     def __init__(
-        self,
-        machine=None,
-        resources=None,
-        python_packages=None,
-        python_executable="python3",
+            self,
+            machine=None,
+            resources=None,
+            python_packages=None,
+            python_executable="python3",
     ):
         """Use DPDispatcher to execute the tool
         Refer to https://docs.deepmodeling.com/projects/dpdispatcher.
@@ -107,7 +110,7 @@ class DispatcherExecutor(BaseExecutor):
         script += "if __name__ == \"__main__\":\n"
         script += "    cwd = os.getcwd()\n"
         script += "    kwargs = jsonpickle.loads(r'''%s''')\n" % \
-            jsonpickle.dumps(kwargs)
+                  jsonpickle.dumps(kwargs)
         script += "    try:\n"
         if import_func_line is not None:
             script += "        " + import_func_line
@@ -118,8 +121,8 @@ class DispatcherExecutor(BaseExecutor):
             script += "        if isinstance(results, dict):\n"
             script += "            for name in results:\n"
             script += "                if isinstance(results[name], Path):\n"
-            script += "                    results[name] = "\
-                "results[name].absolute().relative_to(cwd)\n"
+            script += "                    results[name] = " \
+                      "results[name].absolute().relative_to(cwd)\n"
         script += "    except Exception as e:\n"
         script += "        os.chdir(cwd)\n"
         script += "        with open('err', 'w') as f:\n"
@@ -173,8 +176,8 @@ class DispatcherExecutor(BaseExecutor):
             extra_info = {
                 "bohr_job_id": bohr_job_id,
                 "bohr_group_id": bohr_group_id,
-                "job_link": "https://bohrium.dp.tech/jobs/detail/%s" %
-                bohr_job_id,
+                "job_link": f"{config['bohrium_url']}/jobs/detail/%s" %
+                            bohr_job_id,
             }
             logger.info(extra_info)
             res["extra_info"] = extra_info
@@ -182,8 +185,8 @@ class DispatcherExecutor(BaseExecutor):
             bohr_job_id = submission.belonging_jobs[0].job_id
             extra_info = {
                 "bohr_job_id": bohr_job_id,
-                "job_link": "https://bohrium.dp.tech/jobs/detail/%s" %
-                bohr_job_id,
+                "job_link": f"{config['bohrium_url']}/jobs/detail/%s" %
+                            bohr_job_id,
             }
             logger.info(extra_info)
             res["extra_info"] = extra_info

{bohr_agent_sdk-0.1.18 → bohr_agent_sdk-0.1.20}/src/dp/agent/server/storage/bohrium_storage.py

@@ -50,9 +50,9 @@ def login(username=None, phone=None, password=None, bohrium_url=None):
         password = config["password"]
     if bohrium_url is None:
         bohrium_url = config["bohrium_url"]
-    config["authorization"] = _login(
+    authorization = _login(
         bohrium_url + "/account/login", username, phone, password)
-    return config["authorization"]
+    return authorization
 
 
 def _login(login_url=None, username=None, phone=None, password=None):
@@ -256,6 +256,7 @@ class BohriumStorage(BaseStorage):
         return meta["entityTag"] if "entityTag" in meta else ""
 
     def get_http_url(self, key):
+        key = self.prefixing(key)
         url = self.tiefblue_url + "/api/setacl"
         headers = {
             "Content-type": "application/json",