bohr-agent-sdk 0.1.112__tar.gz → 0.1.114__tar.gz

{bohr_agent_sdk-0.1.112 → bohr_agent_sdk-0.1.114}/PKG-INFO

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

{bohr_agent_sdk-0.1.112 → bohr_agent_sdk-0.1.114}/setup.py

@@ -9,7 +9,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
 
 setup(
     name="bohr-agent-sdk",
-    version="0.1.112",
+    version="0.1.114",
     description="SDK for science agent and mcp tools",
     long_description=long_description,
     long_description_content_type="text/markdown",
@@ -76,4 +76,4 @@ setup(
         "Source": "https://github.com/dptech-corp/bohr-agent-sdk/",
         # "Documentation": "https://bohr-agent-sdk.readthedocs.io/",
     },
-)
+)

{bohr_agent_sdk-0.1.112 → bohr_agent_sdk-0.1.114}/src/bohr_agent_sdk.egg-info/PKG-INFO

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

{bohr_agent_sdk-0.1.112 → bohr_agent_sdk-0.1.114}/src/dp/agent/adapter/adk/client/calculation_mcp_tool.py

@@ -105,6 +105,8 @@ class CalculationMCPTool(MCPTool):
         storage = args["storage"]
         res = await self.submit_tool.run_async(
             args=args, tool_context=tool_context, **kwargs)
+        if isinstance(res, dict):
+            res = types.CallToolResult.model_validate(res)
         if res.isError:
             logger.error(res.content[0].text)
             return res
@@ -125,6 +127,8 @@ class CalculationMCPTool(MCPTool):
             res = await self.query_tool.run_async(
                 args={"job_id": job_id, "executor": executor},
                 tool_context=tool_context, **kwargs)
+            if isinstance(res, dict):
+                res = types.CallToolResult.model_validate(res)
             if res.isError:
                 logger.error(res.content[0].text)
             else:
@@ -138,6 +142,8 @@ class CalculationMCPTool(MCPTool):
         res = await self.results_tool.run_async(
             args={"job_id": job_id, "executor": executor, "storage": storage},
             tool_context=tool_context, **kwargs)
+        if isinstance(res, dict):
+            res = types.CallToolResult.model_validate(res)
         if res.isError:
             await self.log("error", "Job %s failed: %s" % (
                 job_id, res.content[0].text), tool_context)
@@ -246,13 +252,15 @@ class BackgroundJobWatcher:
                     and part.function_response
                     and part.function_response.id in self.long_running_ids
                     and "result" in part.function_response.response
-                    and not part.function_response.response["result"].isError
                 ):
                     result = part.function_response.response["result"]
-                    results = json.loads(result.content[0].text)
-                    job_id = results["job_id"]
-                    self.long_running_jobs[job_id] = part.function_response
-                    self.status[job_id] = "Running"
+                    if isinstance(result, dict):
+                        result = types.CallToolResult.model_validate(result)
+                    if not result.isError:
+                        results = json.loads(result.content[0].text)
+                        job_id = results["job_id"]
+                        self.long_running_jobs[job_id] = part.function_response
+                        self.status[job_id] = "Running"
 
     async def watch_jobs(self):
         for job_id in self.status.keys():
@@ -261,6 +269,8 @@ class BackgroundJobWatcher:
             res = await self.toolset.query_tool.run_async(
                 args={"job_id": job_id, "executor": self.toolset.executor},
                 tool_context=None)
+            if isinstance(res, dict):
+                res = types.CallToolResult.model_validate(res)
             if res.isError:
                 logger.error(res.content[0].text)
                 continue
@@ -270,6 +280,8 @@ class BackgroundJobWatcher:
                     args={"job_id": job_id, "executor": self.toolset.executor,
                           "storage": self.toolset.storage},
                     tool_context=None)
+                if isinstance(res, dict):
+                    res = types.CallToolResult.model_validate(res)
                 job_info = getattr(res.content[0], "job_info", {})
                 response = self.long_running_jobs[job_id]
                 if res.isError:

{bohr_agent_sdk-0.1.112 → bohr_agent_sdk-0.1.114}/src/dp/agent/adapter/adk/utils.py

@@ -31,6 +31,8 @@ def update_session_handler(
     tool_response: dict,
 ) -> Optional[Dict]:
     """Update session state with job and artifact information."""
+    if isinstance(tool_response, dict):
+        tool_response = types.CallToolResult.model_validate(tool_response)
     if len(tool_response.content) == 0 \
             or not hasattr(tool_response.content[0], "text"):
         return None
@@ -83,6 +85,8 @@ def search_error_in_memory_handler(toolset):
         tool: BaseTool, args: Dict[str, Any], tool_context: ToolContext,
         tool_response: dict,
     ) -> Optional[Dict]:
+        if isinstance(tool_response, dict):
+            tool_response = types.CallToolResult.model_validate(tool_response)
         if tool_response.isError:
             err_msg = tool_response.content[0].text
             if err_msg.startswith("Error executing tool"):
@@ -99,6 +103,8 @@ def search_error_in_memory_handler(toolset):
             tools = await toolset.get_tools()
             tool = next(filter(lambda t: t.name == "search_tool_error", tools))
             res = await tool.run_async(args=args, tool_context=None)
+            if isinstance(res, dict):
+                res = types.CallToolResult.model_validate(res)
             result = json.loads(res.content[0].text)
             logger.info("Search tool error result: %s" % result)
             if result.get("results") and result[

{bohr_agent_sdk-0.1.112 → bohr_agent_sdk-0.1.114}/src/dp/agent/server/calculation_mcp_server.py

@@ -17,6 +17,7 @@ from mcp.server.fastmcp.utilities.func_metadata import (
     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
@@ -267,9 +268,35 @@ class SubmitResult(BaseModel):
     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
@@ -277,6 +304,8 @@ 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 = {}