bohr-agent-sdk 0.1.109__tar.gz → 0.1.110__tar.gz

{bohr_agent_sdk-0.1.109 → bohr_agent_sdk-0.1.110}/PKG-INFO

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

{bohr_agent_sdk-0.1.109 → bohr_agent_sdk-0.1.110}/setup.py

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

{bohr_agent_sdk-0.1.109 → bohr_agent_sdk-0.1.110}/src/bohr_agent_sdk.egg-info/PKG-INFO

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

{bohr_agent_sdk-0.1.109 → bohr_agent_sdk-0.1.110}/src/dp/agent/server/executor/base_executor.py

@@ -5,6 +5,8 @@ from collections.abc import Callable
 from typing import Any, Literal, TypedDict
 
 from mcp.server.fastmcp.server import Context
+from mcp.shared.context import RequestContext
+from starlette.requests import Request
 logger = logging.getLogger(__name__)
 
 
@@ -27,6 +29,26 @@ class BaseExecutor(ABC):
     def get_results(self, job_id: str) -> dict:
         pass
 
+    def prune_context(self, kwargs: dict):
+        for key, value in kwargs.items():
+            if isinstance(value, Context):
+                context = Context(request_context=RequestContext(
+                    request_id=value.request_context.request_id,
+                    meta=value.request_context.meta,
+                    session=None,
+                    lifespan_context=value.request_context.lifespan_context,
+                    request=Request(
+                        scope={
+                            k: v for k, v in
+                            value.request_context.request.scope.items()
+                            if k not in ["app", "router", "endpoint",
+                                         "starlette.exception_handlers"]
+                        },
+                    )
+                ))
+                kwargs[key] = context
+        return kwargs
+
     async def async_run(
         self, fn: Callable, kwargs: dict, context: Context,
         trace_id: str) -> TypedDict(

{bohr_agent_sdk-0.1.109 → bohr_agent_sdk-0.1.110}/src/dp/agent/server/executor/dispatcher_executor.py

@@ -129,6 +129,7 @@ class DispatcherExecutor(BaseExecutor):
         self.resources["envs"]["DP_AGENT_RUNNING_MODE"] = "1"
 
     def submit(self, fn, kwargs):
+        kwargs = self.prune_context(kwargs)
         script = ""
         fn_name = fn.__name__
         func_def_script, packages = get_func_def_script(fn)

{bohr_agent_sdk-0.1.109 → bohr_agent_sdk-0.1.110}/src/dp/agent/server/executor/local_executor.py

@@ -88,6 +88,7 @@ class LocalExecutor(BaseExecutor):
         """
         self.env = env or {}
         self.dflow = dflow
+        self.workflow_id = None
 
     def set_env(self):
         old_env = {}
@@ -105,6 +106,7 @@ class LocalExecutor(BaseExecutor):
                 del os.environ[k]
 
     def submit(self, fn, kwargs):
+        kwargs = self.prune_context(kwargs)
         os.environ["DP_AGENT_RUNNING_MODE"] = "1"
         old_env = self.set_env()
         params = {"fn": fn, "kwargs": kwargs}
@@ -123,6 +125,7 @@ class LocalExecutor(BaseExecutor):
                     match_link = re.search(DFLOW_LINK_PATTERN, log)
                     if match_id and match_link:
                         wf_id = match_id.group(1)
+                        self.workflow_id = wf_id
                         wf_uid = match_id.group(2)
                         wf_link = match_link.group(1)
                         extra_info["workflow_id"] = wf_id
@@ -155,11 +158,18 @@ class LocalExecutor(BaseExecutor):
             return "Failed"
 
     def terminate(self, job_id):
+        if self.workflow_id is not None:
+            try:
+                from dflow import Workflow
+                wf = Workflow(id=self.workflow_id)
+                wf.terminate()
+            except Exception as e:
+                logger.error(f"Failed to terminate workflow: {e}")
         try:
             p = psutil.Process(int(job_id))
             p.terminate()
-        except Exception:
-            pass
+        except Exception as e:
+            logger.error(f"Failed to terminate process: {e}")
 
     def get_results(self, job_id):
         if os.path.isfile("%s.txt" % job_id):