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

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

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

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

@@ -62,9 +62,9 @@ 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=vvsf58aKbBtH0AqrG5_qhGqg5g2nEhmXgabDzZKpa6o,18534
+dp/agent/server/calculation_mcp_server.py,sha256=a0hKNVz-WoUbL8y9GhDx1hO830frvmdtvXTBf7V40lI,18478
 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=CZRxbVkLaDvStXhNaMKrKcx2Z0tPPVzIxkU1ufqWgYc,12081
@@ -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.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,,
+bohr_agent_sdk-0.1.113.dist-info/METADATA,sha256=66soACYq7kToAXVjiDX4A_jRkddvoSOLah9QPFrUBxs,11070
+bohr_agent_sdk-0.1.113.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.113.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.113.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.113.dist-info/RECORD,,

dp/agent/server/calculation_mcp_server.py

@@ -7,17 +7,18 @@ from copy import deepcopy
 from datetime import datetime
 from pathlib import Path
 from urllib.parse import urlparse
-from typing import Any, Literal, Optional, TypedDict, List, Dict
+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 mcp.server.sse import SseServerTransport
+from pydantic import BaseModel, Field, create_model
 from starlette.responses import JSONResponse
 from starlette.routing import Route
 
@@ -65,18 +66,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 +82,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 +97,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)
@@ -215,7 +207,8 @@ 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]):
+            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),
@@ -235,7 +228,7 @@ def handle_output_artifacts(results, exec_id, storage):
 # 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:
@@ -245,8 +238,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(
@@ -254,12 +248,55 @@ 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
+
+
+def patch_mcp_close_connection():
+    _mock_orig_handle_post_message = SseServerTransport.handle_post_message
+
+    async def _mock_handle_post_message_with_close(self, scope, receive, send):
+        async def _send(message):
+            if message.get("type") == "http.response.start":
+                headers = list(message.get("headers", []))
+                headers = [
+                    (name, value)
+                    for name, value in headers
+                    if name.lower() != b"connection"
+                ]
+                headers.append((b"connection", b"close"))
+                message["headers"] = headers
+            elif message.get("type") == "http.response.body":
+                message["more_body"] = False
+            await send(message)
+        await _mock_orig_handle_post_message(self, scope, receive, _send)
+
+    if not getattr(SseServerTransport.handle_post_message,
+                   "__patched_close__", False):
+        SseServerTransport.handle_post_message = \
+            _mock_handle_post_message_with_close
+        SseServerTransport.handle_post_message.__patched_close__ = True
 
 
 class CalculationMCPServer:
     def __init__(self, *args, preprocess_func=None, fastmcp_mode=False,
-                 **kwargs):
+                 patch_close_connection=False, **kwargs):
         """
         Args:
             preprocess_func: The preprocess function for all tools
@@ -267,72 +304,50 @@ class CalculationMCPServer:
         """
         self.preprocess_func = preprocess_func
         self.fastmcp_mode = fastmcp_mode
+        if patch_close_connection:
+            patch_mcp_close_connection()
         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))
-                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():
-                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
@@ -341,16 +356,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
 
@@ -361,20 +378,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)
@@ -382,10 +399,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/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: