bohr-agent-sdk 0.1.20__py3-none-any.whl → 0.1.22__py3-none-any.whl

{bohr_agent_sdk-0.1.20.dist-info → bohr_agent_sdk-0.1.22.dist-info}/METADATA

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

{bohr_agent_sdk-0.1.20.dist-info → bohr_agent_sdk-0.1.22.dist-info}/RECORD

@@ -20,21 +20,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=omPcrjVT3k8ek3BbB5Dc3TRjdSC_0IEqjCrK2Ka0YqE,10888
+dp/agent/server/calculation_mcp_server.py,sha256=oumU2xtqrrgp8BKiuF9YsqMNevXba2JACyYsvNy4nQo,11008
 dp/agent/server/preprocessor.py,sha256=XUWu7QOwo_sIDMYS2b1OTrM33EXEVH_73vk-ju1Ok8A,1264
 dp/agent/server/utils.py,sha256=8jgYZEW4XBp86AF2Km6QkwHltBmrnS-soTpHov7ZEJw,4501
 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=urpzmKH_tBOgblBdJEa3y8eEhXqUDrdcdWCnUdJpfZk,9420
-dp/agent/server/executor/local_executor.py,sha256=alNpEPicZQOQMtXc1vqGSo4El_9agXAJ5ldqi-OJGL4,3887
+dp/agent/server/executor/local_executor.py,sha256=9JEtZTEbKsmHAY5nsdLZ36gH5M_rSG6En0W9IrCUwm0,5803
 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.20.dist-info/METADATA,sha256=ZEoxrydj541JG9txn0Wr0iARoxZ_w2KA6E0aJFFGlAY,6329
-bohr_agent_sdk-0.1.20.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.20.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.20.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.20.dist-info/RECORD,,
+bohr_agent_sdk-0.1.22.dist-info/METADATA,sha256=IyYt8bph_mRJ06nyUlf1BQnAm1HBIR2jssavB4nTrwU,6329
+bohr_agent_sdk-0.1.22.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.22.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.22.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.22.dist-info/RECORD,,

dp/agent/server/calculation_mcp_server.py

@@ -29,7 +29,7 @@ def parse_uri(uri):
 
 
 def init_storage(storage_config: Optional[dict] = None):
-    if storage_config is None:
+    if not storage_config:
         storage_config = {"type": "local"}
     storage_config = deepcopy(storage_config)
     storage_type = storage_config.pop("type")
@@ -38,7 +38,7 @@ def init_storage(storage_config: Optional[dict] = None):
 
 
 def init_executor(executor_config: Optional[dict] = None):
-    if executor_config is None:
+    if not executor_config:
         executor_config = {"type": "local"}
     executor_config = deepcopy(executor_config)
     executor_type = executor_config.pop("type")
@@ -227,7 +227,11 @@ class CalculationMCPServer:
                     exec_id = res["job_id"]
                     job_id = "%s/%s" % (trace_id, exec_id)
                     logger.info("Job submitted (ID: %s)" % job_id)
-                return convert_to_content({"job_id": job_id}, job_info={
+                result = {
+                    "job_id": job_id,
+                    "extra_info": res.get("extra_info"),
+                }
+                return convert_to_content(result, job_info={
                     "trace_id": trace_id,
                     "executor_type": executor_type,
                     "job_id": job_id,

dp/agent/server/executor/local_executor.py

@@ -1,19 +1,43 @@
 import asyncio
 import importlib
 import inspect
+import io
 import jsonpickle
 import os
 import psutil
+import re
 import sys
+import time
 import uuid
 from multiprocessing import Process
 from typing import Dict, Optional
 
 from .base_executor import BaseExecutor
+from ..utils import get_logger
 
+DFLOW_ID_PATTERN = r"Workflow has been submitted \(ID: ([^,]*), UID: ([^)]*)\)"
+DFLOW_LINK_PATTERN = r"Workflow link: (.*)"
+logger = get_logger(__name__)
 
-def wrapped_fn(fn, kwargs):
+
+class Tee(io.TextIOBase):
+    def __init__(self, file, stdout):
+        self.file = file
+        self.stdout = stdout
+
+    def write(self, text):
+        self.stdout.write(text)
+        self.file.write(text)
+        self.file.flush()
+        return len(text)
+
+
+def wrapped_fn(fn, kwargs, redirect_file=None):
     pid = os.getpid()
+    if redirect_file:
+        stdout = sys.stdout
+        flog = open(redirect_file, "w")
+        sys.stdout = Tee(flog, stdout)
     try:
         if inspect.iscoroutinefunction(fn):
             result = asyncio.run(fn(**kwargs))
@@ -23,6 +47,10 @@ def wrapped_fn(fn, kwargs):
         with open("err", "w") as f:
             f.write(str(e))
         raise e
+    finally:
+        if redirect_file:
+            sys.stdout = stdout
+            flog.close()
     with open("%s.txt" % pid, "w") as f:
         f.write(jsonpickle.dumps(result))
 
@@ -47,13 +75,16 @@ def reload_dflow_config():
 
 
 class LocalExecutor(BaseExecutor):
-    def __init__(self, env: Optional[Dict[str, str]] = None):
+    def __init__(self, env: Optional[Dict[str, str]] = None,
+                 dflow: bool = False):
         """
         Execute the tool locally
         Args:
             env: The environmental variables at run time
+            dflow: Wait until workflow submitted in submit method
         """
         self.env = env or {}
+        self.dflow = dflow
 
     def set_env(self):
         old_env = {}
@@ -73,10 +104,34 @@ class LocalExecutor(BaseExecutor):
     def submit(self, fn, kwargs):
         os.environ["DP_AGENT_RUNNING_MODE"] = "1"
         old_env = self.set_env()
-        p = Process(target=wrapped_fn, kwargs={"fn": fn, "kwargs": kwargs})
+        params = {"fn": fn, "kwargs": kwargs}
+        if self.dflow:
+            params["redirect_file"] = "log.txt"
+        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:
+                        log = f.read()
+                    match_id = re.search(DFLOW_ID_PATTERN, log)
+                    match_link = re.search(DFLOW_LINK_PATTERN, log)
+                    if match_id and match_link:
+                        wf_id = match_id.group(1)
+                        wf_uid = match_id.group(2)
+                        wf_link = match_link.group(1)
+                        extra_info["workflow_id"] = wf_id
+                        extra_info["workflow_uid"] = wf_uid
+                        extra_info["workflow_link"] = wf_link
+                        break
+                if not alive:
+                    break
+                logger.info("Waiting workflow to be submitted")
+                time.sleep(1)
         self.recover_env(old_env)
-        return {"job_id": str(p.pid)}
+        return {"job_id": str(p.pid), "extra_info": extra_info}
 
     def query_status(self, job_id):
         try: