bohr-agent-sdk 0.1.122__tar.gz → 0.1.124__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (90) hide show
  1. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/PKG-INFO +1 -1
  2. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/setup.py +1 -1
  3. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/bohr_agent_sdk.egg-info/PKG-INFO +1 -1
  4. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/bohr_agent_sdk.egg-info/SOURCES.txt +1 -0
  5. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/calculation_mcp_server.py +180 -94
  6. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/executor/base_executor.py +17 -14
  7. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/executor/dispatcher_executor.py +3 -0
  8. bohr_agent_sdk-0.1.124/tests/test_calculation_mcp_server.py +357 -0
  9. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/README.md +0 -0
  10. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/pyproject.toml +0 -0
  11. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/setup.cfg +0 -0
  12. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/bohr_agent_sdk.egg-info/dependency_links.txt +0 -0
  13. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/bohr_agent_sdk.egg-info/entry_points.txt +0 -0
  14. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/bohr_agent_sdk.egg-info/requires.txt +0 -0
  15. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/bohr_agent_sdk.egg-info/top_level.txt +0 -0
  16. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/__init__.py +0 -0
  17. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/__init__.py +0 -0
  18. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/adapter/adk/__init__.py +0 -0
  19. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/adapter/adk/client/__init__.py +0 -0
  20. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/adapter/adk/client/calculation_mcp_tool.py +0 -0
  21. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/adapter/adk/storage_artifact_service.py +0 -0
  22. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/adapter/adk/utils.py +0 -0
  23. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/adapter/camel/__init__.py +0 -0
  24. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/adapter/camel/client/__init__.py +0 -0
  25. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/adapter/camel/client/calculation_mcp_client.py +0 -0
  26. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/__init__.py +0 -0
  27. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/cli.py +0 -0
  28. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/__init__.py +0 -0
  29. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/calculation/simple.py.template +0 -0
  30. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/device/tescan_device.py.template +0 -0
  31. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/main.py.template +0 -0
  32. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/__init__.py +0 -0
  33. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/__init__.py +0 -0
  34. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/config.py +0 -0
  35. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/constants.py +0 -0
  36. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/debug.py +0 -0
  37. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/files.py +0 -0
  38. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/files_upload.py +0 -0
  39. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/files_user.py +0 -0
  40. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/messages.py +0 -0
  41. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/projects.py +0 -0
  42. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/sessions.py +0 -0
  43. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/utils.py +0 -0
  44. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/api/websocket.py +0 -0
  45. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/config/__init__.py +0 -0
  46. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/config/agent_config.py +0 -0
  47. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/frontend/index.html +0 -0
  48. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/frontend/package.json +0 -0
  49. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/frontend/tsconfig.json +0 -0
  50. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/frontend/tsconfig.node.json +0 -0
  51. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/frontend/ui-static/assets/index-DdAmKhul.js +0 -0
  52. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/frontend/ui-static/assets/index-DfN2raU9.css +0 -0
  53. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/frontend/ui-static/index.html +0 -0
  54. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/frontend/vite.config.ts +0 -0
  55. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/scripts/build_ui.py +0 -0
  56. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/server/__init__.py +0 -0
  57. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/server/app.py +0 -0
  58. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/server/connection.py +0 -0
  59. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/server/file_watcher.py +0 -0
  60. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/server/middleware.py +0 -0
  61. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/server/models.py +0 -0
  62. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/server/session_manager.py +0 -0
  63. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/server/user_files.py +0 -0
  64. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/server/utils.py +0 -0
  65. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/test_download.py +0 -0
  66. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/ui_utils.py +0 -0
  67. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cli/templates/ui/websocket-server.py +0 -0
  68. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/client/__init__.py +0 -0
  69. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/client/mcp_client.py +0 -0
  70. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cloud/__init__.py +0 -0
  71. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cloud/main.py +0 -0
  72. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cloud/mcp.py +0 -0
  73. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/cloud/mqtt.py +0 -0
  74. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/device/__init__.py +0 -0
  75. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/device/device/__init__.py +0 -0
  76. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/device/device/device.py +0 -0
  77. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/device/device/types.py +0 -0
  78. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/device/mqtt_device_twin.py +0 -0
  79. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/__init__.py +0 -0
  80. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/executor/__init__.py +0 -0
  81. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/executor/local_executor.py +0 -0
  82. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/preprocessor.py +0 -0
  83. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/storage/__init__.py +0 -0
  84. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/storage/base_storage.py +0 -0
  85. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/storage/bohrium_storage.py +0 -0
  86. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/storage/http_storage.py +0 -0
  87. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/storage/local_storage.py +0 -0
  88. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/storage/oss_storage.py +0 -0
  89. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/src/dp/agent/server/utils.py +0 -0
  90. {bohr_agent_sdk-0.1.122 → bohr_agent_sdk-0.1.124}/tests/test_cli.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: bohr-agent-sdk
3
- Version: 0.1.122
3
+ Version: 0.1.124
4
4
  Summary: SDK for scientific agents
5
5
  Home-page: https://github.com/dptech-corp/bohr-agent-sdk/
6
6
  Author: DP Technology
@@ -9,7 +9,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
9
9
 
10
10
  setup(
11
11
  name="bohr-agent-sdk",
12
- version="0.1.122",
12
+ version="0.1.124",
13
13
  description="SDK for science agent and mcp tools",
14
14
  long_description=long_description,
15
15
  long_description_content_type="text/markdown",
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: bohr-agent-sdk
3
- Version: 0.1.122
3
+ Version: 0.1.124
4
4
  Summary: SDK for scientific agents
5
5
  Home-page: https://github.com/dptech-corp/bohr-agent-sdk/
6
6
  Author: DP Technology
@@ -84,4 +84,5 @@ src/dp/agent/server/storage/bohrium_storage.py
84
84
  src/dp/agent/server/storage/http_storage.py
85
85
  src/dp/agent/server/storage/local_storage.py
86
86
  src/dp/agent/server/storage/oss_storage.py
87
+ tests/test_calculation_mcp_server.py
87
88
  tests/test_cli.py
@@ -7,7 +7,7 @@ from copy import deepcopy
7
7
  from datetime import datetime
8
8
  from pathlib import Path
9
9
  from urllib.parse import urlparse
10
- from typing import Annotated, Literal, Optional, List, Dict
10
+ from typing import Annotated, Literal, Optional, List, Dict, Union, Any, get_origin, get_args
11
11
 
12
12
  from mcp.server.fastmcp import FastMCP
13
13
  from mcp.server.fastmcp.utilities.context_injection import (
@@ -30,12 +30,18 @@ CALCULATION_MCP_WORKDIR = os.getenv("CALCULATION_MCP_WORKDIR", os.getcwd())
30
30
 
31
31
 
32
32
  def parse_uri(uri):
33
- scheme = urlparse(uri).scheme
33
+ parsed = urlparse(uri)
34
+ scheme = parsed.scheme
34
35
  if scheme == "":
35
36
  key = uri
36
37
  scheme = "local"
37
38
  else:
38
- key = uri[len(scheme)+3:]
39
+ if parsed.netloc:
40
+ key = parsed.netloc + parsed.path
41
+ else:
42
+ key = parsed.path
43
+ if parsed.query:
44
+ key += "?" + parsed.query
39
45
  return scheme, key
40
46
 
41
47
 
@@ -105,96 +111,128 @@ def terminate_job(job_id: str, executor: Optional[dict] = None):
105
111
  logger.info("Job %s is terminated" % job_id)
106
112
 
107
113
 
114
+ def _normalize_annotation(ann):
115
+ if ann is None:
116
+ return None
117
+ origin = get_origin(ann)
118
+ if origin is Annotated:
119
+ return _normalize_annotation(get_args(ann)[0])
120
+ if origin is Union:
121
+ args = get_args(ann)
122
+ if type(None) in args:
123
+ non_none = [a for a in args if a is not type(None)]
124
+ if non_none:
125
+ return _normalize_annotation(non_none[0])
126
+ return ann
127
+
128
+
129
+ def _download_artifact(uri, storage, storage_type, input_artifacts,
130
+ input_name, path_trace):
131
+ scheme, key = parse_uri(uri)
132
+ if scheme == storage_type:
133
+ s = storage
134
+ else:
135
+ s = storage_dict[scheme]()
136
+ rel = "/".join(str(p) for p in path_trace) if path_trace else ""
137
+ download_dir = os.path.join("inputs", input_name, rel) if rel else os.path.join("inputs", input_name)
138
+ os.makedirs(download_dir, exist_ok=True)
139
+ path = s.download(key, download_dir)
140
+ logger.info("Artifact %s downloaded to %s" % (uri, path))
141
+ if input_name not in input_artifacts:
142
+ input_artifacts[input_name] = {"storage_type": scheme, "uri": []}
143
+ if isinstance(input_artifacts[input_name].get("uri"), list):
144
+ input_artifacts[input_name]["uri"].append(uri)
145
+ return path
146
+
147
+
148
+ def _traverse_and_process(value, annotation, storage_type, storage,
149
+ input_artifacts, input_name, path_trace=None):
150
+ if path_trace is None:
151
+ path_trace = []
152
+ ann = _normalize_annotation(annotation)
153
+ if ann is None:
154
+ return value
155
+ origin = get_origin(ann)
156
+ args = get_args(ann)
157
+
158
+ # Path: only when resolved annotation is Path (Optional[Path] is normalized to Path)
159
+ if ann is Path:
160
+ s = str(value)
161
+ if not s:
162
+ return Path(".")
163
+ parsed = urlparse(s)
164
+ if parsed.scheme and len(parsed.scheme) > 1:
165
+ return Path(_download_artifact(
166
+ s, storage, storage_type, input_artifacts, input_name, path_trace))
167
+ return Path(value)
168
+
169
+ # BaseModel: schema-driven traversal over model fields (check before dict so nesting works)
170
+ if isinstance(ann, type) and issubclass(ann, BaseModel):
171
+ # Convert to dict for traversal; re-instantiate to model at the end so callers get objects
172
+ if isinstance(value, BaseModel):
173
+ value = value.model_dump()
174
+ if not isinstance(value, dict):
175
+ return value
176
+ out = dict(value)
177
+ for field_name, field_info in ann.model_fields.items():
178
+ if field_name in out and out[field_name] is not None:
179
+ out[field_name] = _traverse_and_process(
180
+ out[field_name],
181
+ field_info.annotation,
182
+ storage_type,
183
+ storage,
184
+ input_artifacts,
185
+ input_name,
186
+ path_trace + [field_name],
187
+ )
188
+ # Re-instantiate so tool functions receive model instances (dot notation works)
189
+ try:
190
+ return ann.model_validate(out)
191
+ except Exception as e:
192
+ logger.warning("Failed to re-instantiate model %s: %s", ann.__name__, e)
193
+ return out
194
+
195
+ # List: use inner type from type args
196
+ if origin in (list, List) and isinstance(value, (list, tuple)):
197
+ inner = _normalize_annotation(args[0]) if args else Any
198
+ return [
199
+ _traverse_and_process(
200
+ item, inner, storage_type, storage,
201
+ input_artifacts, input_name, path_trace + [i])
202
+ for i, item in enumerate(value)
203
+ ]
204
+
205
+ # Dict: use value type from type args (e.g. Dict[str, Path] processes values as Path)
206
+ if origin in (dict, Dict) and isinstance(value, dict):
207
+ value_type = _normalize_annotation(args[1]) if (args and len(args) > 1) else Any
208
+ return {
209
+ k: _traverse_and_process(
210
+ v, value_type, storage_type, storage,
211
+ input_artifacts, input_name, path_trace + [k])
212
+ for k, v in value.items()
213
+ }
214
+
215
+ return value
216
+
217
+
108
218
  def handle_input_artifacts(fn, kwargs, storage):
109
219
  storage_type, storage = init_storage(storage)
110
- sig = inspect.signature(fn)
220
+ sig = inspect.signature(fn, eval_str=True)
111
221
  input_artifacts = {}
222
+ new_kwargs = {}
112
223
  for name, param in sig.parameters.items():
113
- if param.annotation is Path or (
114
- param.annotation is Optional[Path] and
115
- kwargs.get(name) is not None):
116
- uri = kwargs[name]
117
- scheme, key = parse_uri(uri)
118
- if scheme == storage_type:
119
- s = storage
120
- else:
121
- s = storage_dict[scheme]()
122
- path = s.download(key, "inputs/%s" % name)
123
- logger.info("Artifact %s downloaded to %s" % (
124
- uri, path))
125
- kwargs[name] = Path(path)
126
- input_artifacts[name] = {
127
- "storage_type": scheme,
128
- "uri": uri,
129
- }
130
- elif param.annotation is List[Path] or (
131
- param.annotation is Optional[List[Path]] and
132
- kwargs.get(name) is not None):
133
- uris = kwargs[name]
134
- new_paths = []
135
- for i, uri in enumerate(uris):
136
- scheme, key = parse_uri(uri)
137
- if scheme == storage_type:
138
- s = storage
139
- else:
140
- s = storage_dict[scheme]()
141
- dest_dir = Path("inputs") / name / f"item_{i:03d}"
142
- dest_dir.mkdir(parents=True, exist_ok=True)
143
- path = s.download(key, str(dest_dir))
144
- new_paths.append(Path(path))
145
- logger.info("Artifact %s downloaded to %s" % (
146
- uri, path))
147
- kwargs[name] = new_paths
148
- input_artifacts[name] = {
149
- "storage_type": storage_type,
150
- "uri": uris,
151
- }
152
- elif param.annotation is Dict[str, Path] or (
153
- param.annotation is Optional[Dict[str, Path]] and
154
- kwargs.get(name) is not None):
155
- uris_dict = kwargs[name]
156
- new_paths_dict = {}
157
- for key_name, uri in uris_dict.items():
158
- scheme, key = parse_uri(uri)
159
- if scheme == storage_type:
160
- s = storage
161
- else:
162
- s = storage_dict[scheme]()
163
- path = s.download(key, f"inputs/{name}/{key_name}")
164
- new_paths_dict[key_name] = Path(path)
165
- logger.info("Artifact %s (key=%s) downloaded to %s" % (
166
- uri, key_name, path))
167
- kwargs[name] = new_paths_dict
168
- input_artifacts[name] = {
169
- "storage_type": storage_type,
170
- "uri": uris_dict,
171
- }
172
- elif param.annotation is Dict[str, List[Path]] or (
173
- param.annotation is Optional[Dict[str, List[Path]]] and
174
- kwargs.get(name) is not None):
175
- uris_dict = kwargs[name]
176
- new_paths_dict = {}
177
- for key_name, uris in uris_dict.items():
178
- new_paths = []
179
- for i, uri in enumerate(uris):
180
- scheme, key = parse_uri(uri)
181
- if scheme == storage_type:
182
- s = storage
183
- else:
184
- s = storage_dict[scheme]()
185
- dest_dir = Path("inputs") / name / key_name / f"item_{i:03d}"
186
- dest_dir.mkdir(parents=True, exist_ok=True)
187
- path = s.download(key, str(dest_dir))
188
- new_paths.append(Path(path))
189
- logger.info("Artifact %s (key=%s) downloaded to %s" % (
190
- uri, key_name, path))
191
- new_paths_dict[key_name] = new_paths
192
- kwargs[name] = new_paths_dict
193
- input_artifacts[name] = {
194
- "storage_type": storage_type,
195
- "uri": uris_dict,
196
- }
197
- return kwargs, input_artifacts
224
+ if name not in kwargs:
225
+ if param.default is not inspect.Parameter.empty:
226
+ new_kwargs[name] = param.default
227
+ continue
228
+ val = kwargs[name]
229
+ if val is None and _normalize_annotation(param.annotation) != param.annotation:
230
+ new_kwargs[name] = val
231
+ continue
232
+ new_kwargs[name] = _traverse_and_process(
233
+ val, param.annotation, storage_type, storage,
234
+ input_artifacts, name)
235
+ return new_kwargs, input_artifacts
198
236
 
199
237
 
200
238
  def handle_output_artifacts(results, exec_id, storage):
@@ -269,6 +307,57 @@ annotation_map = {
269
307
  }
270
308
 
271
309
 
310
+ # Cache for schema models derived from BaseModel (Path -> str in JSON schema)
311
+ _schema_model_cache: Dict[type, type] = {}
312
+
313
+ def get_schema_annotation(annotation: Any) -> Any:
314
+ """
315
+ Map an annotation to the type used in JSON schema (e.g. Path -> str).
316
+ For BaseModel, build a schema model with the same structure but Path fields as str.
317
+ Handles List[...], Dict[...], Optional[...] and nested BaseModel recursively.
318
+ """
319
+ if annotation is None or annotation is type(None):
320
+ return annotation
321
+ origin = get_origin(annotation)
322
+ if origin is Annotated:
323
+ return get_schema_annotation(get_args(annotation)[0])
324
+ if origin is Union:
325
+ args = get_args(annotation)
326
+ if type(None) in args:
327
+ non_none = [a for a in args if a is not type(None)]
328
+ if len(non_none) == 1:
329
+ return Optional[get_schema_annotation(non_none[0])]
330
+ if annotation in annotation_map:
331
+ return annotation_map[annotation]
332
+ # List[X] -> List[schema(X)] so e.g. List[BaseModel] becomes List[BaseModelSchema]
333
+ if origin in (list, List):
334
+ type_args = get_args(annotation)
335
+ inner = get_schema_annotation(type_args[0]) if type_args else Any
336
+ return List[inner]
337
+ # Dict[K, V] -> Dict[K, schema(V)]
338
+ if origin in (dict, Dict):
339
+ type_args = get_args(annotation)
340
+ if type_args and len(type_args) >= 2:
341
+ return Dict[type_args[0], get_schema_annotation(type_args[1])]
342
+ return annotation
343
+ if isinstance(annotation, type) and issubclass(annotation, BaseModel):
344
+ if annotation not in _schema_model_cache:
345
+ schema_fields = {}
346
+ for name, field_info in annotation.model_fields.items():
347
+ fa = get_schema_annotation(field_info.annotation)
348
+ if field_info.is_required():
349
+ schema_fields[name] = (fa, Field())
350
+ else:
351
+ default = getattr(field_info, "default", None)
352
+ schema_fields[name] = (fa, Field(default=default))
353
+ _schema_model_cache[annotation] = create_model(
354
+ f"{annotation.__name__}Schema",
355
+ **schema_fields,
356
+ )
357
+ return _schema_model_cache[annotation]
358
+ return annotation
359
+
360
+
272
361
  class SubmitResult(BaseModel):
273
362
  job_id: str
274
363
  extra_info: dict | None = None
@@ -329,11 +418,8 @@ class CalculationMCPServer:
329
418
  for n, annotation in \
330
419
  func_arg_metadata.arg_model.__annotations__.items():
331
420
  param = params[n]
332
- if param.annotation in annotation_map:
333
- model_params[n] = Annotated[
334
- (annotation_map[param.annotation], Field())]
335
- else:
336
- model_params[n] = annotation
421
+ schema_ann = get_schema_annotation(param.annotation)
422
+ model_params[n] = Annotated[(schema_ann, Field())]
337
423
  if param.default is not inspect.Parameter.empty:
338
424
  model_params[n] = (model_params[n], param.default)
339
425
  for n, param in inspect.signature(new_fn).parameters.items():
@@ -32,20 +32,23 @@ class BaseExecutor(ABC):
32
32
  def prune_context(self, kwargs: dict):
33
33
  for key, value in kwargs.items():
34
34
  if isinstance(value, Context):
35
- context = Context(request_context=RequestContext(
36
- request_id=value.request_context.request_id,
37
- meta=value.request_context.meta,
38
- session=None,
39
- lifespan_context=value.request_context.lifespan_context,
40
- request=Request(
41
- scope={
42
- k: v for k, v in
43
- value.request_context.request.scope.items()
44
- if k not in ["app", "router", "endpoint",
45
- "starlette.exception_handlers"]
46
- },
47
- )
48
- ))
35
+ try:
36
+ context = Context(request_context=RequestContext(
37
+ request_id=value.request_context.request_id,
38
+ meta=value.request_context.meta,
39
+ session=None,
40
+ lifespan_context=value.request_context.lifespan_context,
41
+ request=Request(
42
+ scope={
43
+ k: v for k, v in
44
+ value.request_context.request.scope.items()
45
+ if k not in ["app", "router", "endpoint",
46
+ "starlette.exception_handlers"]
47
+ },
48
+ )
49
+ ))
50
+ except Exception:
51
+ context = Context()
49
52
  kwargs[key] = context
50
53
  return kwargs
51
54
 
@@ -200,6 +200,9 @@ class DispatcherExecutor(BaseExecutor):
200
200
  if isinstance(value, Path):
201
201
  forward_files.append(str(value))
202
202
 
203
+ if os.path.isdir("inputs") and "inputs" not in forward_files:
204
+ forward_files.append("inputs")
205
+
203
206
  task = {
204
207
  "task_work_path": "./",
205
208
  "outlog": "log",
@@ -0,0 +1,357 @@
1
+ """
2
+ Unit tests for calculation_mcp_server: get_schema_annotation and _traverse_and_process.
3
+ """
4
+ from pathlib import Path
5
+ from typing import Annotated, Any, Dict, List, Optional
6
+ from unittest.mock import patch, MagicMock
7
+
8
+ import pytest
9
+ from pydantic import BaseModel
10
+
11
+ from dp.agent.server.calculation_mcp_server import (
12
+ get_schema_annotation,
13
+ _schema_model_cache,
14
+ _traverse_and_process,
15
+ )
16
+
17
+
18
+ def _clear_schema_cache():
19
+ _schema_model_cache.clear()
20
+
21
+
22
+ # --- get_schema_annotation ---
23
+
24
+
25
+ def test_get_schema_annotation_none():
26
+ assert get_schema_annotation(None) is None
27
+
28
+
29
+ def test_get_schema_annotation_type_none():
30
+ assert get_schema_annotation(type(None)) is type(None)
31
+
32
+
33
+ def test_get_schema_annotation_annotated_path():
34
+ ann = Annotated[Path, "path field"]
35
+ result = get_schema_annotation(ann)
36
+ assert result is str
37
+
38
+
39
+ def test_get_schema_annotation_annotated_nested():
40
+ ann = Annotated[Annotated[Optional[Path], "x"], "y"]
41
+ result = get_schema_annotation(ann)
42
+ assert result == Optional[str]
43
+
44
+
45
+ def test_get_schema_annotation_optional_path():
46
+ assert get_schema_annotation(Optional[Path]) == Optional[str]
47
+
48
+
49
+ def test_get_schema_annotation_optional_str():
50
+ assert get_schema_annotation(Optional[str]) == Optional[str]
51
+
52
+
53
+ def test_get_schema_annotation_path():
54
+ assert get_schema_annotation(Path) is str
55
+
56
+
57
+ def test_get_schema_annotation_list_path():
58
+ assert get_schema_annotation(List[Path]) == List[str]
59
+
60
+
61
+ def test_get_schema_annotation_optional_list_path():
62
+ assert get_schema_annotation(Optional[List[Path]]) == Optional[List[str]]
63
+
64
+
65
+ def test_get_schema_annotation_dict_str_path():
66
+ assert get_schema_annotation(Dict[str, Path]) == Dict[str, str]
67
+
68
+
69
+ def test_get_schema_annotation_optional_dict_str_path():
70
+ assert get_schema_annotation(Optional[Dict[str, Path]]) == Optional[Dict[str, str]]
71
+
72
+
73
+ def test_get_schema_annotation_dict_str_list_path():
74
+ assert get_schema_annotation(Dict[str, List[Path]]) == Dict[str, List[str]]
75
+
76
+
77
+ def test_get_schema_annotation_list_str():
78
+ assert get_schema_annotation(List[str]) == List[str]
79
+
80
+
81
+ def test_get_schema_annotation_list_any():
82
+ assert get_schema_annotation(List[Any]) == List[Any]
83
+
84
+
85
+ def test_get_schema_annotation_dict_value_schema():
86
+ assert get_schema_annotation(Dict[str, List[Path]]) == Dict[str, List[str]]
87
+
88
+
89
+ def test_get_schema_annotation_dict_no_args_unchanged():
90
+ assert get_schema_annotation(dict) is dict
91
+
92
+
93
+ def test_get_schema_annotation_basemodel_with_path_field():
94
+ class ModelWithPath(BaseModel):
95
+ name: str
96
+ data_path: Path
97
+
98
+ _clear_schema_cache()
99
+ schema_ann = get_schema_annotation(ModelWithPath)
100
+ assert schema_ann is not ModelWithPath
101
+ assert hasattr(schema_ann, "model_fields")
102
+ assert "name" in schema_ann.model_fields
103
+ assert "data_path" in schema_ann.model_fields
104
+ assert schema_ann.model_fields["data_path"].annotation is str
105
+ assert schema_ann.model_fields["name"].annotation is str
106
+
107
+
108
+ def test_get_schema_annotation_basemodel_cached():
109
+ class CachedModel(BaseModel):
110
+ x: Path
111
+
112
+ _clear_schema_cache()
113
+ first = get_schema_annotation(CachedModel)
114
+ second = get_schema_annotation(CachedModel)
115
+ assert first is second
116
+ assert CachedModel in _schema_model_cache
117
+
118
+
119
+ def test_get_schema_annotation_basemodel_nested():
120
+ class Inner(BaseModel):
121
+ path_field: Path
122
+
123
+ class Outer(BaseModel):
124
+ inner: Inner
125
+ top_path: Optional[Path] = None
126
+
127
+ _clear_schema_cache()
128
+ schema_ann = get_schema_annotation(Outer)
129
+ assert "inner" in schema_ann.model_fields
130
+ assert "top_path" in schema_ann.model_fields
131
+ inner_schema = schema_ann.model_fields["inner"].annotation
132
+ assert hasattr(inner_schema, "model_fields")
133
+ assert inner_schema.model_fields["path_field"].annotation is str
134
+ assert schema_ann.model_fields["top_path"].annotation == Optional[str]
135
+
136
+
137
+ def test_get_schema_annotation_plain_types():
138
+ assert get_schema_annotation(str) is str
139
+ assert get_schema_annotation(int) is int
140
+ assert get_schema_annotation(bool) is bool
141
+
142
+
143
+ # --- _traverse_and_process ---
144
+
145
+
146
+ def test_traverse_annotation_none_returns_value():
147
+ value = {"a": 1}
148
+ result = _traverse_and_process(
149
+ value, None, "local", MagicMock(), {}, "input_name"
150
+ )
151
+ assert result == value
152
+
153
+
154
+ def test_traverse_primitive_str_unchanged():
155
+ result = _traverse_and_process(
156
+ "hello", str, "local", MagicMock(), {}, "input_name"
157
+ )
158
+ assert result == "hello"
159
+
160
+
161
+ def test_traverse_primitive_int_unchanged():
162
+ result = _traverse_and_process(
163
+ 42, int, "local", MagicMock(), {}, "input_name"
164
+ )
165
+ assert result == 42
166
+
167
+
168
+ def test_traverse_path_local_string():
169
+ result = _traverse_and_process(
170
+ "/tmp/foo", Path, "local", MagicMock(), {}, "input_name"
171
+ )
172
+ assert result == Path("/tmp/foo")
173
+
174
+
175
+ def test_traverse_path_empty_string_returns_dot():
176
+ result = _traverse_and_process(
177
+ "", Path, "local", MagicMock(), {}, "input_name"
178
+ )
179
+ assert result == Path(".")
180
+
181
+
182
+ def test_traverse_path_uri_calls_download():
183
+ input_artifacts = {}
184
+ mock_storage = MagicMock()
185
+ with patch(
186
+ "dp.agent.server.calculation_mcp_server._download_artifact",
187
+ return_value="/downloaded/path",
188
+ ) as mock_download:
189
+ result = _traverse_and_process(
190
+ "local://bucket/key",
191
+ Path,
192
+ "local",
193
+ mock_storage,
194
+ input_artifacts,
195
+ "data",
196
+ )
197
+ assert result == Path("/downloaded/path")
198
+ mock_download.assert_called_once()
199
+ call_kw = mock_download.call_args
200
+ assert call_kw[0][0] == "local://bucket/key"
201
+ assert call_kw[0][-1] == []
202
+
203
+
204
+ def test_traverse_path_uri_appends_path_trace():
205
+ input_artifacts = {}
206
+ with patch(
207
+ "dp.agent.server.calculation_mcp_server._download_artifact",
208
+ return_value="/downloaded/path",
209
+ ) as mock_download:
210
+ _traverse_and_process(
211
+ "http://example.com/file",
212
+ Path,
213
+ "local",
214
+ MagicMock(),
215
+ input_artifacts,
216
+ "input_name",
217
+ path_trace=["config", "file"],
218
+ )
219
+ call_args = mock_download.call_args[0]
220
+ assert call_args[-1] == ["config", "file"]
221
+
222
+
223
+ def test_traverse_list_of_paths_local():
224
+ value = ["/a", "/b"]
225
+ result = _traverse_and_process(
226
+ value, List[Path], "local", MagicMock(), {}, "input_name"
227
+ )
228
+ assert result == [Path("/a"), Path("/b")]
229
+
230
+
231
+ def test_traverse_list_path_trace():
232
+ value = ["/a", "/b"]
233
+ result = _traverse_and_process(
234
+ value,
235
+ List[Path],
236
+ "local",
237
+ MagicMock(),
238
+ {},
239
+ "input_name",
240
+ path_trace=["items"],
241
+ )
242
+ assert result == [Path("/a"), Path("/b")]
243
+
244
+
245
+ def test_traverse_list_nested_basemodel():
246
+ class Item(BaseModel):
247
+ path: Path
248
+
249
+ value = [{"path": "/p1"}, {"path": "/p2"}]
250
+ result = _traverse_and_process(
251
+ value, List[Item], "local", MagicMock(), {}, "input_name"
252
+ )
253
+ assert len(result) == 2
254
+ assert result[0].path == Path("/p1")
255
+ assert result[1].path == Path("/p2")
256
+
257
+
258
+ def test_traverse_dict_str_path():
259
+ value = {"a": "/path/a", "b": "/path/b"}
260
+ result = _traverse_and_process(
261
+ value, Dict[str, Path], "local", MagicMock(), {}, "input_name"
262
+ )
263
+ assert result == {"a": Path("/path/a"), "b": Path("/path/b")}
264
+
265
+
266
+ def test_traverse_dict_path_trace():
267
+ value = {"x": "/x"}
268
+ result = _traverse_and_process(
269
+ value,
270
+ Dict[str, Path],
271
+ "local",
272
+ MagicMock(),
273
+ {},
274
+ "input_name",
275
+ path_trace=["config"],
276
+ )
277
+ assert result == {"x": Path("/x")}
278
+
279
+
280
+ def test_traverse_basemodel_from_dict():
281
+ class Model(BaseModel):
282
+ name: str
283
+ data_path: Path
284
+
285
+ value = {"name": "test", "data_path": "/tmp/data"}
286
+ result = _traverse_and_process(
287
+ value, Model, "local", MagicMock(), {}, "input_name"
288
+ )
289
+ assert isinstance(result, Model)
290
+ assert result.name == "test"
291
+ assert result.data_path == Path("/tmp/data")
292
+
293
+
294
+ def test_traverse_basemodel_from_instance():
295
+ class Model(BaseModel):
296
+ path: Path
297
+
298
+ inst = Model(path="/foo")
299
+ result = _traverse_and_process(
300
+ inst, Model, "local", MagicMock(), {}, "input_name"
301
+ )
302
+ assert isinstance(result, Model)
303
+ assert result.path == Path("/foo")
304
+
305
+
306
+ def test_traverse_basemodel_nested():
307
+ class Inner(BaseModel):
308
+ p: Path
309
+
310
+ class Outer(BaseModel):
311
+ inner: Inner
312
+ label: str
313
+
314
+ value = {"inner": {"p": "/nested"}, "label": "ok"}
315
+ result = _traverse_and_process(
316
+ value, Outer, "local", MagicMock(), {}, "input_name"
317
+ )
318
+ assert isinstance(result, Outer)
319
+ assert result.inner.p == Path("/nested")
320
+ assert result.label == "ok"
321
+
322
+
323
+ def test_traverse_basemodel_skips_none_fields():
324
+ class Model(BaseModel):
325
+ required: Path
326
+ optional: Optional[Path] = None
327
+
328
+ value = {"required": "/r", "optional": None}
329
+ result = _traverse_and_process(
330
+ value, Model, "local", MagicMock(), {}, "input_name"
331
+ )
332
+ assert result.required == Path("/r")
333
+ assert result.optional is None
334
+
335
+
336
+ def test_traverse_basemodel_non_dict_value_returned_unchanged():
337
+ class Model(BaseModel):
338
+ x: int
339
+
340
+ result = _traverse_and_process(
341
+ "not a dict", Model, "local", MagicMock(), {}, "input_name"
342
+ )
343
+ assert result == "not a dict"
344
+
345
+
346
+ def test_traverse_optional_path_present():
347
+ result = _traverse_and_process(
348
+ "/some/path", Optional[Path], "local", MagicMock(), {}, "input_name"
349
+ )
350
+ assert result == Path("/some/path")
351
+
352
+
353
+ def test_traverse_optional_path_none():
354
+ with pytest.raises(TypeError, match="expected str, bytes or os.PathLike"):
355
+ _traverse_and_process(
356
+ None, Optional[Path], "local", MagicMock(), {}, "input_name"
357
+ )