bohr-agent-sdk 0.1.110__py3-none-any.whl → 0.1.111__py3-none-any.whl

{bohr_agent_sdk-0.1.110.dist-info → bohr_agent_sdk-0.1.111.dist-info}/METADATA

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

{bohr_agent_sdk-0.1.110.dist-info → bohr_agent_sdk-0.1.111.dist-info}/RECORD

@@ -51,7 +51,7 @@ dp/agent/cli/templates/ui/server/session_manager.py,sha256=ZbNHGCFvswa-LKWn6c6RM
 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/client/mcp_client.py,sha256=glbQa-fv2aOBhv_GC2ldwyWHOVcSIEwtVLAjzLZvp0c,6289
 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
@@ -62,12 +62,12 @@ 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=hsTxuguyqgq_4HSeJYS7VnW-jTvqf9DI0mpvRj0wK3w,13566
+dp/agent/server/calculation_mcp_server.py,sha256=vvsf58aKbBtH0AqrG5_qhGqg5g2nEhmXgabDzZKpa6o,18534
 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
 dp/agent/server/executor/base_executor.py,sha256=nR2jI-wFvKoOk8QaK11pnSAkHj2MsE6uyzPWDx-vgJA,3018
-dp/agent/server/executor/dispatcher_executor.py,sha256=-P1HAuHDbfLi-3mkT29I84x2aMvmXiM2h-qfZGEobjI,10872
+dp/agent/server/executor/dispatcher_executor.py,sha256=CZRxbVkLaDvStXhNaMKrKcx2Z0tPPVzIxkU1ufqWgYc,12081
 dp/agent/server/executor/local_executor.py,sha256=pGXlDOrfjfP40hSMzbF-Wls3OnOTsH8PdkQjcDjP6_w,6580
 dp/agent/server/storage/__init__.py,sha256=Sgsyp5hb0_hhIGugAPfQFzBHt_854rS_MuMuE3sn8Gs,389
 dp/agent/server/storage/base_storage.py,sha256=728-oNG6N8isV95gZVnyi4vTznJPJhSjxw9Gl5Y_y5o,2356
@@ -75,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.110.dist-info/METADATA,sha256=zqodSHxsxgJ5AYkjncOWOL4qA5hWhiY8Sx1AvPOfMO0,11070
-bohr_agent_sdk-0.1.110.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.110.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.110.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.110.dist-info/RECORD,,
+bohr_agent_sdk-0.1.111.dist-info/METADATA,sha256=savZekEjjj6ToO6eqTYjuaEqgq7jJF-KH4Xzf_XSukM,11070
+bohr_agent_sdk-0.1.111.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.111.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.111.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.111.dist-info/RECORD,,

dp/agent/client/mcp_client.py

@@ -125,7 +125,7 @@ class MCPClient:
 
         executor = arguments.get("executor")
         storage = arguments.get("storage")
-        res = await self.session.call_tool("submit_" + tool_name, arguments)
+        res = await self.call_tool("submit_" + tool_name, arguments)
         if res.isError:
             logger.error("Failed to submit %s: %s" % (
                 tool_name, res.content[0].text))
@@ -137,7 +137,7 @@ class MCPClient:
             logger.info(job_info["extra_info"])
 
         while True:
-            res = await self.session.call_tool("query_job_status", {
+            res = await self.call_tool("query_job_status", {
                 "job_id": job_id, "executor": executor})
             if res.isError:
                 logger.error(res.content[0].text)
@@ -148,7 +148,7 @@ class MCPClient:
                     break
             await asyncio.sleep(self.query_interval)
 
-        res = await self.session.call_tool("get_job_results", {
+        res = await self.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))

dp/agent/server/calculation_mcp_server.py

@@ -7,7 +7,7 @@ from copy import deepcopy
 from datetime import datetime
 from pathlib import Path
 from urllib.parse import urlparse
-from typing import Any, Literal, Optional, TypedDict
+from typing import Any, Literal, Optional, TypedDict, List, Dict
 
 import mcp
 from mcp.server.fastmcp import FastMCP
@@ -133,6 +133,69 @@ def handle_input_artifacts(fn, kwargs, storage):
                 "storage_type": scheme,
                 "uri": uri,
             }
+        elif param.annotation is List[Path] or (
+            param.annotation is Optional[List[Path]] and
+                kwargs.get(name) is not None):
+            uris = kwargs[name]
+            new_paths = []
+            for uri in uris:
+                scheme, key = parse_uri(uri)
+                if scheme == storage_type:
+                    s = storage
+                else:
+                    s = storage_dict[scheme]()
+                path = s.download(key, "inputs/%s" % name)
+                new_paths.append(Path(path))
+                logger.info("Artifact %s downloaded to %s" % (
+                    uri, path))
+            kwargs[name] = new_paths
+            input_artifacts[name] = {
+                "storage_type": storage_type,
+                "uri": uris,
+            }
+        elif param.annotation is Dict[str, Path] or (
+            param.annotation is Optional[Dict[str, Path]] and
+                kwargs.get(name) is not None):
+            uris_dict = kwargs[name]
+            new_paths_dict = {}
+            for key_name, uri in uris_dict.items():
+                scheme, key = parse_uri(uri)
+                if scheme == storage_type:
+                    s = storage
+                else:
+                    s = storage_dict[scheme]()
+                path = s.download(key, f"inputs/{name}/{key_name}")
+                new_paths_dict[key_name] = Path(path)
+                logger.info("Artifact %s (key=%s) downloaded to %s" % (
+                    uri, key_name, path))
+            kwargs[name] = new_paths_dict
+            input_artifacts[name] = {
+                "storage_type": storage_type,
+                "uri": uris_dict,
+            }
+        elif param.annotation is Dict[str, List[Path]] or (
+            param.annotation is Optional[Dict[str, List[Path]]] and
+                kwargs.get(name) is not None):
+            uris_dict = kwargs[name]
+            new_paths_dict = {}
+            for key_name, uris in uris_dict.items():
+                new_paths = []
+                for uri in uris:
+                    scheme, key = parse_uri(uri)
+                    if scheme == storage_type:
+                        s = storage
+                    else:
+                        s = storage_dict[scheme]()
+                    path = s.download(key, f"inputs/{name}/{key_name}")
+                    new_paths.append(Path(path))
+                    logger.info("Artifact %s (key=%s) downloaded to %s" % (
+                        uri, key_name, path))
+                new_paths_dict[key_name] = new_paths
+            kwargs[name] = new_paths_dict
+            input_artifacts[name] = {
+                "storage_type": storage_type,
+                "uri": uris_dict,
+            }
     return kwargs, input_artifacts
 
 
@@ -152,6 +215,20 @@ def handle_output_artifacts(results, exec_id, storage):
                     "storage_type": storage_type,
                     "uri": uri,
                 }
+            elif isinstance(results[name], list) and all(isinstance(item, Path) for item in results[name]):
+                new_uris = []
+                for item in results[name]:
+                    key = storage.upload("%s/outputs/%s" % (exec_id, name),
+                                         item)
+                    uri = storage_type + "://" + key
+                    logger.info("Artifact %s uploaded to %s" % (
+                        item, uri))
+                    new_uris.append(uri)
+                results[name] = new_uris
+                output_artifacts[name] = {
+                    "storage_type": storage_type,
+                    "uri": new_uris,
+                }
     return results, output_artifacts
 
 
@@ -211,6 +288,30 @@ class CalculationMCPServer:
                     parameters.append(inspect.Parameter(
                         name=param.name, default=param.default,
                         annotation=Optional[str], kind=param.kind))
+                elif param.annotation is List[Path]:
+                    parameters.append(inspect.Parameter(
+                        name=param.name, default=param.default,
+                        annotation=List[str], kind=param.kind))
+                elif param.annotation is Optional[List[Path]]:
+                    parameters.append(inspect.Parameter(
+                        name=param.name, default=param.default,
+                        annotation=Optional[List[str]], kind=param.kind))
+                elif param.annotation is Dict[str, Path]:
+                    parameters.append(inspect.Parameter(
+                        name=param.name, default=param.default,
+                        annotation=Dict[str, str], kind=param.kind))
+                elif param.annotation is Optional[Dict[str, Path]]:
+                    parameters.append(inspect.Parameter(
+                        name=param.name, default=param.default,
+                        annotation=Optional[Dict[str, str]], kind=param.kind))
+                elif param.annotation is Dict[str, List[Path]]:
+                    parameters.append(inspect.Parameter(
+                        name=param.name, default=param.default,
+                        annotation=Dict[str, List[str]], kind=param.kind))
+                elif param.annotation is Optional[Dict[str, List[Path]]]:
+                    parameters.append(inspect.Parameter(
+                        name=param.name, default=param.default,
+                        annotation=Optional[Dict[str, List[str]]], kind=param.kind))
                 else:
                     parameters.append(param)
             for param in new_typed_signature.parameters.values():

dp/agent/server/executor/dispatcher_executor.py

@@ -135,7 +135,7 @@ class DispatcherExecutor(BaseExecutor):
         func_def_script, packages = get_func_def_script(fn)
         self.python_packages.extend(packages)
 
-        script += "import asyncio, jsonpickle, os\n"
+        script += "import asyncio, jsonpickle, os, shutil\n"
         script += "from pathlib import Path\n\n"
         script += "if __name__ == \"__main__\":\n"
         script += "    cwd = os.getcwd()\n"
@@ -149,11 +149,24 @@ class DispatcherExecutor(BaseExecutor):
             script += "        results = asyncio.run(%s(**kwargs))\n" % fn_name
         else:
             script += "        results = %s(**kwargs)\n" % fn_name
+            script += "        result_dir = None\n"
+            script += "        import uuid\n"
             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 += "                    if not results[name].absolute().is_relative_to(cwd):\n"
+            script += "                        if result_dir is None:\n"
+            script += "                            result_dir = Path('result_files_dir_' + str(uuid.uuid4()))\n"
+            script += "                            result_dir.mkdir(parents=True, exist_ok=True)\n"
+            script += "                        dest_path = result_dir / results[name].absolute().relative_to('/')\n"
+            script += "                        dest_path.parent.mkdir(parents=True, exist_ok=True)\n"
+            script += "                        if results[name].is_file():\n"
+            script += "                            shutil.copy2(results[name], dest_path)\n"
+            script += "                        elif results[name].is_dir():\n"
+            script += "                            shutil.copytree(results[name], dest_path, dirs_exist_ok=True)\n"
+            script += "                        results[name] = dest_path.absolute().relative_to(cwd)\n"
+            script += "                    else:\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"