From e6e1c75a7a47b56583d4109285f90404f0d322d0 Mon Sep 17 00:00:00 2001 From: Asadullo Ganiev <62354884+solvemproblr@users.noreply.github.com> Date: Thu, 20 Aug 2026 01:36:00 +0500 Subject: [PATCH] feat(sdk): add typed UI contracts --- hud/environment/__init__.py | 6 ++ hud/environment/arguments.py | 38 +++++++++++ hud/environment/server.py | 19 +++--- hud/environment/tests/test_arguments.py | 85 +++++++++++++++++++++++++ 4 files changed, 139 insertions(+), 9 deletions(-) create mode 100644 hud/environment/arguments.py create mode 100644 hud/environment/tests/test_arguments.py diff --git a/hud/environment/__init__.py b/hud/environment/__init__.py index 3b9900795..1beb9bfb9 100644 --- a/hud/environment/__init__.py +++ b/hud/environment/__init__.py @@ -22,6 +22,7 @@ from hud.capabilities import Capability from hud.utils.modules import iter_modules +from .arguments import DataFileArg, DataFileRef, DataFilesArg, GradingArg, PromptArg from .egress import Peer from .env import Answer, Environment from .workspace import DEFAULT_SYSTEM_MOUNTS, Mount, MountKind, Workspace @@ -89,10 +90,15 @@ def load_environment( "DEFAULT_SYSTEM_MOUNTS", "Answer", "Capability", + "DataFileArg", + "DataFileRef", + "DataFilesArg", "Environment", + "GradingArg", "Mount", "MountKind", "Peer", + "PromptArg", "Workspace", "load_environment", ] diff --git a/hud/environment/arguments.py b/hud/environment/arguments.py new file mode 100644 index 000000000..59cf09472 --- /dev/null +++ b/hud/environment/arguments.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from typing import Annotated, TypeAlias, TypeVar + +from pydantic import BaseModel, ConfigDict, Field + +T = TypeVar("T") + + +class DataFileRef(BaseModel): + model_config = ConfigDict(extra="forbid") + + file_id: str = Field(description="HUD data-file id") + path: str | None = Field(default=None, description="Environment-owned destination path") + + +PromptArg: TypeAlias = Annotated[ + str, + Field(json_schema_extra={"x-hud-hint": "prompt"}), +] + +DataFileArg: TypeAlias = Annotated[ + T, + Field(json_schema_extra={"x-hud-hint": "data-file"}), +] + +DataFilesArg: TypeAlias = Annotated[ + list[T], + Field(json_schema_extra={"x-hud-hint": "data-files"}), +] + +GradingArg: TypeAlias = Annotated[ + list[T], + Field(json_schema_extra={"x-hud-hint": "grading"}), +] + + +__all__ = ["DataFileArg", "DataFileRef", "DataFilesArg", "GradingArg", "PromptArg"] diff --git a/hud/environment/server.py b/hud/environment/server.py index 089aafacd..36ad47818 100644 --- a/hud/environment/server.py +++ b/hud/environment/server.py @@ -72,23 +72,24 @@ def _jsonable(value: Any) -> Any: def _coerce_args(sig: inspect.Signature, args: dict[str, Any]) -> dict[str, Any]: - """Coerce string wire args into the task fn's annotated param types. - - JSON-RPC sends args as JSON scalars/strings; a param annotated with a richer - type (Pydantic model, list, etc.) is validated via a ``TypeAdapter``. Values - that already match (or fail to validate) are passed through unchanged. - """ coerced: dict[str, Any] = {} for name, value in args.items(): param = sig.parameters.get(name) annotation = param.annotation if param is not None else inspect.Parameter.empty - if annotation in (inspect.Parameter.empty, str, Any) or not isinstance(value, str): + if annotation in (inspect.Parameter.empty, Any): coerced[name] = value continue + adapter = TypeAdapter(annotation) try: - coerced[name] = TypeAdapter(annotation).validate_json(value) + coerced[name] = adapter.validate_python(value) except ValidationError: - coerced[name] = value + if not isinstance(value, str): + coerced[name] = value + continue + try: + coerced[name] = adapter.validate_json(value) + except ValidationError: + coerced[name] = value return coerced diff --git a/hud/environment/tests/test_arguments.py b/hud/environment/tests/test_arguments.py new file mode 100644 index 000000000..c4424ae11 --- /dev/null +++ b/hud/environment/tests/test_arguments.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel + +from hud.environment import ( + DataFileArg, + DataFileRef, + DataFilesArg, + Environment, + GradingArg, + PromptArg, +) +from hud.eval import Run + +from .conftest import served + + +class _Attachment(DataFileRef): + expand: bool = False + + +class _Fixture(DataFileRef): + channel: Literal["mail", "calendar"] + + +class _Criterion(BaseModel): + requirement: str + weight: float = 1.0 + guidance: str | None = None + + +def test_argument_types_publish_editor_hints_and_model_schemas() -> None: + env = Environment("typed-args") + + @env.template() + async def task( + prompt: PromptArg, + attachment: DataFileArg[_Attachment], + fixtures: DataFilesArg[_Fixture], + criteria: GradingArg[_Criterion], + ): + yield prompt + yield 1.0 + + schema = task.manifest_entry()["args"] + properties = schema["properties"] + + assert properties["prompt"]["x-hud-hint"] == "prompt" + assert properties["attachment"]["x-hud-hint"] == "data-file" + assert properties["fixtures"]["x-hud-hint"] == "data-files" + assert properties["criteria"]["x-hud-hint"] == "grading" + assert properties["fixtures"]["items"]["$ref"] == "#/$defs/_Fixture" + assert properties["criteria"]["items"]["$ref"] == "#/$defs/_Criterion" + assert "channel" in schema["$defs"]["_Fixture"]["properties"] + assert "guidance" in schema["$defs"]["_Criterion"]["properties"] + + +async def test_typed_argument_values_reach_task_as_models() -> None: + env = Environment("typed-args") + + @env.template() + async def task( + attachments: DataFilesArg[_Attachment], + criteria: GradingArg[_Criterion], + ): + assert isinstance(attachments[0], _Attachment) + assert isinstance(criteria[0], _Criterion) + yield f"{attachments[0].path}:{criteria[0].requirement}" + yield 1.0 + + async with ( + served(env) as client, + Run( + client, + "task", + { + "attachments": [{"file_id": "file-1", "path": "brief.pdf"}], + "criteria": [{"requirement": "Answer the question", "weight": 2.0}], + }, + ) as run, + ): + assert run.prompt_text == "brief.pdf:Answer the question" + run.trace.content = "done"