bohr-agent-sdk 0.1.122__py3-none-any.whl → 0.1.123__py3-none-any.whl

{bohr_agent_sdk-0.1.122.dist-info → bohr_agent_sdk-0.1.123.dist-info}/METADATA

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

{bohr_agent_sdk-0.1.122.dist-info → bohr_agent_sdk-0.1.123.dist-info}/RECORD

@@ -62,12 +62,12 @@ 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=oD93SjlcHgSmGtxGJnl4V6cN1icUJSW3KEYt0Lctc6Q,19756
+dp/agent/server/calculation_mcp_server.py,sha256=z4tzhT2hMN9HRiTp0cg2AwYuqlz146SMFbv7bHdAxZ8,23062
 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=lWofctAt2Nemrd3OF-IMBVWgVLKVanHDo3iMTE5lot0,12529
+dp/agent/server/executor/dispatcher_executor.py,sha256=LA79QIsTu6YY60mkn6o4iW1VWRqpAcJHQhKMENKPN1c,12643
 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
@@ -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.122.dist-info/METADATA,sha256=6PXBgaHU2sAKQJHYdN52tD18u7uIpmFEAgVbMBeSNWA,11070
-bohr_agent_sdk-0.1.122.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.122.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.122.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.122.dist-info/RECORD,,
+bohr_agent_sdk-0.1.123.dist-info/METADATA,sha256=5ptmdilTXVCoFGsYIAplF1uiG7vn9gwHaBX4xpS2sks,11070
+bohr_agent_sdk-0.1.123.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+bohr_agent_sdk-0.1.123.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.123.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.123.dist-info/RECORD,,

{bohr_agent_sdk-0.1.122.dist-info → bohr_agent_sdk-0.1.123.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

dp/agent/server/calculation_mcp_server.py

@@ -7,7 +7,7 @@ from copy import deepcopy
 from datetime import datetime
 from pathlib import Path
 from urllib.parse import urlparse
-from typing import Annotated, Literal, Optional, List, Dict
+from typing import Annotated, Literal, Optional, List, Dict, Union, Any, get_origin, get_args
 
 from mcp.server.fastmcp import FastMCP
 from mcp.server.fastmcp.utilities.context_injection import (
@@ -30,12 +30,18 @@ CALCULATION_MCP_WORKDIR = os.getenv("CALCULATION_MCP_WORKDIR", os.getcwd())
 
 
 def parse_uri(uri):
-    scheme = urlparse(uri).scheme
+    parsed = urlparse(uri)
+    scheme = parsed.scheme
     if scheme == "":
         key = uri
         scheme = "local"
     else:
-        key = uri[len(scheme)+3:]
+        if parsed.netloc:
+            key = parsed.netloc + parsed.path
+        else:
+            key = parsed.path
+        if parsed.query:
+            key += "?" + parsed.query
     return scheme, key
 
 
@@ -105,96 +111,128 @@ def terminate_job(job_id: str, executor: Optional[dict] = None):
         logger.info("Job %s is terminated" % job_id)
 
 
+def _normalize_annotation(ann):
+    if ann is None:
+        return None
+    origin = get_origin(ann)
+    if origin is Annotated:
+        return _normalize_annotation(get_args(ann)[0])
+    if origin is Union:
+        args = get_args(ann)
+        if type(None) in args:
+            non_none = [a for a in args if a is not type(None)]
+            if non_none:
+                return _normalize_annotation(non_none[0])
+    return ann
+
+
+def _download_artifact(uri, storage, storage_type, input_artifacts,
+                       input_name, path_trace):
+    scheme, key = parse_uri(uri)
+    if scheme == storage_type:
+        s = storage
+    else:
+        s = storage_dict[scheme]()
+    rel = "/".join(str(p) for p in path_trace) if path_trace else ""
+    download_dir = os.path.join("inputs", input_name, rel) if rel else os.path.join("inputs", input_name)
+    os.makedirs(download_dir, exist_ok=True)
+    path = s.download(key, download_dir)
+    logger.info("Artifact %s downloaded to %s" % (uri, path))
+    if input_name not in input_artifacts:
+        input_artifacts[input_name] = {"storage_type": scheme, "uri": []}
+    if isinstance(input_artifacts[input_name].get("uri"), list):
+        input_artifacts[input_name]["uri"].append(uri)
+    return path
+
+
+def _traverse_and_process(value, annotation, storage_type, storage,
+                          input_artifacts, input_name, path_trace=None):
+    if path_trace is None:
+        path_trace = []
+    ann = _normalize_annotation(annotation)
+    if ann is None:
+        return value
+    origin = get_origin(ann)
+    args = get_args(ann)
+
+    # Path: only when resolved annotation is Path (Optional[Path] is normalized to Path)
+    if ann is Path:
+        s = str(value)
+        if not s:
+            return Path(".")
+        parsed = urlparse(s)
+        if parsed.scheme and len(parsed.scheme) > 1:
+            return Path(_download_artifact(
+                s, storage, storage_type, input_artifacts, input_name, path_trace))
+        return Path(value)
+
+    # BaseModel: schema-driven traversal over model fields (check before dict so nesting works)
+    if isinstance(ann, type) and issubclass(ann, BaseModel):
+        # Convert to dict for traversal; re-instantiate to model at the end so callers get objects
+        if isinstance(value, BaseModel):
+            value = value.model_dump()
+        if not isinstance(value, dict):
+            return value
+        out = dict(value)
+        for field_name, field_info in ann.model_fields.items():
+            if field_name in out and out[field_name] is not None:
+                out[field_name] = _traverse_and_process(
+                    out[field_name],
+                    field_info.annotation,
+                    storage_type,
+                    storage,
+                    input_artifacts,
+                    input_name,
+                    path_trace + [field_name],
+                )
+        # Re-instantiate so tool functions receive model instances (dot notation works)
+        try:
+            return ann.model_validate(out)
+        except Exception as e:
+            logger.warning("Failed to re-instantiate model %s: %s", ann.__name__, e)
+            return out
+
+    # List: use inner type from type args
+    if origin in (list, List) and isinstance(value, (list, tuple)):
+        inner = _normalize_annotation(args[0]) if args else Any
+        return [
+            _traverse_and_process(
+                item, inner, storage_type, storage,
+                input_artifacts, input_name, path_trace + [i])
+            for i, item in enumerate(value)
+        ]
+
+    # Dict: use value type from type args (e.g. Dict[str, Path] processes values as Path)
+    if origin in (dict, Dict) and isinstance(value, dict):
+        value_type = _normalize_annotation(args[1]) if (args and len(args) > 1) else Any
+        return {
+            k: _traverse_and_process(
+                v, value_type, storage_type, storage,
+                input_artifacts, input_name, path_trace + [k])
+            for k, v in value.items()
+        }
+
+    return value
+
+
 def handle_input_artifacts(fn, kwargs, storage):
     storage_type, storage = init_storage(storage)
-    sig = inspect.signature(fn)
+    sig = inspect.signature(fn, eval_str=True)
     input_artifacts = {}
+    new_kwargs = {}
     for name, param in sig.parameters.items():
-        if param.annotation is Path or (
-            param.annotation is Optional[Path] and
-                kwargs.get(name) is not None):
-            uri = kwargs[name]
-            scheme, key = parse_uri(uri)
-            if scheme == storage_type:
-                s = storage
-            else:
-                s = storage_dict[scheme]()
-            path = s.download(key, "inputs/%s" % name)
-            logger.info("Artifact %s downloaded to %s" % (
-                uri, path))
-            kwargs[name] = Path(path)
-            input_artifacts[name] = {
-                "storage_type": scheme,
-                "uri": uri,
-            }
-        elif param.annotation is List[Path] or (
-            param.annotation is Optional[List[Path]] and
-                kwargs.get(name) is not None):
-            uris = kwargs[name]
-            new_paths = []
-            for i, uri in enumerate(uris):
-                scheme, key = parse_uri(uri)
-                if scheme == storage_type:
-                    s = storage
-                else:
-                    s = storage_dict[scheme]()
-                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))
-            kwargs[name] = new_paths
-            input_artifacts[name] = {
-                "storage_type": storage_type,
-                "uri": uris,
-            }
-        elif param.annotation is Dict[str, Path] or (
-            param.annotation is Optional[Dict[str, Path]] and
-                kwargs.get(name) is not None):
-            uris_dict = kwargs[name]
-            new_paths_dict = {}
-            for key_name, uri in uris_dict.items():
-                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}")
-                new_paths_dict[key_name] = Path(path)
-                logger.info("Artifact %s (key=%s) downloaded to %s" % (
-                    uri, key_name, path))
-            kwargs[name] = new_paths_dict
-            input_artifacts[name] = {
-                "storage_type": storage_type,
-                "uri": uris_dict,
-            }
-        elif param.annotation is Dict[str, List[Path]] or (
-            param.annotation is Optional[Dict[str, List[Path]]] and
-                kwargs.get(name) is not None):
-            uris_dict = kwargs[name]
-            new_paths_dict = {}
-            for key_name, uris in uris_dict.items():
-                new_paths = []
-                for i, uri in enumerate(uris):
-                    scheme, key = parse_uri(uri)
-                    if scheme == storage_type:
-                        s = storage
-                    else:
-                        s = storage_dict[scheme]()
-                    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))
-                new_paths_dict[key_name] = new_paths
-            kwargs[name] = new_paths_dict
-            input_artifacts[name] = {
-                "storage_type": storage_type,
-                "uri": uris_dict,
-            }
-    return kwargs, input_artifacts
+        if name not in kwargs:
+            if param.default is not inspect.Parameter.empty:
+                new_kwargs[name] = param.default
+            continue
+        val = kwargs[name]
+        if val is None and _normalize_annotation(param.annotation) != param.annotation:
+            new_kwargs[name] = val
+            continue
+        new_kwargs[name] = _traverse_and_process(
+            val, param.annotation, storage_type, storage,
+            input_artifacts, name)
+    return new_kwargs, input_artifacts
 
 
 def handle_output_artifacts(results, exec_id, storage):
@@ -269,6 +307,57 @@ annotation_map = {
 }
 
 
+# Cache for schema models derived from BaseModel (Path -> str in JSON schema)
+_schema_model_cache: Dict[type, type] = {}
+
+def get_schema_annotation(annotation: Any) -> Any:
+    """
+    Map an annotation to the type used in JSON schema (e.g. Path -> str).
+    For BaseModel, build a schema model with the same structure but Path fields as str.
+    Handles List[...], Dict[...], Optional[...] and nested BaseModel recursively.
+    """
+    if annotation is None or annotation is type(None):
+        return annotation
+    origin = get_origin(annotation)
+    if origin is Annotated:
+        return get_schema_annotation(get_args(annotation)[0])
+    if origin is Union:
+        args = get_args(annotation)
+        if type(None) in args:
+            non_none = [a for a in args if a is not type(None)]
+            if len(non_none) == 1:
+                return Optional[get_schema_annotation(non_none[0])]
+    if annotation in annotation_map:
+        return annotation_map[annotation]
+    # List[X] -> List[schema(X)] so e.g. List[BaseModel] becomes List[BaseModelSchema]
+    if origin in (list, List):
+        type_args = get_args(annotation)
+        inner = get_schema_annotation(type_args[0]) if type_args else Any
+        return List[inner]
+    # Dict[K, V] -> Dict[K, schema(V)]
+    if origin in (dict, Dict):
+        type_args = get_args(annotation)
+        if type_args and len(type_args) >= 2:
+            return Dict[type_args[0], get_schema_annotation(type_args[1])]
+        return annotation
+    if isinstance(annotation, type) and issubclass(annotation, BaseModel):
+        if annotation not in _schema_model_cache:
+            schema_fields = {}
+            for name, field_info in annotation.model_fields.items():
+                fa = get_schema_annotation(field_info.annotation)
+                if field_info.is_required():
+                    schema_fields[name] = (fa, Field())
+                else:
+                    default = getattr(field_info, "default", None)
+                    schema_fields[name] = (fa, Field(default=default))
+            _schema_model_cache[annotation] = create_model(
+                f"{annotation.__name__}Schema",
+                **schema_fields,
+            )
+        return _schema_model_cache[annotation]
+    return annotation
+
+
 class SubmitResult(BaseModel):
     job_id: str
     extra_info: dict | None = None
@@ -329,11 +418,8 @@ class CalculationMCPServer:
         for n, annotation in \
                 func_arg_metadata.arg_model.__annotations__.items():
             param = params[n]
-            if param.annotation in annotation_map:
-                model_params[n] = Annotated[
-                    (annotation_map[param.annotation], Field())]
-            else:
-                model_params[n] = annotation
+            schema_ann = get_schema_annotation(param.annotation)
+            model_params[n] = Annotated[(schema_ann, Field())]
             if param.default is not inspect.Parameter.empty:
                 model_params[n] = (model_params[n], param.default)
         for n, param in inspect.signature(new_fn).parameters.items():

dp/agent/server/executor/dispatcher_executor.py

@@ -200,6 +200,9 @@ class DispatcherExecutor(BaseExecutor):
             if isinstance(value, Path):
                 forward_files.append(str(value))
 
+        if os.path.isdir("inputs") and "inputs" not in forward_files:
+            forward_files.append("inputs")
+
         task = {
             "task_work_path": "./",
             "outlog": "log",