bohr-agent-sdk 0.1.107__py3-none-any.whl → 0.1.110__py3-none-any.whl

{bohr_agent_sdk-0.1.107.dist-info → bohr_agent_sdk-0.1.110.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.107
+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.107.dist-info → bohr_agent_sdk-0.1.110.dist-info}/RECORD

@@ -62,21 +62,21 @@ 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=wyOutOTlQEBIJUQQFUFBGE18rQL5fUnjKu88NxbIlhQ,12703
+dp/agent/server/calculation_mcp_server.py,sha256=hsTxuguyqgq_4HSeJYS7VnW-jTvqf9DI0mpvRj0wK3w,13566
 dp/agent/server/preprocessor.py,sha256=XUWu7QOwo_sIDMYS2b1OTrM33EXEVH_73vk-ju1Ok8A,1264
 dp/agent/server/utils.py,sha256=ui3lca9EagcGqmYf8BKLsPARIzXxJ3jgN98yuEO3OSQ,1668
 dp/agent/server/executor/__init__.py,sha256=s95M5qKQk39Yi9qaVJZhk_nfj54quSf7EDghR3OCFUA,248
-dp/agent/server/executor/base_executor.py,sha256=EFJBsYVYAvuRbiLAbLOwLTw3h7ScjN025xnSP4uJHrQ,2052
-dp/agent/server/executor/dispatcher_executor.py,sha256=wUJEmCrLzckwinOd8Caf5H9TN-YVfHLpfYgkVFd2OC0,10828
-dp/agent/server/executor/local_executor.py,sha256=wYCclNZFkLb3v7KpW1nCnupO8piBES-esYlDAuz86zk,6120
+dp/agent/server/executor/base_executor.py,sha256=nR2jI-wFvKoOk8QaK11pnSAkHj2MsE6uyzPWDx-vgJA,3018
+dp/agent/server/executor/dispatcher_executor.py,sha256=-P1HAuHDbfLi-3mkT29I84x2aMvmXiM2h-qfZGEobjI,10872
+dp/agent/server/executor/local_executor.py,sha256=pGXlDOrfjfP40hSMzbF-Wls3OnOTsH8PdkQjcDjP6_w,6580
 dp/agent/server/storage/__init__.py,sha256=Sgsyp5hb0_hhIGugAPfQFzBHt_854rS_MuMuE3sn8Gs,389
 dp/agent/server/storage/base_storage.py,sha256=728-oNG6N8isV95gZVnyi4vTznJPJhSjxw9Gl5Y_y5o,2356
 dp/agent/server/storage/bohrium_storage.py,sha256=EsKX4dWWvZTn2TEhZv4zsvihfDK0mmPFecrln-Ytk40,10488
 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.107.dist-info/METADATA,sha256=YDcLpn0uLCRpNsAOCW46X8v7-HRiwp2g_DTIHD_99_4,11070
-bohr_agent_sdk-0.1.107.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.107.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.107.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.107.dist-info/RECORD,,
+bohr_agent_sdk-0.1.110.dist-info/METADATA,sha256=zqodSHxsxgJ5AYkjncOWOL4qA5hWhiY8Sx1AvPOfMO0,11070
+bohr_agent_sdk-0.1.110.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.110.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.110.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.110.dist-info/RECORD,,

dp/agent/server/calculation_mcp_server.py

@@ -1,4 +1,5 @@
 import inspect
+import json
 import os
 from collections.abc import Callable
 from contextlib import contextmanager
@@ -64,6 +65,20 @@ 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 query_job_status(job_id: str, executor: Optional[dict] = None
                      ) -> Literal["Running", "Succeeded", "Failed"]:
     """
@@ -75,6 +90,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 = init_executor(executor)
         status = executor.query_status(exec_id)
         logger.info("Job %s status is %s" % (job_id, status))
@@ -89,6 +105,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 = init_executor(executor)
         executor.terminate(exec_id)
         logger.info("Job %s is terminated" % job_id)
@@ -151,6 +168,8 @@ 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)
         _, executor = init_executor(executor)
         results = executor.get_results(exec_id)
         results, output_artifacts = handle_output_artifacts(
@@ -249,6 +268,12 @@ class CalculationMCPServer:
                     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)
                     kwargs, input_artifacts = handle_input_artifacts(
                         fn, kwargs, storage)
                     executor_type, executor = init_executor(executor)

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(

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)

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):