bohr-agent-sdk 0.1.19__tar.gz → 0.1.21__tar.gz

{bohr_agent_sdk-0.1.19 → bohr_agent_sdk-0.1.21}/PKG-INFO

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

{bohr_agent_sdk-0.1.19 → bohr_agent_sdk-0.1.21}/setup.py

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

{bohr_agent_sdk-0.1.19 → bohr_agent_sdk-0.1.21}/src/bohr_agent_sdk.egg-info/PKG-INFO

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

{bohr_agent_sdk-0.1.19 → bohr_agent_sdk-0.1.21}/src/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,

{bohr_agent_sdk-0.1.19 → bohr_agent_sdk-0.1.21}/src/dp/agent/server/executor/dispatcher_executor.py

@@ -12,12 +12,15 @@ from dpdispatcher import Machine, Resources, Task, Submission
 
 from .base_executor import BaseExecutor
 from .... import __path__
+
 config = {
     "username": os.environ.get("BOHRIUM_USERNAME", ""),
     "password": os.environ.get("BOHRIUM_PASSWORD", ""),
     "project_id": os.environ.get("BOHRIUM_PROJECT_ID", ""),
     "access_key": os.environ.get("BOHRIUM_ACCESS_KEY", ""),
     "app_key": os.environ.get("BOHRIUM_APP_KEY", "agent"),
+    "bohrium_url": os.environ.get("BOHRIUM_BOHRIUM_URL",
+                                  "https://bohrium.dp.tech"),
 }
 logger = logging.getLogger(__name__)
 
@@ -26,17 +29,17 @@ def get_source_code(fn):
     source_lines, start_line = inspect.getsourcelines(fn)
     source_file = inspect.getsourcefile(fn)
     with open(source_file, "r", encoding="utf-8") as fd:
-        pre_lines = fd.readlines()[:start_line-1]
+        pre_lines = fd.readlines()[:start_line - 1]
     return "".join(pre_lines + source_lines) + "\n"
 
 
 class DispatcherExecutor(BaseExecutor):
     def __init__(
-        self,
-        machine=None,
-        resources=None,
-        python_packages=None,
-        python_executable="python3",
+            self,
+            machine=None,
+            resources=None,
+            python_packages=None,
+            python_executable="python3",
     ):
         """Use DPDispatcher to execute the tool
         Refer to https://docs.deepmodeling.com/projects/dpdispatcher.
@@ -107,7 +110,7 @@ class DispatcherExecutor(BaseExecutor):
         script += "if __name__ == \"__main__\":\n"
         script += "    cwd = os.getcwd()\n"
         script += "    kwargs = jsonpickle.loads(r'''%s''')\n" % \
-            jsonpickle.dumps(kwargs)
+                  jsonpickle.dumps(kwargs)
         script += "    try:\n"
         if import_func_line is not None:
             script += "        " + import_func_line
@@ -118,8 +121,8 @@ class DispatcherExecutor(BaseExecutor):
             script += "        if isinstance(results, dict):\n"
             script += "            for name in results:\n"
             script += "                if isinstance(results[name], Path):\n"
-            script += "                    results[name] = "\
-                "results[name].absolute().relative_to(cwd)\n"
+            script += "                    results[name] = " \
+                      "results[name].absolute().relative_to(cwd)\n"
         script += "    except Exception as e:\n"
         script += "        os.chdir(cwd)\n"
         script += "        with open('err', 'w') as f:\n"
@@ -173,8 +176,8 @@ class DispatcherExecutor(BaseExecutor):
             extra_info = {
                 "bohr_job_id": bohr_job_id,
                 "bohr_group_id": bohr_group_id,
-                "job_link": "https://bohrium.dp.tech/jobs/detail/%s" %
-                bohr_job_id,
+                "job_link": f"{config['bohrium_url']}/jobs/detail/%s" %
+                            bohr_job_id,
             }
             logger.info(extra_info)
             res["extra_info"] = extra_info
@@ -182,8 +185,8 @@ class DispatcherExecutor(BaseExecutor):
             bohr_job_id = submission.belonging_jobs[0].job_id
             extra_info = {
                 "bohr_job_id": bohr_job_id,
-                "job_link": "https://bohrium.dp.tech/jobs/detail/%s" %
-                bohr_job_id,
+                "job_link": f"{config['bohrium_url']}/jobs/detail/%s" %
+                            bohr_job_id,
             }
             logger.info(extra_info)
             res["extra_info"] = extra_info

{bohr_agent_sdk-0.1.19 → bohr_agent_sdk-0.1.21}/src/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: