bohr-agent-sdk 0.1.24__py3-none-any.whl → 0.1.101__py3-none-any.whl

{bohr_agent_sdk-0.1.24.dist-info → bohr_agent_sdk-0.1.101.dist-info}/METADATA

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

{bohr_agent_sdk-0.1.24.dist-info → bohr_agent_sdk-0.1.101.dist-info}/RECORD

@@ -4,7 +4,7 @@ dp/agent/adapter/adk/__init__.py,sha256=mxX2SxGgDjwQ7WjfKMw9mQkq4ddl8XlCA6gnUinC
 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/adk/client/calculation_mcp_tool.py,sha256=8_uBP12SLaYUiQL6c0vysVsePEZS_r0YMogvW93ZFHo,12188
 dp/agent/adapter/camel/__init__.py,sha256=RN1NhdmsJyN43fTxTXFld4UKZksjpSV0b2QvFn5gK7o,77
 dp/agent/adapter/camel/client/__init__.py,sha256=ld-r0_WsZLFv6yyrmxjWmR8JgnrQzOw4fX0hwVHzciY,93
 dp/agent/adapter/camel/client/calculation_mcp_client.py,sha256=JZZUYYfMgXvHzK4f6IJp-ia33wn3aYZqDSDVa5yKtdc,1860
@@ -20,7 +20,7 @@ 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=oumU2xtqrrgp8BKiuF9YsqMNevXba2JACyYsvNy4nQo,11008
+dp/agent/server/calculation_mcp_server.py,sha256=eClRP7A-t5hMGyTh81KC3GAKjSPNJIylOrOKyzqwo8o,11459
 dp/agent/server/preprocessor.py,sha256=XUWu7QOwo_sIDMYS2b1OTrM33EXEVH_73vk-ju1Ok8A,1264
 dp/agent/server/utils.py,sha256=8jgYZEW4XBp86AF2Km6QkwHltBmrnS-soTpHov7ZEJw,4501
 dp/agent/server/executor/__init__.py,sha256=s95M5qKQk39Yi9qaVJZhk_nfj54quSf7EDghR3OCFUA,248
@@ -33,8 +33,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.24.dist-info/METADATA,sha256=kmGJv8gXvQUdMe_P2QOBp_gz0IeP5oe9VUjmdSsmuZo,6329
-bohr_agent_sdk-0.1.24.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.24.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.24.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.24.dist-info/RECORD,,
+bohr_agent_sdk-0.1.101.dist-info/METADATA,sha256=aPSBkjXZHVZyRBAqdWQl_BNJnVxAcRbuK-izY7ibVXw,6330
+bohr_agent_sdk-0.1.101.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.101.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.101.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.101.dist-info/RECORD,,

dp/agent/adapter/adk/client/calculation_mcp_tool.py

@@ -49,6 +49,7 @@ class CalculationMCPTool(MCPTool):
         results_tool: Optional[MCPTool] = None,
         query_interval: int = 10,
         logging_callback: Callable = logging_handler,
+        override: bool = True,
     ):
         """Calculation MCP tool
         extended from google.adk.tools.mcp_tool.MCPTool
@@ -71,6 +72,7 @@ class CalculationMCPTool(MCPTool):
             results_tool: The tool of getting job results
             query_interval: Time interval of querying job status
             logging_callback: Callback function for server notifications
+            override: Override storage and executor in tool params or not
         """
         self.executor = executor
         self.storage = storage
@@ -82,6 +84,7 @@ class CalculationMCPTool(MCPTool):
         self.query_interval = query_interval
         self.wait = wait
         self.logging_callback = logging_callback
+        self.override = override
 
     async def log(self, level: str, message: Any, tool_context: ToolContext):
         await self.logging_callback(types.LoggingMessageNotificationParams(
@@ -90,9 +93,9 @@ class CalculationMCPTool(MCPTool):
     async def run_async(self, args, tool_context: ToolContext, **kwargs):
         # TODO: add progress callback when run_async
         args = deepcopy(args)
-        if "executor" not in args:
+        if self.override or "executor" not in args:
             args["executor"] = self.executor
-        if "storage" not in args:
+        if self.override or "storage" not in args:
             args["storage"] = self.storage
         if not self.async_mode and self.wait:
             return await super().run_async(
@@ -155,6 +158,7 @@ class CalculationMCPToolset(MCPToolset):
         async_mode: bool = False,
         wait: bool = True,
         logging_callback: Callable = logging_handler,
+        override: bool = True,
         **kwargs,
     ):
         """
@@ -174,7 +178,9 @@ class CalculationMCPToolset(MCPToolset):
                 tools
             async_mode: Submit and query until the job finishes, instead of
                 waiting in single connection
+            wait: Wait for the job to finish or directly return
             logging_callback: Callback function for server notifications
+            override: Override storage and executor in tool params or not
         """
         super().__init__(**kwargs)
         self.logging_callback = logging_callback
@@ -191,6 +197,7 @@ class CalculationMCPToolset(MCPToolset):
         self.query_tool = None
         self.terminate_tool = None
         self.results_tool = None
+        self.override = override
 
     async def get_tools(self, *args, **kwargs) -> List[CalculationMCPTool]:
         tools = await super().get_tools(*args, **kwargs)
@@ -213,6 +220,7 @@ class CalculationMCPToolset(MCPToolset):
                 terminate_tool=tools.get("terminate_job"),
                 results_tool=tools.get("get_job_results"),
                 logging_callback=self.logging_callback,
+                override=self.override,
             )
             calc_tool.__dict__.update(tool.__dict__)
             calc_tool.is_long_running = not self.wait

dp/agent/server/calculation_mcp_server.py

@@ -11,6 +11,8 @@ from typing import Literal, Optional, get_origin
 from mcp.server.fastmcp import FastMCP
 from mcp.server.fastmcp.server import Context
 from mcp.server.fastmcp.utilities.func_metadata import _get_typed_signature
+from starlette.responses import JSONResponse
+from starlette.routing import Route
 
 from .executor import executor_dict
 from .storage import storage_dict
@@ -282,4 +284,16 @@ class CalculationMCPServer:
     def run(self, **kwargs):
         if os.environ.get("DP_AGENT_RUNNING_MODE") in ["1", "true"]:
             return
+        async def health_check(request) :
+            return JSONResponse({"status": "ok"})
+        
+        self.mcp._custom_starlette_routes.append(
+            Route(
+                "/health",
+                endpoint=health_check,
+                methods=["GET"],
+                name="health_check",
+                include_in_schema=True,
+            )
+        )
         self.mcp.run(**kwargs)