bohr-agent-sdk 0.1.114__py3-none-any.whl → 0.1.116__py3-none-any.whl

{bohr_agent_sdk-0.1.114.dist-info → bohr_agent_sdk-0.1.116.dist-info}/METADATA

@@ -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.dist-info → bohr_agent_sdk-0.1.116.dist-info}/RECORD

@@ -4,7 +4,7 @@ dp/agent/adapter/adk/__init__.py,sha256=mxX2SxGgDjwQ7WjfKMw9mQkq4ddl8XlCA6gnUinC
 dp/agent/adapter/adk/storage_artifact_service.py,sha256=Hd1CmEB8HgEz3x_riHVSSmYdw6W2bDYYa3c5ySCqmCE,5414
 dp/agent/adapter/adk/utils.py,sha256=UPCTpvAUCV2V1o0V8dmmFWbiLuUf6fJLIvVX68M3L-M,7698
 dp/agent/adapter/adk/client/__init__.py,sha256=F1xfFNa4ZG8jV9adeGI2D3YBiSX-5RkvqEdTqNdJce4,209
-dp/agent/adapter/adk/client/calculation_mcp_tool.py,sha256=tY0XXATaawCzUMZzDS8JvVvtXhfWqGYpZdFrAjQH0yk,12792
+dp/agent/adapter/adk/client/calculation_mcp_tool.py,sha256=VDiDSEDsOYduGuCE1mp2cmOoDkdyfHhaTFhNRv0GnMY,12873
 dp/agent/adapter/camel/__init__.py,sha256=RN1NhdmsJyN43fTxTXFld4UKZksjpSV0b2QvFn5gK7o,77
 dp/agent/adapter/camel/client/__init__.py,sha256=ld-r0_WsZLFv6yyrmxjWmR8JgnrQzOw4fX0hwVHzciY,93
 dp/agent/adapter/camel/client/calculation_mcp_client.py,sha256=JZZUYYfMgXvHzK4f6IJp-ia33wn3aYZqDSDVa5yKtdc,1860
@@ -51,7 +51,7 @@ dp/agent/cli/templates/ui/server/session_manager.py,sha256=ZbNHGCFvswa-LKWn6c6RM
 dp/agent/cli/templates/ui/server/user_files.py,sha256=khkiyY2UOOysHqO6JgCPUDqtrInp83G1M62i3Lj-0aY,2995
 dp/agent/cli/templates/ui/server/utils.py,sha256=f4NfwFBq_RdZyFn_KCW6ZThYW8TvQyVruK7PJZ-DA80,1530
 dp/agent/client/__init__.py,sha256=yu7HYZwAkD7g5dL9JttLkGmspBcyOf-6OoCjci4oPDA,59
-dp/agent/client/mcp_client.py,sha256=glbQa-fv2aOBhv_GC2ldwyWHOVcSIEwtVLAjzLZvp0c,6289
+dp/agent/client/mcp_client.py,sha256=ZWsegHF_ar8IdRHfG_w6smdSNgmIc9I25rz8Eh4cdZA,6469
 dp/agent/cloud/__init__.py,sha256=e16ymCZX2f-S8DyGB5jSK8gnQqVObRIsvtLXLALIKxQ,441
 dp/agent/cloud/main.py,sha256=5QIEjpZ1RxWnR8wyLf-vlgz1bn9oOnxCYn158LBaLN4,727
 dp/agent/cloud/mcp.py,sha256=tsAwC3doVMLYr6Oh8PxVqF-qCygYkDZJTIhoF_h8eGQ,4537
@@ -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=a0hKNVz-WoUbL8y9GhDx1hO830frvmdtvXTBf7V40lI,18478
+dp/agent/server/calculation_mcp_server.py,sha256=L07b0JNMkEjF6UvdKWTkRB14yehXo8QoU4or7lCwiNw,18618
 dp/agent/server/preprocessor.py,sha256=XUWu7QOwo_sIDMYS2b1OTrM33EXEVH_73vk-ju1Ok8A,1264
 dp/agent/server/utils.py,sha256=cIKaAg8UaP5yMwvIVTgUVBjy-B3S16bEdnucUf4UDIM,2055
 dp/agent/server/executor/__init__.py,sha256=s95M5qKQk39Yi9qaVJZhk_nfj54quSf7EDghR3OCFUA,248
 dp/agent/server/executor/base_executor.py,sha256=nR2jI-wFvKoOk8QaK11pnSAkHj2MsE6uyzPWDx-vgJA,3018
 dp/agent/server/executor/dispatcher_executor.py,sha256=CZRxbVkLaDvStXhNaMKrKcx2Z0tPPVzIxkU1ufqWgYc,12081
-dp/agent/server/executor/local_executor.py,sha256=pGXlDOrfjfP40hSMzbF-Wls3OnOTsH8PdkQjcDjP6_w,6580
+dp/agent/server/executor/local_executor.py,sha256=qRVnfqhvaCGBXr-NO4uxcUteFja2BE_6NXauVJ8vnoo,6642
 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.114.dist-info/METADATA,sha256=-ygUw-5TxXPWb_ti95WAIlX6mrjLwAukzr4RdOLpoEA,11070
-bohr_agent_sdk-0.1.114.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.114.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.114.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.114.dist-info/RECORD,,
+bohr_agent_sdk-0.1.116.dist-info/METADATA,sha256=KycMnrdj1v0eud8EOYzXhMSPX6ISX0rE8aPNou4oJlk,11070
+bohr_agent_sdk-0.1.116.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.116.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.116.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.116.dist-info/RECORD,,

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

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

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)

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 {}