bohr-agent-sdk 0.1.17__py3-none-any.whl → 0.1.19__py3-none-any.whl

{bohr_agent_sdk-0.1.17.dist-info → bohr_agent_sdk-0.1.19.dist-info}/METADATA

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

{bohr_agent_sdk-0.1.17.dist-info → bohr_agent_sdk-0.1.19.dist-info}/RECORD

@@ -1,7 +1,8 @@
 dp/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 dp/agent/__init__.py,sha256=4kwF1khocxEkqJjnVhBV5cCsWkR_MBGrUthuHk44pT8,153
-dp/agent/adapter/adk/__init__.py,sha256=x3OZoxHzBcSP3_LAK9oDulRN1vYkEfMEm5ZdocTLlQA,395
-dp/agent/adapter/adk/utils.py,sha256=S6Y9IuQngKXMg0fy8mcCiuAUvOGKYm0xzFFYicgc41I,7025
+dp/agent/adapter/adk/__init__.py,sha256=mxX2SxGgDjwQ7WjfKMw9mQkq4ddl8XlCA6gnUinCW0w,493
+dp/agent/adapter/adk/storage_artifact_service.py,sha256=Hd1CmEB8HgEz3x_riHVSSmYdw6W2bDYYa3c5ySCqmCE,5414
+dp/agent/adapter/adk/utils.py,sha256=Hf6zZOmNi0YSI9t5acmR5kzowWaYr-UfyPiN8fefPkI,7359
 dp/agent/adapter/adk/client/__init__.py,sha256=F1xfFNa4ZG8jV9adeGI2D3YBiSX-5RkvqEdTqNdJce4,209
 dp/agent/adapter/adk/client/calculation_mcp_tool.py,sha256=njEf4DvZt7HaAMpMZ1od2emWwlFUP5FtebBJ17pi20U,11774
 dp/agent/adapter/camel/__init__.py,sha256=RN1NhdmsJyN43fTxTXFld4UKZksjpSV0b2QvFn5gK7o,77
@@ -28,12 +29,12 @@ dp/agent/server/executor/dispatcher_executor.py,sha256=zp9RDUoT2RYZGbKXsUsk4HbS0
 dp/agent/server/executor/local_executor.py,sha256=alNpEPicZQOQMtXc1vqGSo4El_9agXAJ5ldqi-OJGL4,3887
 dp/agent/server/storage/__init__.py,sha256=Sgsyp5hb0_hhIGugAPfQFzBHt_854rS_MuMuE3sn8Gs,389
 dp/agent/server/storage/base_storage.py,sha256=728-oNG6N8isV95gZVnyi4vTznJPJhSjxw9Gl5Y_y5o,2356
-dp/agent/server/storage/bohrium_storage.py,sha256=iJqssa6JX3crlaTYGdpmzgOa0X_AbFxYlTESm_k8C88,10474
+dp/agent/server/storage/bohrium_storage.py,sha256=EsKX4dWWvZTn2TEhZv4zsvihfDK0mmPFecrln-Ytk40,10488
 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.17.dist-info/METADATA,sha256=Wn0wUU7TBvmlJZeswoD6wtr7tiYX8WRDDbaVI-Ilwt4,6329
-bohr_agent_sdk-0.1.17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.17.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.17.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.17.dist-info/RECORD,,
+bohr_agent_sdk-0.1.19.dist-info/METADATA,sha256=tDUixjAnwCjzdtcuXXc-71c2ZhUiJgKmRlXVMQguDEE,6329
+bohr_agent_sdk-0.1.19.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.19.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.19.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.19.dist-info/RECORD,,

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"]

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)

dp/agent/adapter/adk/utils.py

@@ -116,6 +116,7 @@ def search_error_in_memory_handler(toolset):
 def extract_job_info(events: List[Event]) -> dict:
     jobs = {}
     artifacts = {}
+    events.sort(key=lambda event: event.timestamp)
     for event in events:
         if event.content and event.content.parts:
             for part in event.content.parts:
@@ -135,39 +136,42 @@ def extract_job_info(events: List[Event]) -> dict:
                         }
                     job = jobs[resp.id]
                     job["timestamp"] = event.timestamp
-                    if "result" in resp.response and isinstance(
-                            resp.response["result"], types.CallToolResult):
-                        res = resp.response["result"]
+                    res = resp.response.get("result")
+                    if isinstance(res, dict) and "content" in res \
+                            and "isError" in res:
+                        res = types.CallToolResult(
+                            content=res["content"], isError=res["isError"])
+                    if isinstance(res, types.CallToolResult):
                         if res.isError:
                             err_msg = res.content[0].text
                             if err_msg.startswith("Error executing tool"):
                                 err_msg = err_msg[err_msg.find(":")+2:]
                             job["err_msg"] = err_msg
-                        else:
+                        elif hasattr(res.content[0], "text"):
                             result = jsonpickle.loads(res.content[0].text)
                             if "job_id" not in result:
                                 job["result"] = result
-                    if hasattr(res.content[0], "job_info"):
-                        job_info = res.content[0].job_info
-                        job.update(job_info)
-                        for name, art in job_info.get(
-                                "input_artifacts", {}).items():
-                            if art["uri"] not in artifacts:
-                                artifacts[art["uri"]] = {
-                                    "type": "input",
-                                    "name": name,
-                                    "job_id": job_info["job_id"],
-                                    **art,
-                                }
-                        for name, art in job_info.get(
-                                "output_artifacts", {}).items():
-                            if art["uri"] not in artifacts:
-                                artifacts[art["uri"]] = {
-                                    "type": "output",
-                                    "name": name,
-                                    "job_id": job_info["job_id"],
-                                    **art,
-                                }
+                        if hasattr(res.content[0], "job_info"):
+                            job_info = res.content[0].job_info
+                            job.update(job_info)
+                            for name, art in job_info.get(
+                                    "input_artifacts", {}).items():
+                                if art["uri"] not in artifacts:
+                                    artifacts[art["uri"]] = {
+                                        "type": "input",
+                                        "name": name,
+                                        "job_id": job_info["job_id"],
+                                        **art,
+                                    }
+                            for name, art in job_info.get(
+                                    "output_artifacts", {}).items():
+                                if art["uri"] not in artifacts:
+                                    artifacts[art["uri"]] = {
+                                        "type": "output",
+                                        "name": name,
+                                        "job_id": job_info["job_id"],
+                                        **art,
+                                    }
     return {
         "jobs": list(jobs.values()),
         "artifacts": list(artifacts.values()),

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",