bohr-agent-sdk 0.1.24__tar.gz → 0.1.101__tar.gz

{bohr_agent_sdk-0.1.24 → bohr_agent_sdk-0.1.101}/PKG-INFO

@@ -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 → bohr_agent_sdk-0.1.101}/setup.py

@@ -9,7 +9,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="bohr-agent-sdk",
-    version="0.1.24",
+    version="0.1.101",
     description="SDK for science agent and mcp tools",
     long_description=long_description,
     long_description_content_type="text/markdown",

{bohr_agent_sdk-0.1.24 → bohr_agent_sdk-0.1.101}/src/bohr_agent_sdk.egg-info/PKG-INFO

@@ -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 → bohr_agent_sdk-0.1.101}/src/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

{bohr_agent_sdk-0.1.24 → bohr_agent_sdk-0.1.101}/src/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)