uenv-sdk 0.1.0__py3-none-any.whl

uenv/sdk/__init__.py

@@ -0,0 +1,413 @@
+"""Unified Python component SDK; Worker owns process and resource isolation."""
+
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from copy import deepcopy
+from dataclasses import dataclass, field, fields
+from typing import BinaryIO, Callable, Iterable
+import json
+import hashlib
+import warnings
+from os import PathLike
+from .modeling import (
+    ArtifactRef,
+    ContentPart,
+    EvaluationPlan,
+    Field,
+    UEnvModel,
+    bind_schema,
+    model_from_envelope,
+    model_json_schema,
+    model_to_envelope,
+)
+from .schema_registry import SchemaRegistry
+from .workspace import ArtifactStore, CommandResult, RemoteWorkspace, Workspace
+from .public_types import (
+    ArtifactRef as ArtifactReference, ContentPart as ContentItem, DatasetRef,
+    RuntimeSpec, TrajectoryEvent, GenerationEvent, Metric, ScoreResultGenerationRewardsItem as GenerationReward,
+)
+from .content import normalize_content
+
+
+def to_wire(value):
+    if isinstance(value, UEnvModel):
+        return model_to_envelope(value)
+    if hasattr(value, "__dataclass_fields__"):
+        return {
+            f.name: to_wire(getattr(value, f.name))
+            for f in fields(value)
+            if getattr(value, f.name) is not None or f.name in {"success", "reward"}
+        }
+    if isinstance(value, dict):
+        return {k: to_wire(v) for k, v in value.items()}
+    if isinstance(value, (tuple, list)):
+        return [to_wire(v) for v in value]
+    return value
+
+
+@dataclass(frozen=True)
+class PreparedSample:
+    sample_id: str
+    input: UEnvModel
+    private_data: UEnvModel | None = None
+    runtime: RuntimeSpec | None = None
+
+
+@dataclass(frozen=True)
+class TaskSpec:
+    """Typed component view; transport-only digest and envelope stay hidden."""
+
+    # 这是数据集作者在 reset(task, seed=...) 中收到的任务对象，不需要作者自行构造。
+    # Adapter 先产生 PreparedSample，SDK 组装传输数据，Worker 校验后还原成本对象。
+    # 把它当作只读任务输入使用；标准答案、隐藏测试等 private_data 不在这里。
+    # 最常用的是 task.input；其具体类型由本数据集 models.py 的公开输入模型决定。
+
+    # 数据来源标识：id（数据集 ID）、revision（固定数据版本）、split（划分）、
+    # subset（子集，未指定时为空字符串）。由准备流程传入，不在 reset 中重新选择。
+    dataset: DatasetRef
+    # Adapter 提供的样本编号，用于定位该数据版本中的一条题目。
+    sample_id: str
+    # Adapter 填写的公开输入。例：GSM8K 为 Gsm8kInput，SWE-Pro 为 SweProInput。
+    # 使用 task.input.instruction 等模型属性，不需要手动解开传输层的 schema_ref/data。
+    input: UEnvModel
+    # 可选的样本运行资源声明；当前有值时为 {"image": "镜像名@sha256:摘要"}。
+    # None 表示样本未单独指定。实际资源选择由系统完成，这不是启动 Backend 的接口。
+    runtime: RuntimeSpec | None = None
+
+
+def task_from_wire(value: dict, input_model: type[UEnvModel]) -> TaskSpec:
+    return TaskSpec(
+        dataset=deepcopy(value["dataset"]),
+        sample_id=value["sample_id"],
+        input=model_from_envelope(input_model, value["input"]),
+        runtime=deepcopy(value.get("runtime")),
+    )
+
+
+class DatasetAdapter(ABC):
+    def bind_preparation(self, *, upload_file=None, run_spec=None):
+        """System-owned preparation capabilities; never give the Adapter Hub credentials."""
+        self._upload_file = upload_file
+        self._preparation_run_spec = deepcopy(run_spec)
+        return self
+
+    def store_text(self, value: str | ArtifactReference, *, media_type: str = "text/plain") -> ArtifactReference:
+        """Convert source text to the ordinary immutable file reference."""
+        if isinstance(value, dict):
+            return deepcopy(value)
+        if not isinstance(value, str):
+            raise TypeError("Source material must be text or ArtifactRef")
+        import tempfile
+        from pathlib import Path
+
+        with tempfile.TemporaryDirectory(prefix="uenv-prepare-") as directory:
+            path = Path(directory) / "material"
+            path.write_bytes(value.encode("utf-8"))
+            return self.store_file(path, media_type=media_type)
+
+    def store_file(self, path: str | PathLike[str], *, media_type: str = "application/octet-stream") -> ArtifactReference:
+        """Store an explicitly selected image/audio/archive/test file.
+
+        The caller provides the upload capability: publication uploads to Hub,
+        direct task preparation uploads to Server. Dataset code needs neither
+        service addresses nor credentials and returns the same ArtifactRef.
+        """
+        upload = getattr(self, "_upload_file", None)
+        if upload is None:
+            raise ValueError("MATERIAL_UPLOAD_REQUIRED")
+        return upload(path, media_type=media_type)
+
+    @staticmethod
+    def read_rows(path: str | PathLike[str]) -> Iterable[dict]:
+        """Read JSONL or Parquet; override for a dataset-specific source layout.
+
+        Runs only during trusted preparation. Hub samples remain standardized
+        JSONL, and Worker never reads the original source file.
+        """
+        from pathlib import Path
+
+        if Path(path).suffix.lower() == ".parquet":
+            try:
+                import pyarrow.parquet as parquet
+            except ImportError as exc:
+                raise RuntimeError(
+                    "Install uenv-sdk[parquet] in the preparation environment"
+                ) from exc
+            # Only file decoding is shared; business fields remain in normalize().
+            with parquet.ParquetFile(path) as source:
+                for batch in source.iter_batches(batch_size=128):
+                    yield from batch.to_pylist()
+            return
+        with open(path, encoding="utf-8") as source:
+            for line in source:
+                if line.strip():
+                    yield json.loads(line)
+
+    @abstractmethod
+    def normalize(self, row: dict) -> PreparedSample:
+        """Deterministically convert one source row; never choose backend or agent."""
+
+
+@dataclass(frozen=True)
+class Observation:
+    """Agent-visible text, file references or package-owned structured data.
+
+    Authors pass a string, a file reference, a UEnvModel, or an ordered list of
+    these values. SDK normalization keeps the internal ContentPart[] protocol.
+    """
+
+    content: str | UEnvModel | dict | list = field(default_factory=list)
+    terminated: bool = False
+    episode_truncated: bool = False
+
+    def __post_init__(self):
+        object.__setattr__(self, "content", normalize_content(self.content))
+
+
+class Environment(ABC):
+    def __init__(self, *, workspace: Workspace, artifacts: ArtifactStore):
+        # 系统创建环境时提供执行能力；每道任务的数据在 reset(task, seed=...) 中传入。
+        # 这里不保存 task、seed 或截止时间；它们属于每一次调用。
+        # 已准备好的工作区接口：workspace、read_file、write_file、run。
+        # 它操作当前后端内的任务工作区，不负责选择或创建容器。
+        self.workspace = workspace
+        # 系统管理的文件引用：read_artifact(ref)、copy_artifact(ref, output)、
+        # put_bytes(content, media_type)。只允许访问本任务公开材料，不含私有答案。
+        self.artifacts = artifacts
+
+    @abstractmethod
+    def reset(self, task: TaskSpec, *, seed: int) -> Observation:
+        """Initialize this task and return the Agent's first observation."""
+        # 【必须实现】task.input 是本数据集公开输入模型；seed 是本次随机种子。
+        # 纯问答返回题目；需准备工作区时调用 self.workspace / self.artifacts。
+        # self.workspace.run(argv) 自动使用本次调用剩余预算，不要求作者计算截止时间。
+        # Worker 负责输入/输出校验和运行限制；这里不调用模型、评分或调度后端。
+
+    def finalize(self, task: TaskSpec, final_answer: list[dict]) -> list[dict]:
+        """Collect the submission after interaction, within the finalization reserve.
+
+        Default: preserve the Agent answer. Dataset environments may snapshot
+        their public workspace here. No model calls or private scoring inputs
+        are available. Errors fail the attempt; they must not become empty answers.
+        """
+        return deepcopy(final_answer)
+
+    def close(self):
+        """Optional plugin cleanup; Worker still owns underlying resources."""
+        # 【可选实现】释放作者自己创建的连接等；底层容器和进程仍由系统回收。
+
+
+class AgentRuntimeError(RuntimeError):
+    """SDK representation of the existing Worker ErrorRecord, not a new wire type."""
+
+    def __init__(self, error: dict):
+        self.error = deepcopy(error)
+        super().__init__(error["code"])
+
+    @property
+    def code(self):
+        return self.error["code"]
+
+
+class AgentContext(ABC):
+    """Worker-provided capability view; it never owns scheduling or budgets."""
+
+    task: TaskSpec
+    observation: Observation
+    tools: tuple[dict, ...]
+    seed: int
+    artifacts: object  # Task-authorized file reads and writes supplied by Worker.
+    tool_services: dict[str, str]  # Worker-created endpoints by selected tool name.
+    # observation is the current reset/step result, including end flags.
+
+    @property
+    def mcp_config(self) -> dict:
+        """Worker-managed MCP connection for selected Python tools.
+
+        Pass this to the framework's MCP client. Calls must preserve arguments
+        returned by generate(); URLs and credentials must not be persisted.
+        An empty tool selection returns an empty configuration.
+        """
+        raise NotImplementedError("The Worker supplies the MCP connection")
+
+    @abstractmethod
+    async def generate(self, messages: list[dict], *, enable_tools: bool = True) -> dict:
+        """Request one generation through Rust within the original budget and trace.
+
+        enable_tools=False requests text-only generation without changing selected
+        bindings or execution permission. Denials raise AgentRuntimeError.
+        """
+
+    @abstractmethod
+    async def call_tool(self, call: dict, *, execute: Callable | None = None) -> Observation:
+        """Return public feedback after Worker admission.
+
+        Framework-native tools may supply execute(arguments, context), returning
+        Observation (or an awaitable). It runs locally only after admission;
+        omitting it uses the selected tool definition. Native objects stay local.
+        """
+
+
+class AgentRunner(ABC):
+    def __init__(self, config: dict):
+        if not isinstance(config, dict):
+            raise TypeError("AgentRunner config must be an object")
+        self.config = deepcopy(config)
+
+    def prepare(
+        self, context: AgentContext, bindings: list[dict], native_factories: dict,
+        python_definitions: dict[str, ToolDefinition], workspace: str, model_id: str,
+    ) -> dict[str, ToolDefinition]:
+        """Synchronously prepare once per attempt, returning selected native tools.
+
+        Worker supplies the resolved tool bindings, selected native factories,
+        Python definitions keyed by selected name, workspace and model identity.
+        No task is bound to context yet; do not generate or execute tools here.
+        Return exactly the selected native ToolDefinitions keyed by public binding
+        name, never implementation ID. Python definitions stay Worker-owned.
+        The default returns no native tools (for example PlainAgent). Publishing
+        checks signatures only; it never calls prepare. On failure, close must
+        still release any resources acquired by this Agent.
+        """
+        return {}
+
+    def close(self):
+        """Release Agent-owned framework resources, including partial setup.
+
+        Called after tool freeze, even when setup/run fails. Must tolerate an
+        unprepared or already closed Agent. Worker still owns backend resources.
+        """
+
+    @abstractmethod
+    async def run(self, context: AgentContext) -> list[dict]:
+        """Own the interaction loop; receive only visible capabilities."""
+
+
+class ToolExecutor(ABC):
+    """Deprecated SDK base; removed in uenv-sdk 0.2.0.
+
+    Use @tool returning Observation, or ToolDefinition for native integration.
+    The package loader does not discover ToolExecutor subclasses. This is not
+    openhands.sdk.tool.ToolExecutor, which remains the framework's own interface.
+    """
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        warnings.warn(
+            "uenv.sdk.ToolExecutor is deprecated and will be removed in 0.2.0; "
+            "use @tool returning Observation or ToolDefinition instead",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+    def __init__(self, config: dict):
+        if not isinstance(config, dict):
+            raise TypeError("ToolExecutor config must be an object")
+        self.config = deepcopy(config)
+
+    @abstractmethod
+    async def execute(self, arguments: dict, context) -> dict:
+        """Execute one selected tool; input and output follow its published schemas."""
+
+
+@dataclass(frozen=True)
+class ScoreInput:
+    task: TaskSpec
+    final_answer: list[ContentItem]
+    trajectory_ref: ArtifactReference
+    private_data: UEnvModel | None = None
+
+
+@dataclass(frozen=True)
+class ScoreResult:
+    success: bool | None = None
+    metrics: list[Metric] = field(default_factory=list)
+    reward: float | None = None
+    evidence: list[ArtifactReference] = field(default_factory=list)
+    generation_rewards: list[GenerationReward] | None = None
+
+    # Only the Rust Worker fills these fields. Scorer-returned values are rejected.
+    status: str | None = None
+    scorer: dict | None = None
+    error: dict | None = None
+
+    @classmethod
+    def binary(cls, passed):
+        if type(passed) is not bool:
+            raise TypeError("Binary score requires an actual bool")
+        value = float(passed)
+        return cls(
+            passed,
+            [
+                {
+                    "name": "accuracy",
+                    "value": value,
+                    "unit": "ratio",
+                    "direction": "higher",
+                }
+            ],
+            value,
+        )
+
+
+binary_score = ScoreResult.binary
+
+
+class Scorer(ABC):
+    def __init__(
+        self,
+        *,
+        workspace: Workspace,
+        workspace_factory: Callable,
+        read_artifact: Callable,
+        read_trajectory: Callable,
+        copy_artifact: Callable | None = None,
+    ):
+        # Worker 创建评分器时提供能力；普通作者沿用基类构造函数即可。
+        # 这些接口只在本次评分期间有效，不提供 Agent 的可写工作区或容器管理权限。
+        self.workspace = workspace
+        self._workspace_factory = workspace_factory
+        self._read_artifact = read_artifact
+        self._read_trajectory = read_trajectory
+        self._copy_artifact = copy_artifact
+
+    @abstractmethod
+    def score(self, request: ScoreInput) -> ScoreResult:
+        """Score the submitted answer and trajectory using optional reference data."""
+        # 【必须实现】只接收本次评分数据，不再接收系统 context。
+        # 简单任务直接比较答案；测试评分或过程评分按需使用下面的系统方法。
+        # 超时、取消、协议校验和文件授权由 Worker/Runtime 负责。
+
+    def create_workspace(self) -> RemoteWorkspace:
+        """Create an additional isolated workspace; use self.workspace for ordinary scoring."""
+        return RemoteWorkspace(self._workspace_factory)
+
+    def read_artifact(self, reference: ArtifactReference) -> bytes:
+        """Read a file authorized for this scoring request, including private material."""
+        return self._read_artifact(reference)
+
+    def copy_artifact(self, reference: ArtifactReference, output: BinaryIO) -> None:
+        """Stream an authorized artifact to a binary writer without exposing private files."""
+        if self._copy_artifact is None:
+            raise RuntimeError("SCORER_ARTIFACT_STREAM_UNAVAILABLE")
+        self._copy_artifact(reference, output)
+
+    def read_trajectory(self, request: ScoreInput) -> list[TrajectoryEvent]:
+        """Read and validate this task's recorded scoring checkpoint events."""
+        # 返回统一轨迹事件，可按 generation_id 计算过程奖励，不生成新回答。
+        return self._read_trajectory(request)
+
+
+def read_reference_text(request: ScoreInput) -> tuple[str, dict]:
+    """Opt-in extraction for rules requiring text and private reference material."""
+    if request.private_data is None:
+        raise ValueError("This scoring rule requires reference material")
+    if any(part["kind"] != "text" for part in request.final_answer):
+        raise ValueError("Text-only scorer cannot discard non-text answer parts")
+    return "".join(
+        part["text"] for part in request.final_answer
+    ), request.private_data.model_dump()
+
+
+# User-defined tools use the same authoring API in independent and dataset packages.
+from .tools import SandboxToolContext, ToolContext, ToolDefinition, tool

uenv/sdk/artifact_io.py

@@ -0,0 +1,26 @@
+"""Internal chunk transport; authors continue to use the same ArtifactRef."""
+
+import hashlib
+
+
+def artifact_chunks(peer, method, reference):
+    size = reference["size_bytes"]
+    offset, hasher = 0, hashlib.sha256()
+    while offset < size or offset == 0:
+        data = bytes(
+            peer.request(method, {"artifact_ref": reference, "offset": offset})
+        )
+        expected = min(256 * 1024, size - offset)
+        if len(data) != expected:
+            raise ValueError("ARTIFACT_STREAM_LENGTH_MISMATCH")
+        hasher.update(data)
+        offset += len(data)
+        yield data
+        if offset == size:
+            break
+    if "sha256:" + hasher.hexdigest() != reference["digest"]:
+        raise ValueError("ARTIFACT_STREAM_DIGEST_MISMATCH")
+
+
+def read_artifact(peer, method, reference):
+    return b"".join(artifact_chunks(peer, method, reference))