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

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

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

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

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

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

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

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

@@ -93,16 +93,19 @@ class CalculationMCPTool(MCPTool):
     async def run_async(self, args, tool_context: ToolContext, **kwargs):
         # TODO: add progress callback when run_async
         args = deepcopy(args)
-        if self.override or "executor" not in args:
+        if self.override:
             args["executor"] = self.executor
-        if self.override or "storage" not in args:
             args["storage"] = self.storage
+        if not args.get("executor"):
+            args.pop("executor", None)
+        if not args.get("storage"):
+            args.pop("storage", None)
         if not self.async_mode and self.wait:
             return await super().run_async(
                 args=args, tool_context=tool_context, **kwargs)
 
-        executor = args["executor"]
-        storage = args["storage"]
+        executor = args.get("executor")
+        storage = args.get("storage")
         res = await self.submit_tool.run_async(
             args=args, tool_context=tool_context, **kwargs)
         if isinstance(res, dict):

{bohr_agent_sdk-0.1.114 → bohr_agent_sdk-0.1.116}/src/dp/agent/client/mcp_client.py

@@ -160,3 +160,10 @@ class MCPClient:
 
     async def cleanup(self):
         await self.exit_stack.aclose()
+
+    async def __aenter__(self):
+        await self.connect_to_server()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.cleanup()

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

@@ -371,7 +371,7 @@ class CalculationMCPServer:
         self.mcp._tool_manager._tools[tool.name] = tool
         return tool
 
-    def tool(self, preprocess_func=None):
+    def tool(self, preprocess_func=None, create_workdir=True):
         if preprocess_func is None:
             preprocess_func = self.preprocess_func
 
@@ -381,7 +381,8 @@ class CalculationMCPServer:
                            **kwargs) -> SubmitResult:
                 trace_id = datetime.today().strftime('%Y-%m-%d-%H:%M:%S.%f')
                 logger.info("Job processing (Trace ID: %s)" % trace_id)
-                with set_directory(trace_id):
+                workdir = trace_id if create_workdir else "."
+                with set_directory(workdir):
                     if preprocess_func is not None:
                         executor, storage, kwargs = preprocess_func(
                             executor, storage, kwargs)
@@ -397,7 +398,7 @@ class CalculationMCPServer:
                     executor_type, executor = init_executor(executor)
                     res = executor.submit(fn, kwargs)
                     exec_id = res["job_id"]
-                    job_id = "%s/%s" % (trace_id, exec_id)
+                    job_id = "%s/%s" % (workdir, exec_id)
                     logger.info("Job submitted (ID: %s)" % job_id)
                 result = SubmitResult(
                     job_id=job_id,
@@ -416,7 +417,8 @@ class CalculationMCPServer:
                 context = self.mcp.get_context()
                 trace_id = datetime.today().strftime('%Y-%m-%d-%H:%M:%S.%f')
                 logger.info("Job processing (Trace ID: %s)" % trace_id)
-                with set_directory(trace_id):
+                workdir = trace_id if create_workdir else "."
+                with set_directory(workdir):
                     if preprocess_func is not None:
                         executor, storage, kwargs = preprocess_func(
                             executor, storage, kwargs)
@@ -424,9 +426,9 @@ class CalculationMCPServer:
                         fn, kwargs, storage)
                     executor_type, executor = init_executor(executor)
                     res = await executor.async_run(
-                        fn, kwargs, context, trace_id)
+                        fn, kwargs, context, workdir)
                     exec_id = res["job_id"]
-                    job_id = "%s/%s" % (trace_id, exec_id)
+                    job_id = "%s/%s" % (workdir, exec_id)
                     results = res["result"]
                     results, output_artifacts = handle_output_artifacts(
                         results, exec_id, storage)

{bohr_agent_sdk-0.1.114 → bohr_agent_sdk-0.1.116}/src/dp/agent/server/executor/local_executor.py

@@ -32,13 +32,13 @@ class Tee(io.TextIOBase):
         return len(text)
 
 
-def wrapped_fn(fn, kwargs, redirect_file=None):
+def wrapped_fn(fn, kwargs, redirect_log=False):
     pid = os.getpid()
     # explicitly reload dflow config
     reload_dflow_config()
-    if redirect_file:
+    if redirect_log:
         stdout = sys.stdout
-        flog = open(redirect_file, "w")
+        flog = open("%s.log" % pid, "w")
         sys.stdout = Tee(flog, stdout)
     try:
         if inspect.iscoroutinefunction(fn):
@@ -46,11 +46,11 @@ def wrapped_fn(fn, kwargs, redirect_file=None):
         else:
             result = fn(**kwargs)
     except Exception as e:
-        with open("err", "w") as f:
+        with open("%s.err" % pid, "w") as f:
             f.write(str(e))
         raise e
     finally:
-        if redirect_file:
+        if redirect_log:
             sys.stdout = stdout
             flog.close()
     with open("%s.txt" % pid, "w") as f:
@@ -111,15 +111,15 @@ class LocalExecutor(BaseExecutor):
         old_env = self.set_env()
         params = {"fn": fn, "kwargs": kwargs}
         if self.dflow:
-            params["redirect_file"] = "log.txt"
+            params["redirect_log"] = True
         p = Process(target=wrapped_fn, kwargs=params)
         p.start()
         extra_info = {}
         if self.dflow:
             while True:
                 alive = p.is_alive()
-                if os.path.isfile("log.txt"):
-                    with open("log.txt", "r") as f:
+                if os.path.isfile("%s.log" % p.pid):
+                    with open("%s.log" % p.pid, "r") as f:
                         log = f.read()
                     match_id = re.search(DFLOW_ID_PATTERN, log)
                     match_link = re.search(DFLOW_LINK_PATTERN, log)
@@ -133,8 +133,8 @@ class LocalExecutor(BaseExecutor):
                         extra_info["workflow_link"] = wf_link
                         break
                 if not alive:
-                    if os.path.isfile("err"):
-                        with open("err", "r") as f:
+                    if os.path.isfile("%s.err" % p.pid):
+                        with open("%s.err" % p.pid, "r") as f:
                             err_msg = f.read()
                     else:
                         err_msg = "No workflow submitted"
@@ -175,8 +175,8 @@ class LocalExecutor(BaseExecutor):
         if os.path.isfile("%s.txt" % job_id):
             with open("%s.txt" % job_id, "r") as f:
                 return jsonpickle.loads(f.read())
-        elif os.path.isfile("err"):
-            with open("err", "r") as f:
+        elif os.path.isfile("%s.err" % job_id):
+            with open("%s.err" % job_id, "r") as f:
                 err_msg = f.read()
             raise RuntimeError(err_msg)
         return {}