From 385d11d7cf1dbb4a3ca20b103ca0b09349ff0186 Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Tue, 28 Jul 2026 20:42:06 +0300 Subject: [PATCH 01/13] Refactoring --- crates/nagare/src/chat/mod.rs | 21 +++++++-------------- crates/uzu/examples/tool_calls.rs | 2 +- crates/uzu/tests/tool/calls.rs | 5 +---- 3 files changed, 9 insertions(+), 19 deletions(-) diff --git a/crates/nagare/src/chat/mod.rs b/crates/nagare/src/chat/mod.rs index 23ddd4917..5dc5ea753 100644 --- a/crates/nagare/src/chat/mod.rs +++ b/crates/nagare/src/chat/mod.rs @@ -162,25 +162,18 @@ impl ChatSession { self.instance.lock().await.peak_memory_usage() } - pub async fn add_tool_function( - &mut self, - function: impl Into, - ) -> Result<(), ChatSessionError> { - self.add_tool_function_definitions(vec![function.into()]).await - } - pub async fn add_tool_descriptor( &mut self, - definition: ToolDescriptor, + descriptor: impl Into, ) -> Result<(), ChatSessionError> { - self.add_tool_function_definitions(vec![definition]).await + self.add_tool_descriptors(vec![descriptor.into()]).await } - pub async fn add_tool_function_definitions( + pub async fn add_tool_descriptors( &mut self, - definitions: Vec, + descriptors: Vec, ) -> Result<(), ChatSessionError> { - if definitions.is_empty() { + if descriptors.is_empty() { return Ok(()); } let Some(registry) = self.tool_registry.as_ref() else { @@ -193,8 +186,8 @@ impl ChatSession { // an incoming name must displace a stale caller-provided declaration too let owned_namespaces = { let mut registry_guard = registry.lock().await; - for def in definitions { - registry_guard.add_function(def) + for desc in descriptors { + registry_guard.add_function(desc) } registry_guard.get_namespaces() }; diff --git a/crates/uzu/examples/tool_calls.rs b/crates/uzu/examples/tool_calls.rs index b56a38631..adf4562d2 100644 --- a/crates/uzu/examples/tool_calls.rs +++ b/crates/uzu/examples/tool_calls.rs @@ -35,7 +35,7 @@ async fn main() -> Result<(), Box> { } let mut session = engine.chat(model, ChatConfig::default()).await?; - session.add_tool_function(get_current_location).await?; + session.add_tool_descriptor(get_current_location).await?; session .add_tool_descriptor(uzu_tool_closure! { /// Returns temperature in provided location diff --git a/crates/uzu/tests/tool/calls.rs b/crates/uzu/tests/tool/calls.rs index f81a07289..fdf7b4410 100644 --- a/crates/uzu/tests/tool/calls.rs +++ b/crates/uzu/tests/tool/calls.rs @@ -117,10 +117,7 @@ async fn run_tool_calls_test( println!("Prompt: {}", case.prompt); let mut session = engine.chat(model.clone(), ChatConfig::default()).await.unwrap(); - session - .add_tool_function_definitions(case.tools.iter().map(|definition| definition()).collect()) - .await - .unwrap(); + session.add_tool_descriptors(case.tools.iter().map(|definition| definition()).collect()).await.unwrap(); let mut messages = Vec::new(); if with_system_message { From 2c6de702c1ecf155ac9e2627147e7e3d6dedb057 Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 17:14:35 +0300 Subject: [PATCH 02/13] Refactoring --- crates/nagare/src/chat/mod.rs | 6 +++--- crates/uzu/examples/tool_calls.rs | 4 ++-- crates/uzu/tests/tool/calls.rs | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/nagare/src/chat/mod.rs b/crates/nagare/src/chat/mod.rs index 5dc5ea753..0e0575161 100644 --- a/crates/nagare/src/chat/mod.rs +++ b/crates/nagare/src/chat/mod.rs @@ -162,14 +162,14 @@ impl ChatSession { self.instance.lock().await.peak_memory_usage() } - pub async fn add_tool_descriptor( + pub async fn add_tool( &mut self, descriptor: impl Into, ) -> Result<(), ChatSessionError> { - self.add_tool_descriptors(vec![descriptor.into()]).await + self.add_tools(vec![descriptor.into()]).await } - pub async fn add_tool_descriptors( + pub async fn add_tools( &mut self, descriptors: Vec, ) -> Result<(), ChatSessionError> { diff --git a/crates/uzu/examples/tool_calls.rs b/crates/uzu/examples/tool_calls.rs index adf4562d2..367ab609e 100644 --- a/crates/uzu/examples/tool_calls.rs +++ b/crates/uzu/examples/tool_calls.rs @@ -35,9 +35,9 @@ async fn main() -> Result<(), Box> { } let mut session = engine.chat(model, ChatConfig::default()).await?; - session.add_tool_descriptor(get_current_location).await?; + session.add_tool(get_current_location).await?; session - .add_tool_descriptor(uzu_tool_closure! { + .add_tool(uzu_tool_closure! { /// Returns temperature in provided location get_current_temperature: | /// Latitude in decimal degrees. diff --git a/crates/uzu/tests/tool/calls.rs b/crates/uzu/tests/tool/calls.rs index fdf7b4410..b33a94481 100644 --- a/crates/uzu/tests/tool/calls.rs +++ b/crates/uzu/tests/tool/calls.rs @@ -117,7 +117,7 @@ async fn run_tool_calls_test( println!("Prompt: {}", case.prompt); let mut session = engine.chat(model.clone(), ChatConfig::default()).await.unwrap(); - session.add_tool_descriptors(case.tools.iter().map(|definition| definition()).collect()).await.unwrap(); + session.add_tools(case.tools.iter().map(|definition| definition()).collect()).await.unwrap(); let mut messages = Vec::new(); if with_system_message { From 217888ac4fc7f8af119bc9773085f0168945b5bf Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:12:37 +0300 Subject: [PATCH 03/13] Add python tool calls implementation --- bindings/python/examples/tool_calls.py | 71 ++++++ bindings/python/pyproject.toml | 2 +- bindings/python/tests/__init__.py | 1 + bindings/python/tests/test_tools.py | 201 +++++++++++++++ bindings/python/uv.lock | 10 +- bindings/python/uzu/__init__.py | 2 + bindings/python/uzu/__init__.pyi | 5 + bindings/python/uzu/_tool.py | 312 ++++++++++++++++++++++++ crates/nagare/src/chat/bindings_pyo3.rs | 136 +++++++++++ crates/nagare/src/chat/mod.rs | 2 + crates/uzu/src/lib.rs | 3 + 11 files changed, 739 insertions(+), 6 deletions(-) create mode 100644 bindings/python/examples/tool_calls.py create mode 100644 bindings/python/tests/__init__.py create mode 100644 bindings/python/tests/test_tools.py create mode 100644 bindings/python/uzu/_tool.py create mode 100644 crates/nagare/src/chat/bindings_pyo3.rs diff --git a/bindings/python/examples/tool_calls.py b/bindings/python/examples/tool_calls.py new file mode 100644 index 000000000..eaea2b9db --- /dev/null +++ b/bindings/python/examples/tool_calls.py @@ -0,0 +1,71 @@ +import asyncio +from typing import Annotated + +from pydantic import BaseModel + +from uzu import ( + ChatConfig, + ChatMessage, + ChatReplyConfig, + Engine, + EngineConfig, + SamplingMethod, + SamplingPolicy, + uzu_tool_function, +) + + +class Coordinate(BaseModel): + """A geographic coordinate.""" + + latitude: Annotated[float, "Latitude in decimal degrees."] + longitude: Annotated[float, "Longitude in decimal degrees."] + + +@uzu_tool_function( + name="get_location", + description="Return the current location in coordinates" +) +def get_current_location() -> Coordinate: + return Coordinate(latitude=51.5074, longitude=-0.1278) + + +@uzu_tool_function +def get_current_temperature( + latitude: Annotated[float, "Latitude in decimal degrees."], + longitude: Annotated[float, "Longitude in decimal degrees."], +) -> float: + """Return the temperature at the provided coordinates.""" + _ = latitude, longitude + return 25.0 + + +async def main() -> None: + engine = await Engine.create(EngineConfig.create()) + model = await engine.model("mlx-community/Qwen3.5-9B-MLX-8bit") + if model is None: + raise RuntimeError("Model not found") + + async for update in (await engine.download(model)).iterator(): + print(f"Download progress: {update.progress}") + + session = await engine.chat(model, ChatConfig.create()) + await session.add_tool(get_current_location) + await session.add_tool(get_current_temperature) + + messages = [ + ChatMessage.system().with_text("You are a helpful assistant"), + ChatMessage.user().with_text("What temperature is it now at my location?"), + ] + config = ChatReplyConfig.create().with_sampling_policy( + SamplingPolicy.Custom(method=SamplingMethod.Greedy()) + ) + replies = await session.reply(messages, config) + if replies: + message = replies[-1].message + print(f"Reasoning: {message.reasoning or ''}") + print(f"Text: {message.text or ''}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/bindings/python/pyproject.toml b/bindings/python/pyproject.toml index 34ccb5a48..705e129f7 100644 --- a/bindings/python/pyproject.toml +++ b/bindings/python/pyproject.toml @@ -27,6 +27,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dynamic = ["version"] +dependencies = ["pydantic>=2.13.3"] [project.urls] Homepage = "https://trymirai.com" @@ -35,7 +36,6 @@ Issues = "https://github.com/trymirai/uzu/issues" [project.optional-dependencies] dev = ["pytest>=8.3.4", "ruff>=0.14.0"] -examples = ["pydantic>=2.13.3"] [tool.maturin] features = ["bindings-pyo3"] diff --git a/bindings/python/tests/__init__.py b/bindings/python/tests/__init__.py new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/bindings/python/tests/__init__.py @@ -0,0 +1 @@ + diff --git a/bindings/python/tests/test_tools.py b/bindings/python/tests/test_tools.py new file mode 100644 index 000000000..5ef9615d5 --- /dev/null +++ b/bindings/python/tests/test_tools.py @@ -0,0 +1,201 @@ +# ruff: noqa: SLF001 + +import asyncio +import json +from typing import Annotated, Literal + +import pytest +from pydantic import BaseModel, Field, ValidationError + +from uzu import ChatSession, UzuToolFunction, uzu_tool_function + + +class Coordinate(BaseModel): + """A geographic coordinate.""" + + latitude: Annotated[float, "Latitude in decimal degrees."] + longitude: float = Field(description="Longitude in decimal degrees.") + + +class ForecastRequest(BaseModel): + """A weather forecast request.""" + + location: Annotated[Coordinate, "Location to forecast."] + days: int = Field(description="Number of days ahead.") + + +class Node(BaseModel): + """A recursive tree node.""" + + value: Annotated[int, "Value stored in this node."] + children: list["Node"] = Field(default_factory=list) + + +class Weather(BaseModel): + """The current weather.""" + + temperature: float = Field(description="Temperature in degrees Celsius.") + summary: Annotated[str | None, "Optional human-readable summary."] + + +@uzu_tool_function +def get_forecast( + request: Annotated[ForecastRequest, "Request supplied to the forecast service."], + note: Annotated[str | None, "Required nullable note."], + unit: Annotated[Literal["celsius", "fahrenheit"], "Temperature unit."] = "celsius", +) -> Weather: + """Get the weather forecast for a location.""" + assert isinstance(request, ForecastRequest) + assert isinstance(request.location, Coordinate) + return Weather( + temperature=request.location.latitude + request.location.longitude, + summary=None if note is None else f"{note} ({unit})", + ) + + +@uzu_tool_function(name="sum_values", description="Add all node values.") +async def add_node_values(root: Node) -> int: + await asyncio.sleep(0) + return root.value + sum(child.value for child in root.children) + + +def _resolve_reference( + root: dict[str, object], + schema: dict[str, object], +) -> dict[str, object]: + reference = schema.get("$ref") + if not isinstance(reference, str): + return schema + + target: object = root + for raw_token in reference.removeprefix("#/").split("/"): + token = raw_token.replace("~1", "/").replace("~0", "~") + assert isinstance(target, dict) + target = target[token] + assert isinstance(target, dict) + return target + + +def test_decorator_builds_schema_from_function_and_pydantic_metadata() -> None: + assert isinstance(get_forecast, UzuToolFunction) + assert get_forecast.name == "get_forecast" + assert get_forecast.description == "Get the weather forecast for a location." + + parameters = get_forecast.parameters_schema + assert parameters is not None + assert parameters["required"] == ["request", "note"] + + properties = parameters["properties"] + request_property = properties["request"] + assert request_property["description"] == "Request supplied to the forecast service." + request = _resolve_reference(parameters, request_property) + location_property = request["properties"]["location"] + assert location_property["description"] == "Location to forecast." + location = _resolve_reference(parameters, location_property) + assert location["properties"]["latitude"]["description"] == "Latitude in decimal degrees." + assert location["properties"]["longitude"]["description"] == "Longitude in decimal degrees." + assert request["properties"]["days"]["description"] == "Number of days ahead." + + assert properties["note"]["anyOf"] == [{"type": "string"}, {"type": "null"}] + assert properties["note"]["description"] == "Required nullable note." + assert properties["unit"]["description"] == "Temperature unit." + assert properties["unit"]["default"] == "celsius" + + +def test_pydantic_return_schema_uses_docstrings_and_field_metadata() -> None: + schema = get_forecast.return_schema + assert schema is not None + assert schema["description"] == "The current weather." + assert schema["properties"]["temperature"]["description"] == "Temperature in degrees Celsius." + assert schema["properties"]["summary"]["description"] == "Optional human-readable summary." + assert schema["required"] == ["temperature", "summary"] + + +def test_sync_tool_constructs_and_serializes_pydantic_models() -> None: + result = get_forecast._invoke_json( + json.dumps( + { + "request": { + "location": { + "latitude": 51.5, + "longitude": -0.1, + }, + "days": 2, + }, + "note": None, + } + ) + ) + + assert isinstance(result, str) + assert json.loads(result) == { + "temperature": 51.4, + "summary": None, + } + + +def test_required_nullable_parameter_cannot_be_omitted() -> None: + with pytest.raises(ValidationError, match="note"): + get_forecast._invoke_json( + json.dumps( + { + "request": { + "location": { + "latitude": 51.5, + "longitude": -0.1, + }, + "days": 2, + } + } + ) + ) + + +def test_async_tool_constructs_recursive_pydantic_model() -> None: + async def run() -> None: + result = add_node_values._invoke_json( + json.dumps( + { + "root": { + "value": 1, + "children": [ + { + "value": 2, + "children": [], + }, + { + "value": 3, + }, + ], + } + } + ) + ) + assert not isinstance(result, str) + assert json.loads(await result) == 6 + + asyncio.run(run()) + + parameters = add_node_values.parameters_schema + assert parameters is not None + root_property = parameters["properties"]["root"] + root = _resolve_reference(parameters, root_property) + assert root["properties"]["children"]["items"]["$ref"] == "#/$defs/Node" + + +def test_configured_decorator_overrides_name_and_description() -> None: + assert add_node_values.name == "sum_values" + assert add_node_values.description == "Add all node values." + + +def test_chat_session_exposes_tool_registration_methods() -> None: + assert callable(ChatSession.add_tool) + assert callable(ChatSession.add_tools) + + +def test_native_definition_matches_decorator_metadata() -> None: + definition = get_forecast._native_definition() + assert definition.name == get_forecast.name + assert definition.description == get_forecast.description + assert json.loads(definition.parameters.json) == get_forecast.parameters_schema + assert json.loads(definition.return_definition.json) == get_forecast.return_schema diff --git a/bindings/python/uv.lock b/bindings/python/uv.lock index 23e45458d..053abd0e5 100644 --- a/bindings/python/uv.lock +++ b/bindings/python/uv.lock @@ -211,20 +211,20 @@ wheels = [ [[package]] name = "uzu" source = { editable = "." } +dependencies = [ + { name = "pydantic" }, +] [package.optional-dependencies] dev = [ { name = "pytest" }, { name = "ruff" }, ] -examples = [ - { name = "pydantic" }, -] [package.metadata] requires-dist = [ - { name = "pydantic", marker = "extra == 'examples'", specifier = ">=2.13.3" }, + { name = "pydantic", specifier = ">=2.13.3" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.3.4" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.14.0" }, ] -provides-extras = ["dev", "examples"] +provides-extras = ["dev"] diff --git a/bindings/python/uzu/__init__.py b/bindings/python/uzu/__init__.py index 90df4cc80..0588d44c6 100644 --- a/bindings/python/uzu/__init__.py +++ b/bindings/python/uzu/__init__.py @@ -1,6 +1,8 @@ import sys as _sys from . import uzu +from ._tool import UzuToolFunction as UzuToolFunction +from ._tool import uzu_tool_function as uzu_tool_function from .uzu import * # noqa: F403 _sys.modules.pop("uzu.uzu", None) diff --git a/bindings/python/uzu/__init__.pyi b/bindings/python/uzu/__init__.pyi index 75cbaa798..9f9d16587 100644 --- a/bindings/python/uzu/__init__.pyi +++ b/bindings/python/uzu/__init__.pyi @@ -5,6 +5,7 @@ import builtins import collections.abc import enum import typing +from uzu._tool import UzuToolFunction, uzu_tool_function __all__ = [ "CancelToken", "ChatConfig", @@ -84,7 +85,9 @@ __all__ = [ "ToolFunction", "ToolNamespace", "TranslationPayload", + "UzuToolFunction", "Value", + "uzu_tool_function", ] @typing.final @@ -414,6 +417,8 @@ class ChatRole: @typing.final class ChatSession: + def add_tool(self, tool: UzuToolFunction[..., typing.Any]) -> collections.abc.Awaitable[None]: ... + def add_tools(self, tools: collections.abc.Sequence[UzuToolFunction[..., typing.Any]]) -> collections.abc.Awaitable[None]: ... def state(self) -> collections.abc.Awaitable[ChatSessionState]: ... def messages(self) -> collections.abc.Awaitable[builtins.list[ChatMessage]]: ... def reset(self) -> collections.abc.Awaitable[None]: ... diff --git a/bindings/python/uzu/_tool.py b/bindings/python/uzu/_tool.py new file mode 100644 index 000000000..c5146ac56 --- /dev/null +++ b/bindings/python/uzu/_tool.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +import collections.abc +import functools +import inspect +import json +import types +import typing +from typing import Annotated, Any, Literal, overload + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, create_model + +_EMPTY = inspect.Parameter.empty +_NONE_TYPE = type(None) + + +class UzuToolFunction[**P, R]: + """A Python function and its Uzu tool definition.""" + + def __init__( + self, + function: collections.abc.Callable[P, R], + *, + name: str | None = None, + description: str | None = None, + ) -> None: + if not callable(function): + raise TypeError("tool must be callable") + + inferred_name = getattr(function, "__name__", None) + if name is None: + if not inferred_name or inferred_name == "": + raise TypeError("a tool name is required for lambdas and unnamed callables") + name = inferred_name + + self._function = function + self.name = name + self.description = description if description is not None else inspect.getdoc(function) or "" + self._signature = inspect.signature(function) + + try: + annotations = typing.get_type_hints(function, include_extras=True) + except (NameError, TypeError) as error: + raise TypeError(f"unable to resolve annotations for tool {name!r}: {error}") from error + + fields: dict[str, tuple[object, object]] = {} + for parameter in self._signature.parameters.values(): + if parameter.kind in {inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.VAR_POSITIONAL}: + raise TypeError(f"tool {name!r} cannot use positional-only or variadic positional parameters") + if parameter.kind is inspect.Parameter.VAR_KEYWORD: + raise TypeError(f"tool {name!r} cannot use variadic keyword parameters") + + annotation = annotations.get(parameter.name, parameter.annotation) + if annotation is _EMPTY: + raise TypeError(f"parameter {parameter.name!r} of tool {name!r} must have a type annotation") + + if annotation_description := _annotation_description(annotation): + annotation = Annotated[annotation, Field(description=annotation_description)] + default = ... if parameter.default is _EMPTY else parameter.default + fields[parameter.name] = (annotation, default) + + self._arguments_model: type[BaseModel] | None = None + self.parameters_schema: dict[str, object] | None = None + if fields: + self._arguments_model = create_model( + f"{name}_arguments", + __config__=ConfigDict(extra="forbid"), + **fields, + ) + self.parameters_schema = self._arguments_model.model_json_schema(mode="validation") + _apply_annotated_descriptions( + self._arguments_model, + self.parameters_schema, + self.parameters_schema, + mode="validation", + visited=set(), + ) + + return_annotation = annotations.get("return", self._signature.return_annotation) + self._return_adapter = TypeAdapter(Any if return_annotation is _EMPTY else return_annotation) + self.return_schema: dict[str, object] | None = None + if return_annotation not in {_EMPTY, None, _NONE_TYPE}: + self.return_schema = self._return_adapter.json_schema(mode="serialization") + _apply_annotated_descriptions( + return_annotation, + self.return_schema, + self.return_schema, + mode="serialization", + visited=set(), + ) + + functools.update_wrapper(self, function) + + def __call__( + self, + *args: P.args, + **kwargs: P.kwargs, + ) -> R: + return self._function(*args, **kwargs) + + def _native_definition(self) -> object: + # Imported lazily because the extension imports this module while initializing. + from . import ToolFunction, Value + + parameters = None if self.parameters_schema is None else Value(json.dumps(self.parameters_schema)) + return_definition = None if self.return_schema is None else Value(json.dumps(self.return_schema)) + return ToolFunction( + name=self.name, + description=self.description, + parameters=parameters, + return_definition=return_definition, + ) + + def _invoke_json(self, arguments_json: str) -> str | collections.abc.Awaitable[str]: + if self._arguments_model is None: + arguments = json.loads(arguments_json) + if arguments not in ({}, None): + raise TypeError(f"tool {self.name!r} does not accept arguments") + kwargs: dict[str, object] = {} + else: + arguments = self._arguments_model.model_validate_json(arguments_json) + kwargs = {name: getattr(arguments, name) for name in self._signature.parameters} + + result = self._function(**kwargs) + if inspect.isawaitable(result): + + async def finish() -> str: + return self._encode_result(await result) + + return finish() + return self._encode_result(result) + + def _encode_result(self, result: object) -> str: + validated = self._return_adapter.validate_python(result) + return self._return_adapter.dump_json(validated).decode() + + +@overload +def uzu_tool_function[**P, R]( + function: collections.abc.Callable[P, R], + *, + name: str | None = None, + description: str | None = None, +) -> UzuToolFunction[P, R]: ... + + +@overload +def uzu_tool_function[**P, R]( + function: None = None, + *, + name: str | None = None, + description: str | None = None, +) -> collections.abc.Callable[[collections.abc.Callable[P, R]], UzuToolFunction[P, R]]: ... + + +def uzu_tool_function[**P, R]( + function: collections.abc.Callable[P, R] | None = None, + *, + name: str | None = None, + description: str | None = None, +) -> UzuToolFunction[P, R] | collections.abc.Callable[[collections.abc.Callable[P, R]], UzuToolFunction[P, R]]: + """Create an Uzu tool from an annotated Python function.""" + + def decorate(target: collections.abc.Callable[P, R]) -> UzuToolFunction[P, R]: + return UzuToolFunction(target, name=name, description=description) + + return decorate if function is None else decorate(function) + + +def _annotation_description(annotation: object) -> str | None: + if typing.get_origin(annotation) is not typing.Annotated: + return None + + description: str | None = None + for metadata in typing.get_args(annotation)[1:]: + candidate = metadata if isinstance(metadata, str) else getattr(metadata, "description", None) + if candidate: + description = str(candidate) + return description + + +def _apply_annotated_descriptions( # noqa: PLR0911 + annotation: object, + schema: dict[str, object], + root_schema: dict[str, object], + *, + mode: Literal["validation", "serialization"], + visited: set[tuple[object, int]], +) -> None: + origin = typing.get_origin(annotation) + arguments = typing.get_args(annotation) + + if origin is typing.Annotated: + if description := _annotation_description(annotation): + schema["description"] = description + _apply_annotated_descriptions( + arguments[0], + schema, + root_schema, + mode=mode, + visited=visited, + ) + return + + if origin in {typing.Union, types.UnionType}: + branches = schema.get("anyOf") + if isinstance(branches, list): + for member, branch in zip(arguments, branches, strict=False): + if isinstance(branch, dict): + _apply_annotated_descriptions( + member, + branch, + root_schema, + mode=mode, + visited=visited, + ) + return + + if origin in {list, set, frozenset, collections.abc.Sequence}: + items = schema.get("items") + if arguments and isinstance(items, dict): + _apply_annotated_descriptions( + arguments[0], + items, + root_schema, + mode=mode, + visited=visited, + ) + return + + if origin is tuple: + items = schema.get("items") + if isinstance(items, dict) and arguments: + _apply_annotated_descriptions( + arguments[0], + items, + root_schema, + mode=mode, + visited=visited, + ) + prefix_items = schema.get("prefixItems") + if isinstance(prefix_items, list): + for member, item in zip(arguments, prefix_items, strict=False): + if isinstance(item, dict): + _apply_annotated_descriptions( + member, + item, + root_schema, + mode=mode, + visited=visited, + ) + return + + if origin in {dict, collections.abc.Mapping}: + values = schema.get("additionalProperties") + if len(arguments) == 2 and isinstance(values, dict): + _apply_annotated_descriptions( + arguments[1], + values, + root_schema, + mode=mode, + visited=visited, + ) + return + + if not inspect.isclass(annotation) or not issubclass(annotation, BaseModel): + return + + model_schema = _resolve_local_reference(schema, root_schema) + marker = (annotation, id(model_schema)) + if marker in visited: + return + visited.add(marker) + + properties = model_schema.get("properties") + if not isinstance(properties, dict): + return + + try: + annotations = typing.get_type_hints(annotation, include_extras=True) + except (NameError, TypeError): + annotations = {name: field.annotation for name, field in annotation.model_fields.items()} + + for name, field in annotation.model_fields.items(): + field_annotation = annotations.get(name, field.annotation) + alias = field.serialization_alias if mode == "serialization" else field.validation_alias + field_schema = properties.get(alias if isinstance(alias, str) else name) + if isinstance(field_schema, dict): + _apply_annotated_descriptions( + field_annotation, + field_schema, + root_schema, + mode=mode, + visited=visited, + ) + + +def _resolve_local_reference( + schema: dict[str, object], + root_schema: dict[str, object], +) -> dict[str, object]: + reference = schema.get("$ref") + if not isinstance(reference, str) or not reference.startswith("#/"): + return schema + + target: object = root_schema + for raw_token in reference[2:].split("/"): + token = raw_token.replace("~1", "/").replace("~0", "~") + if not isinstance(target, dict) or token not in target: + return schema + target = target[token] + return target if isinstance(target, dict) else schema diff --git a/crates/nagare/src/chat/bindings_pyo3.rs b/crates/nagare/src/chat/bindings_pyo3.rs new file mode 100644 index 000000000..8e6a770d9 --- /dev/null +++ b/crates/nagare/src/chat/bindings_pyo3.rs @@ -0,0 +1,136 @@ +use std::{future::Future, pin::Pin, sync::Arc}; + +use pyo3::{ + Bound, Py, PyAny, PyErr, PyResult, Python, + types::{PyAnyMethods, PyString}, +}; +use shoji::types::basic::{ToolFunction, Value}; + +use super::ChatSession; +use crate::tool::func_def::{ErrorFuture, ToolDescriptor}; + +type PythonFuture = Pin>> + Send>>; + +enum PythonCall { + Ready(Py), + Awaitable(PythonFuture), +} + +#[pyo3_stub_gen::derive::gen_stub_pymethods] +#[pyo3::pymethods] +impl ChatSession { + #[pyo3(name = "add_tool")] + #[gen_stub( + override_return_type( + type_repr = "collections.abc.Awaitable[None]", + imports = ("collections.abc") + ) + )] + fn add_tool_bindings_pyo3<'py>( + &self, + py: Python<'py>, + #[gen_stub( + override_type( + type_repr = "UzuToolFunction[..., typing.Any]", + imports = ("typing") + ) + )] + tool: Py, + ) -> PyResult> { + let task_locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; + let descriptor = descriptor_from_python(py, tool, task_locals)?; + let mut session = self.clone(); + + pyo3_async_runtimes::tokio::future_into_py( + py, + async move { session.add_tool(descriptor).await.map_err(Into::into) }, + ) + } + + #[pyo3(name = "add_tools")] + #[gen_stub( + override_return_type( + type_repr = "collections.abc.Awaitable[None]", + imports = ("collections.abc") + ) + )] + fn add_tools_bindings_pyo3<'py>( + &self, + py: Python<'py>, + #[gen_stub( + override_type( + type_repr = "collections.abc.Sequence[UzuToolFunction[..., typing.Any]]", + imports = ("collections.abc", "typing") + ) + )] + tools: Vec>, + ) -> PyResult> { + let task_locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; + let descriptors = tools + .into_iter() + .map(|tool| descriptor_from_python(py, tool, task_locals.clone())) + .collect::>>()?; + let mut session = self.clone(); + + pyo3_async_runtimes::tokio::future_into_py(py, async move { + session.add_tools(descriptors).await.map_err(Into::into) + }) + } +} + +fn descriptor_from_python( + py: Python<'_>, + tool: Py, + task_locals: pyo3_async_runtimes::TaskLocals, +) -> PyResult { + let tool_function: ToolFunction = tool.bind(py).call_method0("_native_definition")?.extract()?; + let tool = Arc::new(tool); + Ok(ToolDescriptor::new( + tool_function.name, + tool_function.description, + tool_function.parameters, + tool_function.return_definition, + Box::new(move |arguments| { + let tool = Arc::clone(&tool); + let task_locals = task_locals.clone(); + Box::new(call_python_tool(tool, task_locals, arguments)) + }), + )) +} + +async fn call_python_tool( + tool: Arc>, + task_locals: pyo3_async_runtimes::TaskLocals, + arguments: Value, +) -> Result { + let call = Python::attach(|py| -> PyResult { + let result = tool.bind(py).call_method1("_invoke_json", (arguments.json,))?; + if result.is_instance_of::() { + return Ok(PythonCall::Ready(result.unbind())); + } + + let inspect = pyo3::types::PyModule::import(py, "inspect")?; + if inspect.call_method1("isawaitable", (&result,))?.is_truthy()? { + let future = pyo3_async_runtimes::into_future_with_locals(&task_locals, result)?; + Ok(PythonCall::Awaitable(Box::pin(future))) + } else { + Err(pyo3::exceptions::PyTypeError::new_err( + "decorated tool invocation must return JSON text or an awaitable producing JSON text", + )) + } + }) + .map_err(python_error)?; + + let result = match call { + PythonCall::Ready(result) => result, + PythonCall::Awaitable(future) => future.await.map_err(python_error)?, + }; + + let json = Python::attach(|py| result.extract::(py)).map_err(python_error)?; + let value = serde_json::from_str::(&json)?; + Ok(Value::from(value)) +} + +fn python_error(error: PyErr) -> ErrorFuture { + error.to_string().into() +} diff --git a/crates/nagare/src/chat/mod.rs b/crates/nagare/src/chat/mod.rs index 0e0575161..3d2640bea 100644 --- a/crates/nagare/src/chat/mod.rs +++ b/crates/nagare/src/chat/mod.rs @@ -1,3 +1,5 @@ +#[cfg(feature = "bindings-pyo3")] +mod bindings_pyo3; mod error; pub mod message; pub mod token; diff --git a/crates/uzu/src/lib.rs b/crates/uzu/src/lib.rs index db18ad8e8..248ff279b 100644 --- a/crates/uzu/src/lib.rs +++ b/crates/uzu/src/lib.rs @@ -28,6 +28,9 @@ fn uzu(m: &pyo3::Bound<'_, pyo3::types::PyModule>) -> pyo3::PyResult<()> { #[cfg(feature = "bindings-pyo3")] pyo3_stub_gen::define_stub_info_gatherer!(pyo3_bindings_annotations); +#[cfg(feature = "bindings-pyo3")] +pyo3_stub_gen::reexport_module_members!("uzu" from "uzu._tool"; "UzuToolFunction", "uzu_tool_function"); + #[cfg(all(feature = "capability-cli", not(target_family = "wasm")))] pub mod cli; #[cfg(not(target_family = "wasm"))] From 389fb700e1da596fe568906d5fe4d61631d9b6e6 Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:40:25 +0300 Subject: [PATCH 04/13] Resolve task locals when the callback runs --- crates/nagare/src/chat/bindings_pyo3.rs | 30 ++++++++++++++++--------- crates/nagare/src/chat/mod.rs | 6 ++++- 2 files changed, 24 insertions(+), 12 deletions(-) diff --git a/crates/nagare/src/chat/bindings_pyo3.rs b/crates/nagare/src/chat/bindings_pyo3.rs index 8e6a770d9..294d5fb7f 100644 --- a/crates/nagare/src/chat/bindings_pyo3.rs +++ b/crates/nagare/src/chat/bindings_pyo3.rs @@ -37,8 +37,7 @@ impl ChatSession { )] tool: Py, ) -> PyResult> { - let task_locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; - let descriptor = descriptor_from_python(py, tool, task_locals)?; + let descriptor = descriptor_from_python(py, tool)?; let mut session = self.clone(); pyo3_async_runtimes::tokio::future_into_py( @@ -65,11 +64,8 @@ impl ChatSession { )] tools: Vec>, ) -> PyResult> { - let task_locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; - let descriptors = tools - .into_iter() - .map(|tool| descriptor_from_python(py, tool, task_locals.clone())) - .collect::>>()?; + let descriptors = + tools.into_iter().map(|tool| descriptor_from_python(py, tool)).collect::>>()?; let mut session = self.clone(); pyo3_async_runtimes::tokio::future_into_py(py, async move { @@ -81,7 +77,6 @@ impl ChatSession { fn descriptor_from_python( py: Python<'_>, tool: Py, - task_locals: pyo3_async_runtimes::TaskLocals, ) -> PyResult { let tool_function: ToolFunction = tool.bind(py).call_method0("_native_definition")?.extract()?; let tool = Arc::new(tool); @@ -92,15 +87,13 @@ fn descriptor_from_python( tool_function.return_definition, Box::new(move |arguments| { let tool = Arc::clone(&tool); - let task_locals = task_locals.clone(); - Box::new(call_python_tool(tool, task_locals, arguments)) + Box::new(call_python_tool(tool, arguments)) }), )) } async fn call_python_tool( tool: Arc>, - task_locals: pyo3_async_runtimes::TaskLocals, arguments: Value, ) -> Result { let call = Python::attach(|py| -> PyResult { @@ -111,6 +104,8 @@ async fn call_python_tool( let inspect = pyo3::types::PyModule::import(py, "inspect")?; if inspect.call_method1("isawaitable", (&result,))?.is_truthy()? { + // The turn task carries the locals captured by the Python reply invocation. + let task_locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; let future = pyo3_async_runtimes::into_future_with_locals(&task_locals, result)?; Ok(PythonCall::Awaitable(Box::pin(future))) } else { @@ -134,3 +129,16 @@ async fn call_python_tool( fn python_error(error: PyErr) -> ErrorFuture { error.to_string().into() } + +pub(super) fn spawn_with_current_task_locals(future: F) -> tokio::task::JoinHandle<()> +where + F: Future + Send + 'static, +{ + // `future_into_py` scopes the exported reply future with its Python loop and + // context. Preserve that scope when the core session detaches the turn. + let task_locals = Python::try_attach(|py| pyo3_async_runtimes::tokio::get_current_locals(py).ok()).flatten(); + match task_locals { + Some(task_locals) => tokio::spawn(pyo3_async_runtimes::tokio::scope(task_locals, future)), + None => tokio::spawn(future), + } +} diff --git a/crates/nagare/src/chat/mod.rs b/crates/nagare/src/chat/mod.rs index 3d2640bea..80854b52a 100644 --- a/crates/nagare/src/chat/mod.rs +++ b/crates/nagare/src/chat/mod.rs @@ -573,7 +573,11 @@ impl ChatSession { let cancel_token = cancel_token_to_return.inner().clone(); let (sender, receiver) = mpsc::unbounded_channel::, ChatSessionError>>(); let session = self.clone(); - tokio::spawn(async move { session.execute_turn(sender, input, config, cancel_token).await }); + let turn = async move { session.execute_turn(sender, input, config, cancel_token).await }; + #[cfg(feature = "bindings-pyo3")] + bindings_pyo3::spawn_with_current_task_locals(turn); + #[cfg(not(feature = "bindings-pyo3"))] + tokio::spawn(turn); ChatSessionStream { receiver: Arc::new(Mutex::new(receiver)), cancel_token: cancel_token_to_return, From ebea6aec0840a832ad904412553b5f223bade022 Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:47:55 +0300 Subject: [PATCH 05/13] Fixed synchronous callback dispatch --- bindings/python/tests/test_tools.py | 42 +++++++++++++++++++++++-- bindings/python/uzu/_tool.py | 4 +++ crates/nagare/src/chat/bindings_pyo3.rs | 39 +++++------------------ 3 files changed, 51 insertions(+), 34 deletions(-) diff --git a/bindings/python/tests/test_tools.py b/bindings/python/tests/test_tools.py index 5ef9615d5..a90a9cf02 100644 --- a/bindings/python/tests/test_tools.py +++ b/bindings/python/tests/test_tools.py @@ -1,7 +1,9 @@ # ruff: noqa: SLF001 import asyncio +import contextvars import json +import threading from typing import Annotated, Literal import pytest @@ -9,6 +11,8 @@ from uzu import ChatSession, UzuToolFunction, uzu_tool_function +_TOOL_CONTEXT = contextvars.ContextVar[str]("tool_context") + class Coordinate(BaseModel): """A geographic coordinate.""" @@ -59,6 +63,15 @@ async def add_node_values(root: Node) -> int: return root.value + sum(child.value for child in root.children) +@uzu_tool_function +def current_runtime_state() -> dict[str, int | str]: + return { + "context": _TOOL_CONTEXT.get(), + "loop": id(asyncio.get_running_loop()), + "thread": threading.get_ident(), + } + + def _resolve_reference( root: dict[str, object], schema: dict[str, object], @@ -153,7 +166,7 @@ def test_required_nullable_parameter_cannot_be_omitted() -> None: def test_async_tool_constructs_recursive_pydantic_model() -> None: async def run() -> None: - result = add_node_values._invoke_json( + result = await add_node_values._invoke_json_on_loop( json.dumps( { "root": { @@ -171,8 +184,7 @@ async def run() -> None: } ) ) - assert not isinstance(result, str) - assert json.loads(await result) == 6 + assert json.loads(result) == 6 asyncio.run(run()) @@ -183,6 +195,30 @@ async def run() -> None: assert root["properties"]["children"]["items"]["$ref"] == "#/$defs/Node" +def test_sync_tool_trampoline_runs_on_python_loop_and_context() -> None: + async def run() -> None: + loop = asyncio.get_running_loop() + loop_thread = threading.get_ident() + token = _TOOL_CONTEXT.set("reply-context") + try: + worker_thread, invocation = await asyncio.to_thread( + lambda: ( + threading.get_ident(), + current_runtime_state._invoke_json_on_loop("{}"), + ) + ) + assert worker_thread != loop_thread + assert json.loads(await invocation) == { + "context": "reply-context", + "loop": id(loop), + "thread": loop_thread, + } + finally: + _TOOL_CONTEXT.reset(token) + + asyncio.run(run()) + + def test_configured_decorator_overrides_name_and_description() -> None: assert add_node_values.name == "sum_values" assert add_node_values.description == "Add all node values." diff --git a/bindings/python/uzu/_tool.py b/bindings/python/uzu/_tool.py index c5146ac56..02cba9e69 100644 --- a/bindings/python/uzu/_tool.py +++ b/bindings/python/uzu/_tool.py @@ -130,6 +130,10 @@ async def finish() -> str: return finish() return self._encode_result(result) + async def _invoke_json_on_loop(self, arguments_json: str) -> str: + result = self._invoke_json(arguments_json) + return await result if inspect.isawaitable(result) else result + def _encode_result(self, result: object) -> str: validated = self._return_adapter.validate_python(result) return self._return_adapter.dump_json(validated).decode() diff --git a/crates/nagare/src/chat/bindings_pyo3.rs b/crates/nagare/src/chat/bindings_pyo3.rs index 294d5fb7f..dbd20ef93 100644 --- a/crates/nagare/src/chat/bindings_pyo3.rs +++ b/crates/nagare/src/chat/bindings_pyo3.rs @@ -1,9 +1,6 @@ use std::{future::Future, pin::Pin, sync::Arc}; -use pyo3::{ - Bound, Py, PyAny, PyErr, PyResult, Python, - types::{PyAnyMethods, PyString}, -}; +use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, types::PyAnyMethods}; use shoji::types::basic::{ToolFunction, Value}; use super::ChatSession; @@ -11,11 +8,6 @@ use crate::tool::func_def::{ErrorFuture, ToolDescriptor}; type PythonFuture = Pin>> + Send>>; -enum PythonCall { - Ready(Py), - Awaitable(PythonFuture), -} - #[pyo3_stub_gen::derive::gen_stub_pymethods] #[pyo3::pymethods] impl ChatSession { @@ -96,30 +88,15 @@ async fn call_python_tool( tool: Arc>, arguments: Value, ) -> Result { - let call = Python::attach(|py| -> PyResult { - let result = tool.bind(py).call_method1("_invoke_json", (arguments.json,))?; - if result.is_instance_of::() { - return Ok(PythonCall::Ready(result.unbind())); - } - - let inspect = pyo3::types::PyModule::import(py, "inspect")?; - if inspect.call_method1("isawaitable", (&result,))?.is_truthy()? { - // The turn task carries the locals captured by the Python reply invocation. - let task_locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; - let future = pyo3_async_runtimes::into_future_with_locals(&task_locals, result)?; - Ok(PythonCall::Awaitable(Box::pin(future))) - } else { - Err(pyo3::exceptions::PyTypeError::new_err( - "decorated tool invocation must return JSON text or an awaitable producing JSON text", - )) - } + let future = Python::attach(|py| -> PyResult { + // Creating this coroutine does not invoke the tool. Its body runs on the + // reply's Python loop and context once `into_future_with_locals` schedules it. + let task_locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; + let invocation = tool.bind(py).call_method1("_invoke_json_on_loop", (arguments.json,))?; + Ok(Box::pin(pyo3_async_runtimes::into_future_with_locals(&task_locals, invocation)?)) }) .map_err(python_error)?; - - let result = match call { - PythonCall::Ready(result) => result, - PythonCall::Awaitable(future) => future.await.map_err(python_error)?, - }; + let result = future.await.map_err(python_error)?; let json = Python::attach(|py| result.extract::(py)).map_err(python_error)?; let value = serde_json::from_str::(&json)?; From 31168cd33c14dce36c1f7b5d536a468a9a403cec Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 18:57:24 +0300 Subject: [PATCH 06/13] Preserve parameters that collide with BaseModel --- bindings/python/tests/test_tools.py | 27 +++++++++++++++++++++++++++ bindings/python/uzu/_tool.py | 20 ++++++++++++++++---- 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/bindings/python/tests/test_tools.py b/bindings/python/tests/test_tools.py index a90a9cf02..892f0532b 100644 --- a/bindings/python/tests/test_tools.py +++ b/bindings/python/tests/test_tools.py @@ -72,6 +72,17 @@ def current_runtime_state() -> dict[str, int | str]: } +@uzu_tool_function +def reserved_parameter_names( + model_config: Annotated[int, "Configuration number."], + model_dump: Annotated[str, "Dump label."], +) -> dict[str, int | str]: + return { + "model_config": model_config, + "model_dump": model_dump, + } + + def _resolve_reference( root: dict[str, object], schema: dict[str, object], @@ -219,6 +230,22 @@ async def run() -> None: asyncio.run(run()) +def test_parameters_can_use_reserved_base_model_names() -> None: + parameters = reserved_parameter_names.parameters_schema + assert parameters is not None + assert parameters["required"] == ["model_config", "model_dump"] + assert set(parameters["properties"]) == {"model_config", "model_dump"} + assert parameters["properties"]["model_config"]["description"] == "Configuration number." + assert parameters["properties"]["model_dump"]["description"] == "Dump label." + + result = reserved_parameter_names._invoke_json('{"model_config":7,"model_dump":"ready"}') + assert isinstance(result, str) + assert json.loads(result) == { + "model_config": 7, + "model_dump": "ready", + } + + def test_configured_decorator_overrides_name_and_description() -> None: assert add_node_values.name == "sum_values" assert add_node_values.description == "Add all node values." diff --git a/bindings/python/uzu/_tool.py b/bindings/python/uzu/_tool.py index 02cba9e69..e5054b065 100644 --- a/bindings/python/uzu/_tool.py +++ b/bindings/python/uzu/_tool.py @@ -44,6 +44,7 @@ def __init__( raise TypeError(f"unable to resolve annotations for tool {name!r}: {error}") from error fields: dict[str, tuple[object, object]] = {} + self._argument_fields: dict[str, str] = {} for parameter in self._signature.parameters.values(): if parameter.kind in {inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.VAR_POSITIONAL}: raise TypeError(f"tool {name!r} cannot use positional-only or variadic positional parameters") @@ -54,10 +55,18 @@ def __init__( if annotation is _EMPTY: raise TypeError(f"parameter {parameter.name!r} of tool {name!r} must have a type annotation") - if annotation_description := _annotation_description(annotation): - annotation = Annotated[annotation, Field(description=annotation_description)] + field_metadata = ( + Field(alias=parameter.name, description=annotation_description) + if (annotation_description := _annotation_description(annotation)) + else Field(alias=parameter.name) + ) + annotation = Annotated[annotation, field_metadata] default = ... if parameter.default is _EMPTY else parameter.default - fields[parameter.name] = (annotation, default) + internal_name = f"uzu_argument_{len(fields)}" + while hasattr(BaseModel, internal_name): + internal_name += "_" + fields[internal_name] = (annotation, default) + self._argument_fields[parameter.name] = internal_name self._arguments_model: type[BaseModel] | None = None self.parameters_schema: dict[str, object] | None = None @@ -119,7 +128,10 @@ def _invoke_json(self, arguments_json: str) -> str | collections.abc.Awaitable[s kwargs: dict[str, object] = {} else: arguments = self._arguments_model.model_validate_json(arguments_json) - kwargs = {name: getattr(arguments, name) for name in self._signature.parameters} + kwargs = { + parameter_name: getattr(arguments, internal_name) + for parameter_name, internal_name in self._argument_fields.items() + } result = self._function(**kwargs) if inspect.isawaitable(result): From cc035c5d583d0757b01390d27d487d8fe482e9ad Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 19:14:39 +0300 Subject: [PATCH 07/13] Fixed the unhashable return annotation issue --- bindings/python/tests/test_tools.py | 11 +++++++++++ bindings/python/uzu/_tool.py | 2 +- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/bindings/python/tests/test_tools.py b/bindings/python/tests/test_tools.py index 892f0532b..1eff32dde 100644 --- a/bindings/python/tests/test_tools.py +++ b/bindings/python/tests/test_tools.py @@ -83,6 +83,11 @@ def reserved_parameter_names( } +@uzu_tool_function +def unhashable_return_metadata() -> Annotated[int, {"tag": "x"}]: + return 42 + + def _resolve_reference( root: dict[str, object], schema: dict[str, object], @@ -246,6 +251,12 @@ def test_parameters_can_use_reserved_base_model_names() -> None: } +def test_return_annotation_can_have_unhashable_metadata() -> None: + assert unhashable_return_metadata.return_schema is not None + assert unhashable_return_metadata.return_schema["type"] == "integer" + assert unhashable_return_metadata._invoke_json("{}") == "42" + + def test_configured_decorator_overrides_name_and_description() -> None: assert add_node_values.name == "sum_values" assert add_node_values.description == "Add all node values." diff --git a/bindings/python/uzu/_tool.py b/bindings/python/uzu/_tool.py index e5054b065..411562770 100644 --- a/bindings/python/uzu/_tool.py +++ b/bindings/python/uzu/_tool.py @@ -88,7 +88,7 @@ def __init__( return_annotation = annotations.get("return", self._signature.return_annotation) self._return_adapter = TypeAdapter(Any if return_annotation is _EMPTY else return_annotation) self.return_schema: dict[str, object] | None = None - if return_annotation not in {_EMPTY, None, _NONE_TYPE}: + if return_annotation is not _EMPTY and return_annotation is not None and return_annotation is not _NONE_TYPE: self.return_schema = self._return_adapter.json_schema(mode="serialization") _apply_annotated_descriptions( return_annotation, From b7d2024d20fbd8eeb94035798ee94d889a91c6ba Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 19:38:02 +0300 Subject: [PATCH 08/13] Fixes --- bindings/python/tests/test_tools.py | 50 +++++++++++++++++++++++ bindings/python/uzu/_tool.py | 31 ++++++++++++++- crates/nagare/src/chat/bindings_pyo3.rs | 53 +++++++++++++++++++++---- 3 files changed, 126 insertions(+), 8 deletions(-) diff --git a/bindings/python/tests/test_tools.py b/bindings/python/tests/test_tools.py index 1eff32dde..36b584d38 100644 --- a/bindings/python/tests/test_tools.py +++ b/bindings/python/tests/test_tools.py @@ -257,6 +257,56 @@ def test_return_annotation_can_have_unhashable_metadata() -> None: assert unhashable_return_metadata._invoke_json("{}") == "42" +def test_wrapper_metadata_does_not_replace_tool_state() -> None: + def target(value: int) -> int: + return value + + target.__dict__.update( + { + "_function": lambda value: value + 100, + "description": "Wrong description.", + "name": "wrong_name", + "parameters_schema": {"wrong": True}, + } + ) + tool = uzu_tool_function(target, name="configured_name", description="Configured description.") + + assert tool.name == "configured_name" + assert tool.description == "Configured description." + assert tool.parameters_schema is not None + assert set(tool.parameters_schema["properties"]) == {"value"} + assert tool._invoke_json('{"value":5}') == "5" + + +def test_cancelling_invocation_cancels_async_python_tool() -> None: + async def run() -> None: + started = asyncio.Event() + cleaned_up = asyncio.Event() + side_effect = False + + @uzu_tool_function + async def waiting_tool() -> None: + nonlocal side_effect + started.set() + try: + await asyncio.Event().wait() + side_effect = True + finally: + cleaned_up.set() + + invocation = waiting_tool._new_json_invocation("{}") + task = asyncio.create_task(invocation.run()) + await started.wait() + await asyncio.to_thread(invocation.cancel) + + with pytest.raises(asyncio.CancelledError): + await task + await asyncio.wait_for(cleaned_up.wait(), timeout=1) + assert not side_effect + + asyncio.run(run()) + + def test_configured_decorator_overrides_name_and_description() -> None: assert add_node_values.name == "sum_values" assert add_node_values.description == "Add all node values." diff --git a/bindings/python/uzu/_tool.py b/bindings/python/uzu/_tool.py index 411562770..fbad190b6 100644 --- a/bindings/python/uzu/_tool.py +++ b/bindings/python/uzu/_tool.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import collections.abc import functools import inspect @@ -98,7 +99,7 @@ def __init__( visited=set(), ) - functools.update_wrapper(self, function) + functools.update_wrapper(self, function, updated=()) def __call__( self, @@ -146,11 +147,39 @@ async def _invoke_json_on_loop(self, arguments_json: str) -> str: result = self._invoke_json(arguments_json) return await result if inspect.isawaitable(result) else result + def _new_json_invocation(self, arguments_json: str) -> _ToolInvocation: + return _ToolInvocation(self, arguments_json) + def _encode_result(self, result: object) -> str: validated = self._return_adapter.validate_python(result) return self._return_adapter.dump_json(validated).decode() +class _ToolInvocation: + def __init__( + self, + tool: UzuToolFunction, + arguments_json: str, + ) -> None: + self._tool = tool + self._arguments_json = arguments_json + self._cancelled = False + self._loop: asyncio.AbstractEventLoop | None = None + self._task: asyncio.Task[object] | None = None + + async def run(self) -> str: + self._loop = asyncio.get_running_loop() + self._task = asyncio.current_task() + if self._cancelled: + raise asyncio.CancelledError + return await self._tool._invoke_json_on_loop(self._arguments_json) # noqa: SLF001 + + def cancel(self) -> None: + self._cancelled = True + if self._loop is not None and self._task is not None: + self._loop.call_soon_threadsafe(self._task.cancel) + + @overload def uzu_tool_function[**P, R]( function: collections.abc.Callable[P, R], diff --git a/crates/nagare/src/chat/bindings_pyo3.rs b/crates/nagare/src/chat/bindings_pyo3.rs index dbd20ef93..d1ee33a0c 100644 --- a/crates/nagare/src/chat/bindings_pyo3.rs +++ b/crates/nagare/src/chat/bindings_pyo3.rs @@ -1,4 +1,9 @@ -use std::{future::Future, pin::Pin, sync::Arc}; +use std::{ + future::Future, + pin::Pin, + sync::Arc, + task::{Context, Poll}, +}; use pyo3::{Bound, Py, PyAny, PyErr, PyResult, Python, types::PyAnyMethods}; use shoji::types::basic::{ToolFunction, Value}; @@ -8,6 +13,37 @@ use crate::tool::func_def::{ErrorFuture, ToolDescriptor}; type PythonFuture = Pin>> + Send>>; +struct PythonInvocation { + future: PythonFuture, + cancellation: Option>, +} + +impl Future for PythonInvocation { + type Output = PyResult>; + + fn poll( + mut self: Pin<&mut Self>, + context: &mut Context<'_>, + ) -> Poll { + let result = self.future.as_mut().poll(context); + if result.is_ready() { + self.cancellation = None; + } + result + } +} + +impl Drop for PythonInvocation { + fn drop(&mut self) { + let Some(cancellation) = self.cancellation.take() else { + return; + }; + let _ = Python::try_attach(|py| { + let _ = cancellation.bind(py).call_method0("cancel"); + }); + } +} + #[pyo3_stub_gen::derive::gen_stub_pymethods] #[pyo3::pymethods] impl ChatSession { @@ -88,15 +124,18 @@ async fn call_python_tool( tool: Arc>, arguments: Value, ) -> Result { - let future = Python::attach(|py| -> PyResult { - // Creating this coroutine does not invoke the tool. Its body runs on the - // reply's Python loop and context once `into_future_with_locals` schedules it. + let invocation = Python::attach(|py| -> PyResult { let task_locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; - let invocation = tool.bind(py).call_method1("_invoke_json_on_loop", (arguments.json,))?; - Ok(Box::pin(pyo3_async_runtimes::into_future_with_locals(&task_locals, invocation)?)) + let cancellation = tool.bind(py).call_method1("_new_json_invocation", (arguments.json,))?; + let awaitable = cancellation.call_method0("run")?; + let future = Box::pin(pyo3_async_runtimes::into_future_with_locals(&task_locals, awaitable)?); + Ok(PythonInvocation { + future, + cancellation: Some(cancellation.unbind()), + }) }) .map_err(python_error)?; - let result = future.await.map_err(python_error)?; + let result = invocation.await.map_err(python_error)?; let json = Python::attach(|py| result.extract::(py)).map_err(python_error)?; let value = serde_json::from_str::(&json)?; From 4d13e613e6a954049e50117ea9a335db9718f084 Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 20:06:15 +0300 Subject: [PATCH 09/13] Serialize tool results with schema aliases --- bindings/python/tests/test_tools.py | 20 ++++++++++++++++++++ bindings/python/uzu/_tool.py | 2 +- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/bindings/python/tests/test_tools.py b/bindings/python/tests/test_tools.py index 36b584d38..01db07a16 100644 --- a/bindings/python/tests/test_tools.py +++ b/bindings/python/tests/test_tools.py @@ -42,6 +42,10 @@ class Weather(BaseModel): summary: Annotated[str | None, "Optional human-readable summary."] +class AliasedResult(BaseModel): + value: int = Field(serialization_alias="answer") + + @uzu_tool_function def get_forecast( request: Annotated[ForecastRequest, "Request supplied to the forecast service."], @@ -88,6 +92,11 @@ def unhashable_return_metadata() -> Annotated[int, {"tag": "x"}]: return 42 +@uzu_tool_function +def get_aliased_result() -> AliasedResult: + return AliasedResult(value=1) + + def _resolve_reference( root: dict[str, object], schema: dict[str, object], @@ -163,6 +172,17 @@ def test_sync_tool_constructs_and_serializes_pydantic_models() -> None: } +def test_return_serialization_uses_schema_aliases() -> None: + schema = get_aliased_result.return_schema + assert schema is not None + assert set(schema["properties"]) == {"answer"} + assert schema["required"] == ["answer"] + + result = get_aliased_result._invoke_json("{}") + assert isinstance(result, str) + assert json.loads(result) == {"answer": 1} + + def test_required_nullable_parameter_cannot_be_omitted() -> None: with pytest.raises(ValidationError, match="note"): get_forecast._invoke_json( diff --git a/bindings/python/uzu/_tool.py b/bindings/python/uzu/_tool.py index fbad190b6..d6b8bf18b 100644 --- a/bindings/python/uzu/_tool.py +++ b/bindings/python/uzu/_tool.py @@ -152,7 +152,7 @@ def _new_json_invocation(self, arguments_json: str) -> _ToolInvocation: def _encode_result(self, result: object) -> str: validated = self._return_adapter.validate_python(result) - return self._return_adapter.dump_json(validated).decode() + return self._return_adapter.dump_json(validated, by_alias=True).decode() class _ToolInvocation: From 10a9a5d6ddae28e2e3bef9a3987b60074524cbe5 Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 20:08:50 +0300 Subject: [PATCH 10/13] Fixed discriminated-union description propagation. --- bindings/python/tests/test_tools.py | 26 ++++++++++++++++++++++++++ bindings/python/uzu/_tool.py | 2 ++ 2 files changed, 28 insertions(+) diff --git a/bindings/python/tests/test_tools.py b/bindings/python/tests/test_tools.py index 01db07a16..106849cb5 100644 --- a/bindings/python/tests/test_tools.py +++ b/bindings/python/tests/test_tools.py @@ -46,6 +46,16 @@ class AliasedResult(BaseModel): value: int = Field(serialization_alias="answer") +class Cat(BaseModel): + kind: Literal["cat"] + lives: Annotated[int, "Lives remaining."] + + +class Dog(BaseModel): + kind: Literal["dog"] + tricks: Annotated[list[str], "Tricks the dog knows."] + + @uzu_tool_function def get_forecast( request: Annotated[ForecastRequest, "Request supplied to the forecast service."], @@ -97,6 +107,11 @@ def get_aliased_result() -> AliasedResult: return AliasedResult(value=1) +@uzu_tool_function +def describe_pet(pet: Annotated[Cat | Dog, Field(discriminator="kind")]) -> str: + return pet.kind + + def _resolve_reference( root: dict[str, object], schema: dict[str, object], @@ -140,6 +155,17 @@ def test_decorator_builds_schema_from_function_and_pydantic_metadata() -> None: assert properties["unit"]["default"] == "celsius" +def test_discriminated_union_schema_includes_annotated_field_descriptions() -> None: + parameters = describe_pet.parameters_schema + assert parameters is not None + + pet = parameters["properties"]["pet"] + assert "oneOf" in pet + cat, dog = (_resolve_reference(parameters, branch) for branch in pet["oneOf"]) + assert cat["properties"]["lives"]["description"] == "Lives remaining." + assert dog["properties"]["tricks"]["description"] == "Tricks the dog knows." + + def test_pydantic_return_schema_uses_docstrings_and_field_metadata() -> None: schema = get_forecast.return_schema assert schema is not None diff --git a/bindings/python/uzu/_tool.py b/bindings/python/uzu/_tool.py index d6b8bf18b..9463e561c 100644 --- a/bindings/python/uzu/_tool.py +++ b/bindings/python/uzu/_tool.py @@ -249,6 +249,8 @@ def _apply_annotated_descriptions( # noqa: PLR0911 if origin in {typing.Union, types.UnionType}: branches = schema.get("anyOf") + if not isinstance(branches, list): + branches = schema.get("oneOf") if isinstance(branches, list): for member, branch in zip(arguments, branches, strict=False): if isinstance(branch, dict): From 9d3fae91210614b12ed8dc101d3651e08958a288 Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Wed, 29 Jul 2026 20:38:46 +0300 Subject: [PATCH 11/13] Add uv sync --- crates/cli-tools/src/languages/python.rs | 1 + crates/cli-tools/src/types/command.rs | 12 +++++++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/crates/cli-tools/src/languages/python.rs b/crates/cli-tools/src/languages/python.rs index 7af03911b..85ed39532 100644 --- a/crates/cli-tools/src/languages/python.rs +++ b/crates/cli-tools/src/languages/python.rs @@ -51,6 +51,7 @@ impl LanguageBackend for PythonLanguageBackend { if !bindings_path.join(".venv").exists() { Command::uv_venv().with_current_path(&bindings_path).run()?; } + Command::uv_sync_params(true).with_current_path(&bindings_path).run()?; Command::uv_pip_install_wheel(wheel_path.clone()) .with_current_path(&bindings_path) .with_envs(envs.clone()) diff --git a/crates/cli-tools/src/types/command.rs b/crates/cli-tools/src/types/command.rs index a3872704b..2ec9c59df 100644 --- a/crates/cli-tools/src/types/command.rs +++ b/crates/cli-tools/src/types/command.rs @@ -257,7 +257,17 @@ impl Command { } pub fn uv_sync() -> Self { - Self::new("uv").with_argument("sync").with_argument("--reinstall") + Self::uv_sync_params(false) + } + + pub fn uv_sync_params(no_install_project: bool) -> Self { + let mut cmd = Self::new("uv").with_argument("sync"); + if no_install_project { + cmd = cmd.with_argument("--no-install-project") + } else { + cmd = cmd.with_argument("--reinstall") + } + cmd } pub fn uv_venv() -> Self { From d81e4cbd4a844c6f6938503521097641c46a6ff1 Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Thu, 30 Jul 2026 14:02:37 +0300 Subject: [PATCH 12/13] Allow documenting fields/arguments by parsing docstrings --- bindings/python/examples/tool_calls.py | 27 ++++++++++------- bindings/python/pyproject.toml | 5 ++- bindings/python/tests/test_tools.py | 42 ++++++++++++++++++++++++++ bindings/python/uv.lock | 11 +++++++ bindings/python/uzu/_tool.py | 30 ++++++++++++++++-- 5 files changed, 100 insertions(+), 15 deletions(-) diff --git a/bindings/python/examples/tool_calls.py b/bindings/python/examples/tool_calls.py index eaea2b9db..0ce68694f 100644 --- a/bindings/python/examples/tool_calls.py +++ b/bindings/python/examples/tool_calls.py @@ -16,26 +16,33 @@ class Coordinate(BaseModel): - """A geographic coordinate.""" + """A geographic coordinate. - latitude: Annotated[float, "Latitude in decimal degrees."] + Attributes: + latitude: Latitude in decimal degrees. + longitude: Longitude in decimal degrees. + """ + + latitude: float longitude: Annotated[float, "Longitude in decimal degrees."] -@uzu_tool_function( - name="get_location", - description="Return the current location in coordinates" -) +@uzu_tool_function(name="get_location", description="Return the current location in coordinates") def get_current_location() -> Coordinate: return Coordinate(latitude=51.5074, longitude=-0.1278) @uzu_tool_function def get_current_temperature( - latitude: Annotated[float, "Latitude in decimal degrees."], + latitude: float, longitude: Annotated[float, "Longitude in decimal degrees."], ) -> float: - """Return the temperature at the provided coordinates.""" + """Return the temperature at the provided coordinates. + + Args: + latitude: Latitude in decimal degrees. + longitude: This is overridden by the Annotated description. + """ _ = latitude, longitude return 25.0 @@ -57,9 +64,7 @@ async def main() -> None: ChatMessage.system().with_text("You are a helpful assistant"), ChatMessage.user().with_text("What temperature is it now at my location?"), ] - config = ChatReplyConfig.create().with_sampling_policy( - SamplingPolicy.Custom(method=SamplingMethod.Greedy()) - ) + config = ChatReplyConfig.create().with_sampling_policy(SamplingPolicy.Custom(method=SamplingMethod.Greedy())) replies = await session.reply(messages, config) if replies: message = replies[-1].message diff --git a/bindings/python/pyproject.toml b/bindings/python/pyproject.toml index 705e129f7..6f9bca71c 100644 --- a/bindings/python/pyproject.toml +++ b/bindings/python/pyproject.toml @@ -27,7 +27,10 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dynamic = ["version"] -dependencies = ["pydantic>=2.13.3"] +dependencies = [ + "docstring-parser>=0.17.0", + "pydantic>=2.13.3", +] [project.urls] Homepage = "https://trymirai.com" diff --git a/bindings/python/tests/test_tools.py b/bindings/python/tests/test_tools.py index 106849cb5..40202f8e8 100644 --- a/bindings/python/tests/test_tools.py +++ b/bindings/python/tests/test_tools.py @@ -42,6 +42,18 @@ class Weather(BaseModel): summary: Annotated[str | None, "Optional human-readable summary."] +class DocstringCoordinate(BaseModel): + """A coordinate documented through its class docstring. + + Attributes: + latitude: Latitude from the class docstring. + longitude: This must not replace explicit field metadata. + """ + + latitude: float + longitude: float = Field(description="Explicit longitude description.") + + class AliasedResult(BaseModel): value: int = Field(serialization_alias="answer") @@ -112,6 +124,22 @@ def describe_pet(pet: Annotated[Cat | Dog, Field(discriminator="kind")]) -> str: return pet.kind +@uzu_tool_function +def locate( + coordinate: DocstringCoordinate, + label: Annotated[str, "Explicit label description."], +) -> None: + """Locate a named coordinate. + + The coordinate and label are validated before use. + + Args: + coordinate: Coordinate from the function docstring. + label: This must not replace the Annotated description. + """ + _ = coordinate, label + + def _resolve_reference( root: dict[str, object], schema: dict[str, object], @@ -155,6 +183,20 @@ def test_decorator_builds_schema_from_function_and_pydantic_metadata() -> None: assert properties["unit"]["default"] == "celsius" +def test_docstrings_describe_tool_arguments_and_model_fields() -> None: + assert locate.description == ("Locate a named coordinate.\n\nThe coordinate and label are validated before use.") + + parameters = locate.parameters_schema + assert parameters is not None + properties = parameters["properties"] + assert properties["coordinate"]["description"] == "Coordinate from the function docstring." + assert properties["label"]["description"] == "Explicit label description." + + coordinate = _resolve_reference(parameters, properties["coordinate"]) + assert coordinate["properties"]["latitude"]["description"] == "Latitude from the class docstring." + assert coordinate["properties"]["longitude"]["description"] == "Explicit longitude description." + + def test_discriminated_union_schema_includes_annotated_field_descriptions() -> None: parameters = describe_pet.parameters_schema assert parameters is not None diff --git a/bindings/python/uv.lock b/bindings/python/uv.lock index 053abd0e5..5f1a8cbb2 100644 --- a/bindings/python/uv.lock +++ b/bindings/python/uv.lock @@ -20,6 +20,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "docstring-parser" +version = "0.18.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/4d/f332313098c1de1b2d2ff91cf2674415cc7cddab2ca1b01ae29774bd5fdf/docstring_parser-0.18.0.tar.gz", hash = "sha256:292510982205c12b1248696f44959db3cdd1740237a968ea1e2e7a900eeb2015", size = 29341, upload-time = "2026-04-14T04:09:19.867Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/5f/ed01f9a3cdffbd5a008556fc7b2a08ddb1cc6ace7effa7340604b1d16699/docstring_parser-0.18.0-py3-none-any.whl", hash = "sha256:b3fcbed555c47d8479be0796ef7e19c2670d428d72e96da63f3a40122860374b", size = 22484, upload-time = "2026-04-14T04:09:18.638Z" }, +] + [[package]] name = "iniconfig" version = "2.3.0" @@ -212,6 +221,7 @@ wheels = [ name = "uzu" source = { editable = "." } dependencies = [ + { name = "docstring-parser" }, { name = "pydantic" }, ] @@ -223,6 +233,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "docstring-parser", specifier = ">=0.17.0" }, { name = "pydantic", specifier = ">=2.13.3" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.3.4" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.14.0" }, diff --git a/bindings/python/uzu/_tool.py b/bindings/python/uzu/_tool.py index 9463e561c..9f1d263be 100644 --- a/bindings/python/uzu/_tool.py +++ b/bindings/python/uzu/_tool.py @@ -9,6 +9,7 @@ import typing from typing import Annotated, Any, Literal, overload +from docstring_parser import parse as parse_docstring from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, create_model _EMPTY = inspect.Parameter.empty @@ -36,7 +37,8 @@ def __init__( self._function = function self.name = name - self.description = description if description is not None else inspect.getdoc(function) or "" + inferred_description, argument_descriptions = _docstring_hints(function) + self.description = description if description is not None else inferred_description self._signature = inspect.signature(function) try: @@ -56,9 +58,15 @@ def __init__( if annotation is _EMPTY: raise TypeError(f"parameter {parameter.name!r} of tool {name!r} must have a type annotation") + annotation_description = _annotation_description(annotation) + field_description = ( + annotation_description + if annotation_description is not None + else argument_descriptions.get(parameter.name) + ) field_metadata = ( - Field(alias=parameter.name, description=annotation_description) - if (annotation_description := _annotation_description(annotation)) + Field(alias=parameter.name, description=field_description) + if field_description is not None else Field(alias=parameter.name) ) annotation = Annotated[annotation, field_metadata] @@ -212,6 +220,19 @@ def decorate(target: collections.abc.Callable[P, R]) -> UzuToolFunction[P, R]: return decorate if function is None else decorate(function) +def _docstring_hints(target: object) -> tuple[str, dict[str, str]]: + docstring = inspect.getdoc(target) + if docstring is None: + return "", {} + + parsed = parse_docstring(docstring) + description = "\n\n".join(part for part in (parsed.short_description, parsed.long_description) if part) + parameter_descriptions = { + parameter.arg_name: parameter.description for parameter in parsed.params if parameter.description + } + return description, parameter_descriptions + + def _annotation_description(annotation: object) -> str | None: if typing.get_origin(annotation) is not typing.Annotated: return None @@ -327,12 +348,15 @@ def _apply_annotated_descriptions( # noqa: PLR0911 annotations = typing.get_type_hints(annotation, include_extras=True) except (NameError, TypeError): annotations = {name: field.annotation for name, field in annotation.model_fields.items()} + _, field_descriptions = _docstring_hints(annotation) for name, field in annotation.model_fields.items(): field_annotation = annotations.get(name, field.annotation) alias = field.serialization_alias if mode == "serialization" else field.validation_alias field_schema = properties.get(alias if isinstance(alias, str) else name) if isinstance(field_schema, dict): + if description := field_descriptions.get(name): + field_schema.setdefault("description", description) _apply_annotated_descriptions( field_annotation, field_schema, From 464faa8e6605d74023737d29c4166300881a8ff4 Mon Sep 17 00:00:00 2001 From: Aleksandr Golokoz <5617556+agolokoz@users.noreply.github.com> Date: Thu, 30 Jul 2026 14:21:22 +0300 Subject: [PATCH 13/13] Stop requesting the removed examples extra --- crates/cli-tools/src/types/command.rs | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/crates/cli-tools/src/types/command.rs b/crates/cli-tools/src/types/command.rs index 2ec9c59df..d268cba31 100644 --- a/crates/cli-tools/src/types/command.rs +++ b/crates/cli-tools/src/types/command.rs @@ -300,11 +300,7 @@ impl Command { } pub fn uv_python_file(path: PathBuf) -> Self { - Self::new("uv") - .with_argument("run") - .with_arguments(vec!["--extra".to_string(), "examples".to_string()]) - .with_argument("python") - .with_argument(&path.to_string_lossy()) + Self::new("uv").with_argument("run").with_argument("python").with_argument(&path.to_string_lossy()) } pub fn maturin_build(