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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.110
+Version: 0.1.112
 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.112.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=dMIFUQttO7yZyivcuf4BWEZra0Fkh7M7fB2b8tpFix0,17203
 dp/agent/server/preprocessor.py,sha256=XUWu7QOwo_sIDMYS2b1OTrM33EXEVH_73vk-ju1Ok8A,1264
-dp/agent/server/utils.py,sha256=ui3lca9EagcGqmYf8BKLsPARIzXxJ3jgN98yuEO3OSQ,1668
+dp/agent/server/utils.py,sha256=cIKaAg8UaP5yMwvIVTgUVBjy-B3S16bEdnucUf4UDIM,2055
 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.112.dist-info/METADATA,sha256=qnbHAl8Ey0kVMpxD_3Q6qWhOZCC0uocIZXwUQZk11_s,11070
+bohr_agent_sdk-0.1.112.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.112.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.112.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.112.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,17 +7,17 @@ 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 Annotated, Literal, Optional, List, Dict
 
-import mcp
 from mcp.server.fastmcp import FastMCP
 from mcp.server.fastmcp.utilities.context_injection import (
     find_context_parameter,
 )
 from mcp.server.fastmcp.utilities.func_metadata import (
-    _get_typed_signature,
+    ArgModelBase,
     func_metadata,
 )
+from pydantic import BaseModel, Field, create_model
 from starlette.responses import JSONResponse
 from starlette.routing import Route
 
@@ -65,18 +65,9 @@ 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 load_job_info():
+    with open("job.json", "r") as f:
+        return json.load(f)
 
 
 def query_job_status(job_id: str, executor: Optional[dict] = None
@@ -90,7 +81,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 = load_job_info()["executor"] or executor
         _, executor = init_executor(executor)
         status = executor.query_status(exec_id)
         logger.info("Job %s status is %s" % (job_id, status))
@@ -105,7 +96,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 = load_job_info()["executor"] or executor
         _, executor = init_executor(executor)
         executor.terminate(exec_id)
         logger.info("Job %s is terminated" % job_id)
@@ -133,6 +124,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,13 +206,28 @@ 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
 
 
 # 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) -> Optional[Any]:
+                    storage: Optional[dict] = None):
     """
     Get results of a calculation job
     Args:
@@ -168,8 +237,9 @@ def get_job_results(job_id: str, executor: Optional[dict] = None,
     """
     trace_id, exec_id = job_id.split("/")
     with set_directory(trace_id):
-        executor = load_executor(executor)
-        storage = load_storage(storage)
+        job_info = load_job_info()
+        executor = job_info["executor"] or executor
+        storage = job_info["storage"] or storage
         _, executor = init_executor(executor)
         results = executor.get_results(exec_id)
         results, output_artifacts = handle_output_artifacts(
@@ -177,7 +247,24 @@ def get_job_results(job_id: str, executor: Optional[dict] = None,
         logger.info("Job %s result is %s" % (job_id, results))
     return JobResult(result=results, job_info={
         "output_artifacts": output_artifacts,
-    })
+    }, tool_name=job_info["tool_name"])
+
+
+annotation_map = {
+    Path: str,
+    Optional[Path]: Optional[str],
+    List[Path]: List[str],
+    Optional[List[Path]]: Optional[List[str]],
+    Dict[str, Path]: Dict[str, str],
+    Optional[Dict[str, Path]]: Optional[Dict[str, str]],
+    Dict[str, List[Path]]: Dict[str, List[str]],
+    Optional[Dict[str, List[Path]]]: Optional[Dict[str, List[str]]],
+}
+
+
+class SubmitResult(BaseModel):
+    job_id: str
+    extra_info: dict | None = None
 
 
 class CalculationMCPServer:
@@ -191,47 +278,47 @@ class CalculationMCPServer:
         self.preprocess_func = preprocess_func
         self.fastmcp_mode = fastmcp_mode
         self.mcp = FastMCP(*args, **kwargs)
+        self.fn_metadata_map = {}
 
     def add_patched_tool(self, fn, new_fn, name, is_async=False, doc=None,
                          override_return_annotation=False):
         """patch the metadata of the tool"""
         context_kwarg = find_context_parameter(fn)
-
-        def _get_typed_signature_patched(call):
-            """patch parameters"""
-            typed_signature = _get_typed_signature(call)
-            new_typed_signature = _get_typed_signature(new_fn)
-            parameters = []
-            for param in typed_signature.parameters.values():
-                if param.annotation is Path:
-                    parameters.append(inspect.Parameter(
-                        name=param.name, default=param.default,
-                        annotation=str, kind=param.kind))
-                elif param.annotation is Optional[Path]:
-                    parameters.append(inspect.Parameter(
-                        name=param.name, default=param.default,
-                        annotation=Optional[str], kind=param.kind))
-                else:
-                    parameters.append(param)
-            for param in new_typed_signature.parameters.values():
-                if param.name != "kwargs":
-                    parameters.append(param)
-            return inspect.Signature(
-                parameters,
-                return_annotation=(new_typed_signature.return_annotation
-                                   if override_return_annotation
-                                   else typed_signature.return_annotation))
-
-        # Due to the frequent changes of MCP, we use a patching style here
-        mcp.server.fastmcp.utilities.func_metadata._get_typed_signature = \
-            _get_typed_signature_patched
         func_arg_metadata = func_metadata(
             fn,
             skip_names=[context_kwarg] if context_kwarg is not None else [],
-            structured_output=None,
         )
-        mcp.server.fastmcp.utilities.func_metadata._get_typed_signature = \
-            _get_typed_signature
+        self.fn_metadata_map[name] = func_arg_metadata
+        model_params = {}
+        params = inspect.signature(fn, eval_str=True).parameters
+        for n, annotation in \
+                func_arg_metadata.arg_model.__annotations__.items():
+            param = params[n]
+            if param.annotation in annotation_map:
+                model_params[n] = Annotated[
+                    (annotation_map[param.annotation], Field())]
+            else:
+                model_params[n] = annotation
+            if param.default is not inspect.Parameter.empty:
+                model_params[n] = (model_params[n], param.default)
+        for n, param in inspect.signature(new_fn).parameters.items():
+            if n == "kwargs":
+                continue
+            model_params[n] = Annotated[(param.annotation, Field())]
+            if param.default is not inspect.Parameter.empty:
+                model_params[n] = (model_params[n], param.default)
+
+        func_arg_metadata.arg_model = create_model(
+            f"{fn.__name__}Arguments",
+            __base__=ArgModelBase,
+            **model_params,
+        )
+        if override_return_annotation:
+            new_func_arg_metadata = func_metadata(new_fn)
+            func_arg_metadata.output_model = new_func_arg_metadata.output_model
+            func_arg_metadata.output_schema = \
+                new_func_arg_metadata.output_schema
+            func_arg_metadata.wrap_output = new_func_arg_metadata.wrap_output
         if self.fastmcp_mode and func_arg_metadata.wrap_output:
             # Only simulate behavior of fastmcp for output_schema
             func_arg_metadata.output_schema["x-fastmcp-wrap-result"] = True
@@ -240,16 +327,18 @@ class CalculationMCPServer:
         tool = Tool(
             fn=new_fn,
             name=name,
-            description=doc or fn.__doc__,
+            description=doc or fn.__doc__ or "",
             parameters=parameters,
             fn_metadata=func_arg_metadata,
             is_async=is_async,
             context_kwarg=context_kwarg,
+            fn_metadata_map=self.fn_metadata_map,
         )
         self.mcp._tool_manager._tools[name] = tool
 
     def add_tool(self, fn, *args, **kwargs):
-        tool = Tool.from_function(fn, *args, **kwargs)
+        tool = Tool.from_function(
+            fn, *args, fn_metadata_map=self.fn_metadata_map, **kwargs)
         self.mcp._tool_manager._tools[tool.name] = tool
         return tool
 
@@ -260,20 +349,20 @@ class CalculationMCPServer:
         def decorator(fn: Callable) -> Callable:
             def submit_job(executor: Optional[dict] = None,
                            storage: Optional[dict] = None,
-                           **kwargs) -> TypedDict("results", {
-                               "job_id": str, "extra_info": Optional[dict]}):
+                           **kwargs) -> SubmitResult:
                 trace_id = datetime.today().strftime('%Y-%m-%d-%H:%M:%S.%f')
                 logger.info("Job processing (Trace ID: %s)" % trace_id)
                 with set_directory(trace_id):
                     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)
+                    job = {
+                        "tool_name": fn.__name__,
+                        "executor": executor,
+                        "storage": storage,
+                    }
+                    with open("job.json", "w") as f:
+                        json.dump(job, f, indent=4)
                     kwargs, input_artifacts = handle_input_artifacts(
                         fn, kwargs, storage)
                     executor_type, executor = init_executor(executor)
@@ -281,10 +370,10 @@ class CalculationMCPServer:
                     exec_id = res["job_id"]
                     job_id = "%s/%s" % (trace_id, exec_id)
                     logger.info("Job submitted (ID: %s)" % job_id)
-                result = {
-                    "job_id": job_id,
-                    "extra_info": res.get("extra_info"),
-                }
+                result = SubmitResult(
+                    job_id=job_id,
+                    extra_info=res.get("extra_info"),
+                )
                 return JobResult(result=result, job_info={
                     "trace_id": trace_id,
                     "executor_type": executor_type,

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"

dp/agent/server/utils.py

@@ -21,6 +21,7 @@ def get_logger(name, level="INFO",
 class JobResult(BaseModel):
     result: Any
     job_info: dict
+    tool_name: str | None = None
 
 
 class Tool(mcp.server.fastmcp.tools.Tool):
@@ -28,13 +29,23 @@ class Tool(mcp.server.fastmcp.tools.Tool):
     Workaround MCP server cannot print traceback
     Add job info to first unstructured content
     """
+    fn_metadata_map: dict | None = None
+
+    @classmethod
+    def from_function(cls, *args, fn_metadata_map=None, **kwargs):
+        tool = super().from_function(*args, **kwargs)
+        tool.fn_metadata_map = fn_metadata_map
+        return tool
+
     async def run(self, *args, **kwargs):
         try:
             kwargs["convert_result"] = False
             result = await super().run(*args, **kwargs)
             if isinstance(result, JobResult):
                 job_info = result.job_info
-                result = self.fn_metadata.convert_result(result.result)
+                fn_metadata = self.fn_metadata_map.get(
+                    result.tool_name, self.fn_metadata)
+                result = fn_metadata.convert_result(result.result)
                 if isinstance(result, tuple) and len(result) == 2:
                     unstructured_content, _ = result
                 else: