bohr-agent-sdk 0.1.117__py3-none-any.whl → 0.1.119__py3-none-any.whl

{bohr_agent_sdk-0.1.117.dist-info → bohr_agent_sdk-0.1.119.dist-info}/METADATA

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

{bohr_agent_sdk-0.1.117.dist-info → bohr_agent_sdk-0.1.119.dist-info}/RECORD

@@ -17,7 +17,7 @@ dp/agent/cli/templates/device/tescan_device.py.template,sha256=c7RIwWhfwp5WiPy1S
 dp/agent/cli/templates/ui/__init__.py,sha256=QgYZneX7gKyajZIEvbHIEMD3av3bKj64E_eeDPBeuvY,14
 dp/agent/cli/templates/ui/test_download.py,sha256=RzZ59pVd-4fYIbUbW6ZO1b1Twbycu5G6yb5yLFuT0oI,3002
 dp/agent/cli/templates/ui/ui_utils.py,sha256=OMgQxZ-UvIJB_a58YBMypQKTI7N9B3AQBOR7kOTudRE,9774
-dp/agent/cli/templates/ui/websocket-server.py,sha256=YLVnXyBelrZLo9EcU-17_qts4-UVscMM8-cX08X6XYY,2809
+dp/agent/cli/templates/ui/websocket-server.py,sha256=vhLsDuhQFxqp9RL7lWq3M8CP196VDZuoI3y_5YPqXQA,2968
 dp/agent/cli/templates/ui/api/__init__.py,sha256=eq1djuBJYvWlGTgxXvT9A386Eu21FNCcB7gV2lU46dU,8
 dp/agent/cli/templates/ui/api/config.py,sha256=ZhHApcXDrmO1JjI-Ok09ETYB-BVis390ZrMRt140DhQ,1018
 dp/agent/cli/templates/ui/api/constants.py,sha256=mHY81zf4ksfKEOsSiVsWD_A_KcQ6wveLvAImCay_kVM,1391
@@ -47,7 +47,7 @@ dp/agent/cli/templates/ui/server/connection.py,sha256=DVlhx49gMjKVPxOpi7spAZiHZA
 dp/agent/cli/templates/ui/server/file_watcher.py,sha256=t0e7IMgf4E1y0CxeboCqsVeZPBo1CGSw_qoO8wG6Wy4,2923
 dp/agent/cli/templates/ui/server/middleware.py,sha256=bNR5aFqKjxv2einQ0e26aq6tiUW4SMses3IHkrmHDCk,1474
 dp/agent/cli/templates/ui/server/models.py,sha256=obfMkFFB9qZZXDdwoC25u3tBlUC6Rz7sjeyuh7a_UWk,1680
-dp/agent/cli/templates/ui/server/session_manager.py,sha256=ZbNHGCFvswa-LKWn6c6RMHWDdfeOvY5L6A8CfrsF0OE,46536
+dp/agent/cli/templates/ui/server/session_manager.py,sha256=GviOjYZtdbd4Cpdxc0zfb1NV38uBvyJnzgd9uUusb9k,46528
 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
@@ -62,7 +62,7 @@ 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=VI1TI51QNTZEY9SPdlFFC0PFBcL8zPzL0PV6QTZNPOo,19167
+dp/agent/server/calculation_mcp_server.py,sha256=Nq1lv-H1E8J1yNz85TCbl1GtSQXQGZfAIB1qKajZcNY,19664
 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
@@ -75,8 +75,8 @@ dp/agent/server/storage/bohrium_storage.py,sha256=EsKX4dWWvZTn2TEhZv4zsvihfDK0mm
 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.117.dist-info/METADATA,sha256=kyyDZnKLQQFKuOF5ZV0ftTzlEp02fPnZhH06MmUzYyE,11070
-bohr_agent_sdk-0.1.117.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.117.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.117.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.117.dist-info/RECORD,,
+bohr_agent_sdk-0.1.119.dist-info/METADATA,sha256=oG4bst9eor1ytVutI5UGppqY8aJ-C9c12jKOGBoHh2g,11070
+bohr_agent_sdk-0.1.119.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.119.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.119.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.119.dist-info/RECORD,,

dp/agent/cli/templates/ui/server/session_manager.py

@@ -632,7 +632,7 @@ class SessionManager:
     ):
         """Process message stream - using ADK native event handling"""
         streaming_text = ""  # Accumulate streaming text
-        
+
         # Run Runner
         async for event in runner.run_async(
             new_message=content,

dp/agent/cli/templates/ui/websocket-server.py

@@ -48,11 +48,13 @@ if __name__ == "__main__":
     
     # uvicorn 始终监听 0.0.0.0 以支持所有配置的主机
     uvicorn.run(
-        app, 
-        host="0.0.0.0", 
+        app,
+        host="0.0.0.0",
         port=port,
         log_level="info",  # 使用 info 级别，过滤掉 warning
         access_log=False,  # 禁用访问日志，减少噪音
+        ws_ping_interval=20,  # Send WebSocket ping every 20s to keep connection alive
+        ws_ping_timeout=30,   # Disconnect if no pong received within 30s
         # 添加自定义的日志配置
         log_config={
             "version": 1,

dp/agent/server/calculation_mcp_server.py

@@ -26,6 +26,7 @@ from .executor import executor_dict
 from .storage import storage_dict
 from .utils import get_logger, JobResult, Tool
 logger = get_logger(__name__)
+CALCULATION_MCP_WORKDIR = os.getenv("CALCULATION_MCP_WORKDIR", os.getcwd())
 
 
 def parse_uri(uri):
@@ -81,7 +82,8 @@ def query_job_status(job_id: str, executor: Optional[dict] = None
         status (str): One of "Running", "Succeeded" or "Failed"
     """
     trace_id, exec_id = job_id.split("/")
-    with set_directory(trace_id):
+    workdir = os.path.join(CALCULATION_MCP_WORKDIR, trace_id)
+    with set_directory(workdir):
         executor = load_job_info()["executor"] or executor
         _, executor = init_executor(executor)
         status = executor.query_status(exec_id)
@@ -96,7 +98,8 @@ def terminate_job(job_id: str, executor: Optional[dict] = None):
         job_id (str): The ID of the calculation job
     """
     trace_id, exec_id = job_id.split("/")
-    with set_directory(trace_id):
+    workdir = os.path.join(CALCULATION_MCP_WORKDIR, trace_id)
+    with set_directory(workdir):
         executor = load_job_info()["executor"] or executor
         _, executor = init_executor(executor)
         executor.terminate(exec_id)
@@ -130,13 +133,15 @@ def handle_input_artifacts(fn, kwargs, storage):
                 kwargs.get(name) is not None):
             uris = kwargs[name]
             new_paths = []
-            for uri in uris:
+            for i, uri in enumerate(uris):
                 scheme, key = parse_uri(uri)
                 if scheme == storage_type:
                     s = storage
                 else:
                     s = storage_dict[scheme]()
-                path = s.download(key, "inputs/%s" % name)
+                dest_dir = Path("inputs") / name / f"item_{i:03d}"
+                dest_dir.mkdir(parents=True, exist_ok=True)
+                path = s.download(key, str(dest_dir))
                 new_paths.append(Path(path))
                 logger.info("Artifact %s downloaded to %s" % (
                     uri, path))
@@ -172,13 +177,15 @@ def handle_input_artifacts(fn, kwargs, storage):
             new_paths_dict = {}
             for key_name, uris in uris_dict.items():
                 new_paths = []
-                for uri in uris:
+                for i, uri in enumerate(uris):
                     scheme, key = parse_uri(uri)
                     if scheme == storage_type:
                         s = storage
                     else:
                         s = storage_dict[scheme]()
-                    path = s.download(key, f"inputs/{name}/{key_name}")
+                    dest_dir = Path("inputs") / name / key_name / f"item_{i:03d}"
+                    dest_dir.mkdir(parents=True, exist_ok=True)
+                    path = s.download(key, str(dest_dir))
                     new_paths.append(Path(path))
                     logger.info("Artifact %s (key=%s) downloaded to %s" % (
                         uri, key_name, path))
@@ -237,7 +244,8 @@ def get_job_results(job_id: str, executor: Optional[dict] = None,
         results (Any): results of the calculation job
     """
     trace_id, exec_id = job_id.split("/")
-    with set_directory(trace_id):
+    workdir = os.path.join(CALCULATION_MCP_WORKDIR, trace_id)
+    with set_directory(workdir):
         job_info = load_job_info()
         executor = job_info["executor"] or executor
         storage = job_info["storage"] or storage
@@ -383,17 +391,14 @@ 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)
-                if preprocess_func is not None:
-                    executor, storage, kwargs = preprocess_func(
-                        executor, storage, kwargs)
-                executor_type, executor = init_executor(executor)
-                if create_workdir is False or (
-                    create_workdir is None and inspect.iscoroutinefunction(fn)
-                        and executor_type == "local"):
+                if create_workdir is False:
                     workdir = "."
                 else:
-                    workdir = trace_id
+                    workdir = os.path.join(CALCULATION_MCP_WORKDIR, trace_id)
                 with set_directory(workdir):
+                    if preprocess_func is not None:
+                        executor, storage, kwargs = preprocess_func(
+                            executor, storage, kwargs)
                     job = {
                         "tool_name": fn.__name__,
                         "executor": executor,
@@ -403,6 +408,7 @@ class CalculationMCPServer:
                         json.dump(job, f, indent=4)
                     kwargs, input_artifacts = handle_input_artifacts(
                         fn, kwargs, storage)
+                    executor_type, executor = init_executor(executor)
                     res = executor.submit(fn, kwargs)
                     exec_id = res["job_id"]
                     job_id = "%s/%s" % (workdir, exec_id)
@@ -433,7 +439,7 @@ class CalculationMCPServer:
                         and executor_type == "local"):
                     workdir = "."
                 else:
-                    workdir = trace_id
+                    workdir = os.path.join(CALCULATION_MCP_WORKDIR, trace_id)
                 with set_directory(workdir):
                     kwargs, input_artifacts = handle_input_artifacts(
                         fn, kwargs, storage)