From f533f6a01339c8efb25a7f6c10a13a8d248255be Mon Sep 17 00:00:00 2001 From: Elias Posen Date: Thu, 7 May 2026 16:49:00 -0400 Subject: [PATCH 1/5] input_schema as base model or schema --- pctx-py/pyproject.toml | 1 + pctx-py/src/pctx_client/_convert.py | 18 ++- pctx-py/src/pctx_client/_tool.py | 50 ++++++- pctx-py/src/pctx_client/_websocket_client.py | 3 +- pctx-py/tests/test_tool_decorator.py | 149 ++++++++++++++++++- pctx-py/uv.lock | 2 + 6 files changed, 211 insertions(+), 12 deletions(-) diff --git a/pctx-py/pyproject.toml b/pctx-py/pyproject.toml index 1b760135..0e6387bf 100644 --- a/pctx-py/pyproject.toml +++ b/pctx-py/pyproject.toml @@ -11,6 +11,7 @@ requires-python = ">=3.10,<3.15" dependencies = [ "docstring-parser>=0.17.0", "httpx>=0.28.1", + "jsonschema>=4.26.0", "pydantic>=2.7.2", "websockets>=15.0.1", ] diff --git a/pctx-py/src/pctx_client/_convert.py b/pctx-py/src/pctx_client/_convert.py index e3d4cc86..cac592fc 100644 --- a/pctx-py/src/pctx_client/_convert.py +++ b/pctx-py/src/pctx_client/_convert.py @@ -1,6 +1,8 @@ from collections.abc import Callable from typing import Any, overload +from pydantic import BaseModel + from pctx_client._tool import AsyncTool, Tool @@ -10,6 +12,7 @@ def tool( *args: Any, namespace: str = "tools", description: str | None = None, + input_schema: type[BaseModel] | dict[str, Any] | None = None, ) -> Callable[[Callable], Tool | AsyncTool]: ... @overload def tool( @@ -17,6 +20,7 @@ def tool( *args: Any, namespace: str = "tools", description: str | None = None, + input_schema: type[BaseModel] | dict[str, Any] | None = None, ) -> Tool | AsyncTool: ... @@ -25,6 +29,7 @@ def tool( *args: Any, namespace: str = "tools", description: str | None = None, + input_schema: type[BaseModel] | dict[str, Any] | None = None, ) -> Tool | AsyncTool | Callable[[Callable], Tool | AsyncTool]: """ Decorator that converts a function into a Tool or AsyncTool instance. @@ -33,11 +38,15 @@ def tool( - @tool - Uses function name as tool name - @tool("custom_name") - Uses custom name for the tool - @tool(namespace="custom", description="...") - With additional options + - @tool(input_schema=MyModel) or @tool(input_schema={...}) - Override + signature inference with an explicit Pydantic model or JSON Schema dict Args: name_or_callable: Either a custom tool name (str) or the function to wrap (Callable) namespace: The namespace the tool belongs to (default: "tools") description: Optional description override (default: uses function docstring) + input_schema: Optional explicit input schema. When provided, signature + inference is skipped and this schema is used directly. Returns: Either a Tool/AsyncTool instance or a decorator function that creates one @@ -51,6 +60,10 @@ def tool( >>> @tool("custom_name", namespace="math") ... def add_two(x: int) -> int: ... return x + 2 + + >>> @tool(input_schema={"type": "object", "properties": {"x": {"type": "integer"}}, "required": ["x"]}) + ... def from_jsonschema(**kwargs) -> int: + ... return kwargs["x"] + 1 """ def _crate_tool_factory(tool_name: str) -> Callable[[Callable], Tool | AsyncTool]: @@ -65,13 +78,12 @@ def _crate_tool_factory(tool_name: str) -> Callable[[Callable], Tool | AsyncTool """ def _tool_factory(fn: Callable) -> Tool | AsyncTool: - tool_desc = description - return Tool.from_func( func=fn, name=tool_name, namespace=namespace, - description=tool_desc, + description=description, + input_schema=input_schema, ) return _tool_factory diff --git a/pctx-py/src/pctx_client/_tool.py b/pctx-py/src/pctx_client/_tool.py index a53ab025..f8702625 100644 --- a/pctx-py/src/pctx_client/_tool.py +++ b/pctx-py/src/pctx_client/_tool.py @@ -7,13 +7,16 @@ from docstring_parser import Docstring from docstring_parser import parse as parse_docstring +from jsonschema import Draft202012Validator from pydantic import ( BaseModel, ConfigDict, Field, + PrivateAttr, SkipValidation, TypeAdapter, create_model, + model_validator, ) @@ -33,16 +36,42 @@ class BaseTool(BaseModel): Longer-form text which instructs the model how/why/when to use the tool. """ - input_schema: Annotated[type[BaseModel] | None, SkipValidation] = Field( - default=None, description="The tool schema." + input_schema: Annotated[ + type[BaseModel] | dict[str, Any] | None, SkipValidation + ] = Field( + default=None, + description="The tool input schema. Either a Pydantic BaseModel class or a JSON Schema dict.", ) output_schema: Annotated[Any | None, SkipValidation] = Field( default=None, description="The return type schema." ) + _input_validator: Draft202012Validator | None = PrivateAttr(default=None) + + @model_validator(mode="after") + def _compile_input_validator(self) -> "BaseTool": + """ + When ``input_schema`` is a JSON Schema dict, validate it against the + Draft 2020-12 metaschema and cache a compiled validator on the + instance so per-call ``validate_input`` doesn't recompile. + + Raises: + jsonschema.SchemaError: If the provided dict is not a valid + JSON Schema. + """ + if isinstance(self.input_schema, dict): + Draft202012Validator.check_schema(self.input_schema) + self._input_validator = Draft202012Validator(self.input_schema) + return self + def validate_input(self, obj: Any): - if self.input_schema is not None: + if self.input_schema is None: + return + if isinstance(self.input_schema, dict): + assert self._input_validator is not None + self._input_validator.validate(obj) + else: self.input_schema.model_validate(obj) def validate_output(self, obj: Any): @@ -53,7 +82,8 @@ def validate_output(self, obj: Any): def input_json_schema(self) -> dict[str, Any] | None: if self.input_schema is None: return None - + if isinstance(self.input_schema, dict): + return self.input_schema return self.input_schema.model_json_schema() def output_json_schema(self) -> dict[str, Any] | None: @@ -70,19 +100,25 @@ def from_func( name: str | None = None, namespace: str = "tools", description: str | None = None, + input_schema: type[BaseModel] | dict[str, Any] | None = None, ) -> "Tool | AsyncTool": """ Creates a tool from a given function. + + If ``input_schema`` is provided (either a Pydantic BaseModel class or a + JSON Schema dict), it is used as-is and signature inference is skipped. + Otherwise the schema is inferred from the function signature. """ docstring = parse_docstring(textwrap.dedent(description or func.__doc__ or "")) name_ = name or func.__name__ - in_schema = create_input_schema(f"{name_}_Input", func, docstring=docstring) - out_schema = create_output_schema(func, docstring=docstring) + if input_schema is None: + in_schema = create_input_schema(f"{name_}_Input", func, docstring=docstring) + input_schema = None if is_empty_schema(in_schema) else in_schema - input_schema = None if is_empty_schema(in_schema) else in_schema + out_schema = create_output_schema(func, docstring=docstring) output_schema = out_schema # Create concrete tool classes based on sync vs async diff --git a/pctx-py/src/pctx_client/_websocket_client.py b/pctx-py/src/pctx_client/_websocket_client.py index fa4bcf17..7fede708 100644 --- a/pctx-py/src/pctx_client/_websocket_client.py +++ b/pctx-py/src/pctx_client/_websocket_client.py @@ -10,6 +10,7 @@ import uuid from typing import Any, Union +import jsonschema import pydantic import websockets from websockets.asyncio.client import ClientConnection @@ -233,7 +234,7 @@ async def _handle_execute_tool( return ExecuteToolResponse( id=req.id, result=ExecuteToolResult(output=output) ) - except pydantic.ValidationError as e: + except (pydantic.ValidationError, jsonschema.ValidationError) as e: return JsonRpcError( id=req.id, error=ErrorData( diff --git a/pctx-py/tests/test_tool_decorator.py b/pctx-py/tests/test_tool_decorator.py index 6295dd90..9d445d6c 100644 --- a/pctx-py/tests/test_tool_decorator.py +++ b/pctx-py/tests/test_tool_decorator.py @@ -2,8 +2,9 @@ from __future__ import annotations +import jsonschema import pytest -from pydantic import ValidationError +from pydantic import BaseModel, ValidationError from pctx_client import Tool, tool from pctx_client._tool import AsyncTool @@ -656,3 +657,149 @@ def calculate_distance( assert "description" in output_schema assert "Euclidean distance" in output_schema["description"] assert "floating point" in output_schema["description"] + + +# ============================================================================ +# SECTION: EXPLICIT input_schema (skips signature inference) +# ============================================================================ + + +def test_input_schema_dict_skips_inference() -> None: + """An explicit JSON Schema dict is used as-is and signature is not inspected.""" + + schema: dict = { + "type": "object", + "properties": {"x": {"type": "integer"}}, + "required": ["x"], + "additionalProperties": False, + } + + @tool("from_dict", input_schema=schema) + def from_dict(**kwargs) -> int: + return kwargs["x"] + 1 + + assert isinstance(from_dict, Tool) + # Schema is the exact dict the user passed, not derived from the signature. + assert from_dict.input_json_schema() == schema + assert from_dict.invoke(x=4) == 5 + + +def test_input_schema_dict_rejects_invalid_input() -> None: + """Dict-defined input_schema validates via jsonschema and rejects bad input.""" + + @tool( + "add_one", + input_schema={ + "type": "object", + "properties": {"x": {"type": "integer"}}, + "required": ["x"], + }, + ) + def add_one(**kwargs) -> int: + return kwargs["x"] + 1 + + with pytest.raises(jsonschema.ValidationError): + add_one.invoke(x="not-an-int") + + with pytest.raises(jsonschema.ValidationError): + add_one.invoke() # missing required `x` + + +def test_input_schema_dict_malformed_raises_at_construction() -> None: + """A malformed JSON Schema is rejected when the tool is built, not on first call.""" + + with pytest.raises(jsonschema.SchemaError): + + @tool("bad", input_schema={"type": "not-a-real-type"}) + def bad(**kwargs) -> int: + return 0 + + +def test_input_schema_pydantic_model_skips_inference() -> None: + """An explicit Pydantic BaseModel class is used as-is, ignoring the signature.""" + + class Args(BaseModel): + y: float + + @tool("from_model", input_schema=Args) + def from_model(**kwargs) -> float: + return kwargs["y"] * 2 + + assert from_model.input_schema is Args + schema = from_model.input_json_schema() + assert schema is not None + assert schema["properties"] == {"y": {"title": "Y", "type": "number"}} + assert from_model.invoke(y=1.5) == 3.0 + + +def test_input_schema_pydantic_model_rejects_invalid_input() -> None: + """Pydantic-defined input_schema raises pydantic.ValidationError on bad input.""" + + class Args(BaseModel): + y: float + + @tool("from_model", input_schema=Args) + def from_model(**kwargs) -> float: + return kwargs["y"] * 2 + + with pytest.raises(ValidationError): + from_model.invoke(y="not-a-float") + + +def test_input_schema_explicit_overrides_signature() -> None: + """When input_schema is provided, the function's own annotations are ignored.""" + + schema: dict = { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + } + + # Function signature says `n: int`, but we pass an unrelated schema. + # The explicit schema wins; signature inference is skipped entirely. + @tool("search", input_schema=schema) + def search(n: int = 0, **kwargs) -> str: + return kwargs.get("q", "") + + assert search.input_json_schema() == schema + assert search.invoke(q="hello") == "hello" + + +async def test_input_schema_dict_async() -> None: + """Explicit JSON Schema dict works for async tools too.""" + + @tool( + "add_one_async", + input_schema={ + "type": "object", + "properties": {"x": {"type": "integer"}}, + "required": ["x"], + }, + ) + async def add_one(**kwargs) -> int: + return kwargs["x"] + 1 + + assert isinstance(add_one, AsyncTool) + assert await add_one.ainvoke(x=4) == 5 + + with pytest.raises(jsonschema.ValidationError): + await add_one.ainvoke(x="bad") + + +def test_input_schema_with_custom_name_and_namespace() -> None: + """Explicit input_schema composes with custom name + namespace.""" + + schema: dict = { + "type": "object", + "properties": {"n": {"type": "integer"}}, + "required": ["n"], + } + + @tool("add_one", namespace="math", input_schema=schema) + def named(**kwargs) -> int: + return kwargs["n"] + 1 + + assert named.name == "add_one" + assert named.namespace == "math" + assert named.input_json_schema() == schema + assert named.invoke(n=2) == 3 diff --git a/pctx-py/uv.lock b/pctx-py/uv.lock index 0e5b4052..b3c94f66 100644 --- a/pctx-py/uv.lock +++ b/pctx-py/uv.lock @@ -3344,6 +3344,7 @@ source = { editable = "." } dependencies = [ { name = "docstring-parser" }, { name = "httpx" }, + { name = "jsonschema" }, { name = "pydantic" }, { name = "websockets" }, ] @@ -3398,6 +3399,7 @@ requires-dist = [ { name = "crewai", marker = "extra == 'crewai'", specifier = ">=1.6.1" }, { name = "docstring-parser", specifier = ">=0.17.0" }, { name = "httpx", specifier = ">=0.28.1" }, + { name = "jsonschema", specifier = ">=4.26.0" }, { name = "langchain", marker = "extra == 'langchain'", specifier = ">=1.1.2" }, { name = "openai-agents", marker = "extra == 'openai'", specifier = ">=0.12.0" }, { name = "pydantic", specifier = ">=2.7.2" }, From a7bc944b4dd079bd82f52d679947e4eb1450f2a5 Mon Sep 17 00:00:00 2001 From: Elias Posen Date: Thu, 7 May 2026 16:58:58 -0400 Subject: [PATCH 2/5] output_schema as python type or schema --- pctx-py/src/pctx_client/_convert.py | 9 ++ pctx-py/src/pctx_client/_tool.py | 62 +++++++---- pctx-py/tests/test_tool_decorator.py | 150 +++++++++++++++++++++++++++ 3 files changed, 199 insertions(+), 22 deletions(-) diff --git a/pctx-py/src/pctx_client/_convert.py b/pctx-py/src/pctx_client/_convert.py index cac592fc..e62fd36c 100644 --- a/pctx-py/src/pctx_client/_convert.py +++ b/pctx-py/src/pctx_client/_convert.py @@ -13,6 +13,7 @@ def tool( namespace: str = "tools", description: str | None = None, input_schema: type[BaseModel] | dict[str, Any] | None = None, + output_schema: Any | None = None, ) -> Callable[[Callable], Tool | AsyncTool]: ... @overload def tool( @@ -21,6 +22,7 @@ def tool( namespace: str = "tools", description: str | None = None, input_schema: type[BaseModel] | dict[str, Any] | None = None, + output_schema: Any | None = None, ) -> Tool | AsyncTool: ... @@ -30,6 +32,7 @@ def tool( namespace: str = "tools", description: str | None = None, input_schema: type[BaseModel] | dict[str, Any] | None = None, + output_schema: Any | None = None, ) -> Tool | AsyncTool | Callable[[Callable], Tool | AsyncTool]: """ Decorator that converts a function into a Tool or AsyncTool instance. @@ -40,6 +43,9 @@ def tool( - @tool(namespace="custom", description="...") - With additional options - @tool(input_schema=MyModel) or @tool(input_schema={...}) - Override signature inference with an explicit Pydantic model or JSON Schema dict + - @tool(output_schema={...}) or @tool(output_schema=MyModel) - Override + return-annotation inference with an explicit JSON Schema dict or + Python type / typing construct Args: name_or_callable: Either a custom tool name (str) or the function to wrap (Callable) @@ -47,6 +53,8 @@ def tool( description: Optional description override (default: uses function docstring) input_schema: Optional explicit input schema. When provided, signature inference is skipped and this schema is used directly. + output_schema: Optional explicit output schema. When provided, return- + annotation inference is skipped and this schema is used directly. Returns: Either a Tool/AsyncTool instance or a decorator function that creates one @@ -84,6 +92,7 @@ def _tool_factory(fn: Callable) -> Tool | AsyncTool: namespace=namespace, description=description, input_schema=input_schema, + output_schema=output_schema, ) return _tool_factory diff --git a/pctx-py/src/pctx_client/_tool.py b/pctx-py/src/pctx_client/_tool.py index f8702625..e3d80900 100644 --- a/pctx-py/src/pctx_client/_tool.py +++ b/pctx-py/src/pctx_client/_tool.py @@ -36,33 +36,42 @@ class BaseTool(BaseModel): Longer-form text which instructs the model how/why/when to use the tool. """ - input_schema: Annotated[ - type[BaseModel] | dict[str, Any] | None, SkipValidation - ] = Field( - default=None, - description="The tool input schema. Either a Pydantic BaseModel class or a JSON Schema dict.", + input_schema: Annotated[type[BaseModel] | dict[str, Any] | None, SkipValidation] = ( + Field( + default=None, + description="The tool input schema. Either a Pydantic BaseModel class or a JSON Schema dict.", + ) ) output_schema: Annotated[Any | None, SkipValidation] = Field( - default=None, description="The return type schema." + default=None, + description=( + "The return type schema. Either a JSON Schema dict, or any Python " + "type / typing construct accepted by pydantic.TypeAdapter " + "(e.g. int, list[int], Annotated[Foo, Field(...)], BaseModel subclass)." + ), ) _input_validator: Draft202012Validator | None = PrivateAttr(default=None) + _output_validator: Draft202012Validator | None = PrivateAttr(default=None) @model_validator(mode="after") - def _compile_input_validator(self) -> "BaseTool": + def _compile_schema_validators(self) -> "BaseTool": """ - When ``input_schema`` is a JSON Schema dict, validate it against the - Draft 2020-12 metaschema and cache a compiled validator on the - instance so per-call ``validate_input`` doesn't recompile. + For any schema field that's a JSON Schema dict, validate it against + the Draft 2020-12 metaschema and cache a compiled validator on the + instance so per-call validation doesn't recompile. Raises: - jsonschema.SchemaError: If the provided dict is not a valid - JSON Schema. + jsonschema.SchemaError: If a provided dict is not a valid JSON + Schema. """ if isinstance(self.input_schema, dict): Draft202012Validator.check_schema(self.input_schema) self._input_validator = Draft202012Validator(self.input_schema) + if isinstance(self.output_schema, dict): + Draft202012Validator.check_schema(self.output_schema) + self._output_validator = Draft202012Validator(self.output_schema) return self def validate_input(self, obj: Any): @@ -75,7 +84,12 @@ def validate_input(self, obj: Any): self.input_schema.model_validate(obj) def validate_output(self, obj: Any): - if self.output_schema is not None: + if self.output_schema is None: + return + if isinstance(self.output_schema, dict): + assert self._output_validator is not None + self._output_validator.validate(obj) + else: adapter = TypeAdapter(self.output_schema) adapter.validate_python(obj) @@ -89,7 +103,8 @@ def input_json_schema(self) -> dict[str, Any] | None: def output_json_schema(self) -> dict[str, Any] | None: if self.output_schema is None: return None - + if isinstance(self.output_schema, dict): + return self.output_schema adapter = TypeAdapter(self.output_schema) return adapter.json_schema() @@ -101,13 +116,16 @@ def from_func( namespace: str = "tools", description: str | None = None, input_schema: type[BaseModel] | dict[str, Any] | None = None, + output_schema: Any | None = None, ) -> "Tool | AsyncTool": """ Creates a tool from a given function. - If ``input_schema`` is provided (either a Pydantic BaseModel class or a - JSON Schema dict), it is used as-is and signature inference is skipped. - Otherwise the schema is inferred from the function signature. + If ``input_schema`` or ``output_schema`` is provided, it is used as-is + and the corresponding inference step is skipped. ``input_schema`` + accepts a Pydantic BaseModel class or a JSON Schema dict; + ``output_schema`` accepts a JSON Schema dict or any Python type / + typing construct that pydantic.TypeAdapter accepts. """ docstring = parse_docstring(textwrap.dedent(description or func.__doc__ or "")) @@ -115,11 +133,11 @@ def from_func( name_ = name or func.__name__ if input_schema is None: - in_schema = create_input_schema(f"{name_}_Input", func, docstring=docstring) + in_schema = infer_input_model(f"{name_}_Input", func, docstring=docstring) input_schema = None if is_empty_schema(in_schema) else in_schema - out_schema = create_output_schema(func, docstring=docstring) - output_schema = out_schema + if output_schema is None: + output_schema = infer_output_type(func, docstring=docstring) # Create concrete tool classes based on sync vs async if asyncio.iscoroutinefunction(func): @@ -247,7 +265,7 @@ async def ainvoke(self, **kwargs: Any) -> Any: _MODEL_CONFIG: ConfigDict = {"extra": "forbid", "arbitrary_types_allowed": True} -def create_input_schema( +def infer_input_model( model_name: str, func: Callable, docstring: Docstring | None = None ) -> type[BaseModel]: """ @@ -302,7 +320,7 @@ def create_input_schema( return create_model(model_name, __config__=_MODEL_CONFIG, **fields) -def create_output_schema(func: Callable, docstring: Docstring | None = None) -> Any: +def infer_output_type(func: Callable, docstring: Docstring | None = None) -> Any: """ Extracts the return type annotation from a function. diff --git a/pctx-py/tests/test_tool_decorator.py b/pctx-py/tests/test_tool_decorator.py index 9d445d6c..2dd52785 100644 --- a/pctx-py/tests/test_tool_decorator.py +++ b/pctx-py/tests/test_tool_decorator.py @@ -803,3 +803,153 @@ def named(**kwargs) -> int: assert named.namespace == "math" assert named.input_json_schema() == schema assert named.invoke(n=2) == 3 + + +# ============================================================================ +# SECTION: EXPLICIT output_schema (skips return-annotation inference) +# ============================================================================ + + +def test_output_schema_dict_skips_inference() -> None: + """An explicit JSON Schema dict for output is used as-is.""" + + schema: dict = {"type": "integer", "minimum": 0} + + # Function annotation says `str`, but explicit dict overrides it. + @tool("returns_int", output_schema=schema) + def returns_int() -> str: + return 42 # type: ignore[return-value] + + assert returns_int.output_json_schema() == schema + assert returns_int.invoke() == 42 + + +def test_output_schema_dict_rejects_invalid_output() -> None: + """Dict-defined output_schema validates via jsonschema and rejects bad output.""" + + @tool( + "produces", + output_schema={"type": "string", "minLength": 3}, + ) + def produces(value): + return value + + assert produces.invoke(value="hello") == "hello" + + with pytest.raises(jsonschema.ValidationError): + produces.invoke(value=123) # not a string + + with pytest.raises(jsonschema.ValidationError): + produces.invoke(value="hi") # too short + + +def test_output_schema_dict_malformed_raises_at_construction() -> None: + """A malformed output JSON Schema is rejected when the tool is built.""" + + with pytest.raises(jsonschema.SchemaError): + + @tool("bad_out", output_schema={"type": "not-a-real-type"}) + def _bad_out() -> int: # pyright: ignore[reportUnusedFunction] + return 0 + + +def test_output_schema_pydantic_model_skips_inference() -> None: + """An explicit Pydantic BaseModel class for output is used as-is.""" + + class Result(BaseModel): + value: int + label: str + + @tool("returns_model", output_schema=Result) + def returns_model() -> dict: + return {"value": 1, "label": "one"} + + schema = returns_model.output_json_schema() + assert schema is not None + assert schema["properties"] == { + "value": {"title": "Value", "type": "integer"}, + "label": {"title": "Label", "type": "string"}, + } + # TypeAdapter validates the dict against the Pydantic model. + assert returns_model.invoke() == {"value": 1, "label": "one"} + + +def test_output_schema_pydantic_model_rejects_invalid_output() -> None: + """Pydantic-defined output_schema raises pydantic.ValidationError on bad output.""" + + class Result(BaseModel): + value: int + + @tool("returns_model", output_schema=Result) + def returns_model(payload): + return payload + + with pytest.raises(ValidationError): + returns_model.invoke(payload={"value": "not-an-int"}) + + +def test_output_schema_overrides_return_annotation() -> None: + """When output_schema is provided, the function's return annotation is ignored.""" + + # Function annotated as -> int, but explicit schema says string. + @tool("override", output_schema={"type": "string"}) + def override() -> int: + return "hello" # type: ignore[return-value] + + assert override.output_json_schema() == {"type": "string"} + assert override.invoke() == "hello" + + +def test_output_schema_plain_python_type() -> None: + """A plain python type (e.g. int) passed as output_schema works via TypeAdapter.""" + + @tool("counter", output_schema=int) + def counter(n): + return n + + assert counter.output_json_schema() == {"type": "integer"} + assert counter.invoke(n=5) == 5 + + with pytest.raises(ValidationError): + counter.invoke(n="not-an-int") + + +async def test_output_schema_dict_async() -> None: + """Explicit output JSON Schema works for async tools too.""" + + @tool("async_int", output_schema={"type": "integer"}) + async def async_int(n): + return n + + assert isinstance(async_int, AsyncTool) + assert await async_int.ainvoke(n=7) == 7 + + with pytest.raises(jsonschema.ValidationError): + await async_int.ainvoke(n="bad") + + +def test_input_and_output_schema_both_explicit() -> None: + """Both schemas can be explicit dicts at once.""" + + @tool( + "echo", + input_schema={ + "type": "object", + "properties": {"msg": {"type": "string"}}, + "required": ["msg"], + }, + output_schema={"type": "string"}, + ) + def echo(**kwargs) -> str: + return kwargs["msg"] + + assert echo.input_json_schema() == { + "type": "object", + "properties": {"msg": {"type": "string"}}, + "required": ["msg"], + } + assert echo.output_json_schema() == {"type": "string"} + assert echo.invoke(msg="hi") == "hi" + + with pytest.raises(jsonschema.ValidationError): + echo.invoke(msg=123) # input is not a string From c3c6d191da156c2a719da6adaeeca9a585a93f5b Mon Sep 17 00:00:00 2001 From: Elias Posen Date: Thu, 7 May 2026 17:14:18 -0400 Subject: [PATCH 3/5] make @tool name keyword-only --- pctx-py/src/pctx_client/_convert.py | 97 +++++++++----------- pctx-py/tests/scripts/manual_code_mode.py | 6 +- pctx-py/tests/test_create_input_schema.py | 30 +++---- pctx-py/tests/test_integration.py | 2 +- pctx-py/tests/test_output_schema.py | 26 +++--- pctx-py/tests/test_tool_decorator.py | 105 +++++++++++++--------- 6 files changed, 136 insertions(+), 130 deletions(-) diff --git a/pctx-py/src/pctx_client/_convert.py b/pctx-py/src/pctx_client/_convert.py index e62fd36c..bb07e5e6 100644 --- a/pctx-py/src/pctx_client/_convert.py +++ b/pctx-py/src/pctx_client/_convert.py @@ -8,27 +8,30 @@ @overload def tool( - name_or_callable: str, - *args: Any, + fn: Callable, + *, + name: str | None = None, namespace: str = "tools", description: str | None = None, input_schema: type[BaseModel] | dict[str, Any] | None = None, output_schema: Any | None = None, -) -> Callable[[Callable], Tool | AsyncTool]: ... +) -> Tool | AsyncTool: ... @overload def tool( - name_or_callable: Callable, - *args: Any, + fn: None = None, + *, + name: str | None = None, namespace: str = "tools", description: str | None = None, input_schema: type[BaseModel] | dict[str, Any] | None = None, output_schema: Any | None = None, -) -> Tool | AsyncTool: ... +) -> Callable[[Callable], Tool | AsyncTool]: ... def tool( - name_or_callable: str | Callable, - *args: Any, + fn: Callable | None = None, + *, + name: str | None = None, namespace: str = "tools", description: str | None = None, input_schema: type[BaseModel] | dict[str, Any] | None = None, @@ -39,7 +42,7 @@ def tool( Can be used with or without parameters: - @tool - Uses function name as tool name - - @tool("custom_name") - Uses custom name for the tool + - @tool(name="custom_name") - Uses custom name for the tool - @tool(namespace="custom", description="...") - With additional options - @tool(input_schema=MyModel) or @tool(input_schema={...}) - Override signature inference with an explicit Pydantic model or JSON Schema dict @@ -48,16 +51,20 @@ def tool( Python type / typing construct Args: - name_or_callable: Either a custom tool name (str) or the function to wrap (Callable) - namespace: The namespace the tool belongs to (default: "tools") - description: Optional description override (default: uses function docstring) + fn: The function to wrap. Only set when used as a bare ``@tool`` + decorator; in the parameterized form ``@tool(...)`` it is None + and the function is supplied on the second call. + name: Optional custom tool name (default: function's ``__name__``). + namespace: The namespace the tool belongs to (default: "tools"). + description: Optional description override (default: uses function docstring). input_schema: Optional explicit input schema. When provided, signature inference is skipped and this schema is used directly. output_schema: Optional explicit output schema. When provided, return- annotation inference is skipped and this schema is used directly. Returns: - Either a Tool/AsyncTool instance or a decorator function that creates one + Either a Tool/AsyncTool instance (bare form) or a decorator function + that creates one (parameterized form). Examples: >>> @tool @@ -65,7 +72,7 @@ def tool( ... '''Adds one to x''' ... return x + 1 - >>> @tool("custom_name", namespace="math") + >>> @tool(name="custom_name", namespace="math") ... def add_two(x: int) -> int: ... return x + 2 @@ -74,45 +81,25 @@ def tool( ... return kwargs["x"] + 1 """ - def _crate_tool_factory(tool_name: str) -> Callable[[Callable], Tool | AsyncTool]: - """ - Creates a decorator which takes the callable & returns the tool - - Args: - tool_name: the unique name of the tool - - Returns: - A function that takes a callable & returns a base tool - """ - - def _tool_factory(fn: Callable) -> Tool | AsyncTool: - return Tool.from_func( - func=fn, - name=tool_name, - namespace=namespace, - description=description, - input_schema=input_schema, - output_schema=output_schema, - ) - - return _tool_factory - - if len(args) != 0: - raise ValueError("Too many arguments for @tool decorator") - - if isinstance(name_or_callable, str): - # decorator used with params - # @tool("other_tool") - # def some_tool(): - # pass - return _crate_tool_factory(name_or_callable) - elif callable(name_or_callable) and hasattr(name_or_callable, "__name__"): - # decorator used without params - # @tool - # def some_tool(): - # pass - return _crate_tool_factory(name_or_callable.__name__)(name_or_callable) - else: - raise ValueError( - f"The first arg of the tool decorator must be a string or a callable with a __name__ attribute. Got {type(name_or_callable)}" + def _factory(f: Callable) -> Tool | AsyncTool: + return Tool.from_func( + func=f, + name=name, + namespace=namespace, + description=description, + input_schema=input_schema, + output_schema=output_schema, ) + + if fn is None: + # Parameterized form: @tool(name=..., input_schema=...) — return the + # factory so Python applies it to the decorated function on the next call. + return _factory + + if not callable(fn): + raise TypeError( + f"@tool's positional argument must be the decorated callable, got {type(fn).__name__}" + ) + + # Bare form: @tool — fn is the decorated callable, build the Tool now. + return _factory(fn) diff --git a/pctx-py/tests/scripts/manual_code_mode.py b/pctx-py/tests/scripts/manual_code_mode.py index 42e676fa..bd282614 100755 --- a/pctx-py/tests/scripts/manual_code_mode.py +++ b/pctx-py/tests/scripts/manual_code_mode.py @@ -22,13 +22,13 @@ def search_logs(query: str = "", level: str = "info", limit: int = 100) -> list[ ] -@tool("add", namespace="my_math") +@tool(name="add", namespace="my_math") def add(a: float, b: float) -> float: """adds two numbers""" return a + b -@tool("subtract", namespace="my_math") +@tool(name="subtract", namespace="my_math") def subtract(a: float, b: float) -> float: """subtracts b from a""" return a - b @@ -39,7 +39,7 @@ class MultiplyOutput(BaseModel): result: float -@tool("multiply", namespace="my_math") +@tool(name="multiply", namespace="my_math") def multiply(a: float, b: float) -> MultiplyOutput: """multiplies a and b""" return MultiplyOutput(message=f"Show your work! {a} * {b} = {a * b}", result=a * b) diff --git a/pctx-py/tests/test_create_input_schema.py b/pctx-py/tests/test_create_input_schema.py index b60810ec..f2c50847 100644 --- a/pctx-py/tests/test_create_input_schema.py +++ b/pctx-py/tests/test_create_input_schema.py @@ -1,9 +1,9 @@ -"""Tests for create_input_schema""" +"""Tests for infer_input_model""" import pytest from pydantic import BaseModel, ValidationError -from pctx_client._tool import create_input_schema +from pctx_client._tool import infer_input_model def test_simple_function_signature(): @@ -12,7 +12,7 @@ def test_simple_function_signature(): def add(a: int, b: int) -> int: return a + b - Model = create_input_schema("AddModel", add) + Model = infer_input_model("AddModel", add) # Test valid input instance = Model(a=5, b=10) @@ -34,7 +34,7 @@ def test_function_with_defaults(): def greet(name: str, greeting: str = "Hello") -> str: return f"{greeting}, {name}!" - Model = create_input_schema("GreetModel", greet) + Model = infer_input_model("GreetModel", greet) # Test with all parameters instance1 = Model(name="Alice", greeting="Hi") @@ -53,7 +53,7 @@ def test_function_without_type_hints(): def no_types(x, y=5): return x + y - Model = create_input_schema("NoTypesModel", no_types) + Model = infer_input_model("NoTypesModel", no_types) # Should work with Any type instance = Model(x="test", y=10) @@ -67,7 +67,7 @@ def test_function_with_args_kwargs(): def flexible(a: int, *args, b: str = "default", **kwargs) -> None: pass - Model = create_input_schema("FlexibleModel", flexible) + Model = infer_input_model("FlexibleModel", flexible) # Only 'a' and 'b' should be in the model instance = Model(a=42, b="test") @@ -91,7 +91,7 @@ def test_function_with_only_args(): def only_args(*args: int) -> None: pass - Model = create_input_schema("OnlyArgsModel", only_args) + Model = infer_input_model("OnlyArgsModel", only_args) # Model should have no fields schema = Model.model_json_schema() @@ -108,7 +108,7 @@ def test_function_with_only_kwargs(): def only_kwargs(**kwargs: str) -> None: pass - Model = create_input_schema("OnlyKwargsModel", only_kwargs) + Model = infer_input_model("OnlyKwargsModel", only_kwargs) # Model should have no fields schema = Model.model_json_schema() @@ -125,7 +125,7 @@ def test_function_with_positional_only_and_args(): def mixed(a: int, b: str, /, c: float = 1.0, *args) -> None: pass - Model = create_input_schema("MixedModel", mixed) + Model = infer_input_model("MixedModel", mixed) # Should include regular params but not args schema = Model.model_json_schema() @@ -146,7 +146,7 @@ def test_complex_types(): def process(items: list[str], count: int = 0) -> None: pass - Model = create_input_schema("ProcessModel", process) + Model = infer_input_model("ProcessModel", process) instance = Model(items=["a", "b", "c"], count=3) assert getattr(instance, "items") == ["a", "b", "c"] @@ -163,7 +163,7 @@ def test_model_is_basemodel(): def dummy(x: int) -> None: pass - Model = create_input_schema("DummyModel", dummy) + Model = infer_input_model("DummyModel", dummy) assert issubclass(Model, BaseModel) assert getattr(Model, "__name__") == "DummyModel" @@ -175,7 +175,7 @@ def test_model_json_schema(): def example(name: str, age: int, active: bool = True) -> None: pass - Model = create_input_schema("ExampleModel", example) + Model = infer_input_model("ExampleModel", example) schema = Model.model_json_schema() assert "properties" in schema @@ -192,7 +192,7 @@ async def fetch_data(url: str, timeout: int = 30) -> str: """Fetches data from a URL""" return f"Data from {url}" - Model = create_input_schema("FetchModel", fetch_data) + Model = infer_input_model("FetchModel", fetch_data) # Test with all parameters instance1 = Model(url="https://example.com", timeout=60) @@ -226,7 +226,7 @@ def create_user( ) -> None: pass - Model = create_input_schema("UserModel", create_user) + Model = infer_input_model("UserModel", create_user) # Test with valid nested models address_data = Address(street="123 Main St", city="Boston", zipcode="02101") @@ -269,7 +269,7 @@ def process_coordinates( ) -> None: pass - Model = create_input_schema("CoordinatesModel", process_coordinates) + Model = infer_input_model("CoordinatesModel", process_coordinates) # Test with valid tuples instance1 = Model(point=(10, 20), color=(100, 150, 200)) diff --git a/pctx-py/tests/test_integration.py b/pctx-py/tests/test_integration.py index b68481a5..afa21ff0 100644 --- a/pctx-py/tests/test_integration.py +++ b/pctx-py/tests/test_integration.py @@ -77,7 +77,7 @@ def greet(name: str, greeting: str = "Hello") -> str: """Greet someone with a custom greeting""" return f"{greeting}, {name}!" - @tool("foo_bar", namespace="namespaced_with_underscore") + @tool(name="foo_bar", namespace="namespaced_with_underscore") def namespaced_fn(val: str) -> str: return f"Hello {val}" diff --git a/pctx-py/tests/test_output_schema.py b/pctx-py/tests/test_output_schema.py index d487c769..1c9b83ac 100644 --- a/pctx-py/tests/test_output_schema.py +++ b/pctx-py/tests/test_output_schema.py @@ -1,10 +1,10 @@ -"""Tests for create_output_schema""" +"""Tests for infer_output_type""" from typing import Annotated, get_args, get_origin from pydantic import BaseModel, TypeAdapter -from pctx_client._tool import create_output_schema +from pctx_client._tool import infer_output_type def _unwrap_annotated(typ): @@ -20,7 +20,7 @@ def test_output_schema_simple_type(): def returns_int() -> int: return 42 - typ = create_output_schema(returns_int) + typ = infer_output_type(returns_int) adapter = TypeAdapter(typ) assert adapter.json_schema() == {"type": "integer"} @@ -31,7 +31,7 @@ def test_output_schema_string_type(): def returns_str() -> str: return "hello" - typ = create_output_schema(returns_str) + typ = infer_output_type(returns_str) adapter = TypeAdapter(typ) assert adapter.json_schema() == {"type": "string"} @@ -42,7 +42,7 @@ def test_output_schema_complex_type(): def returns_list() -> list[str]: return ["a", "b", "c"] - typ = create_output_schema(returns_list) + typ = infer_output_type(returns_list) adapter = TypeAdapter(typ) assert adapter.json_schema() == {"type": "array", "items": {"type": "string"}} @@ -53,7 +53,7 @@ def test_output_schema_dict_type(): def returns_dict() -> dict[str, int]: return {"a": 1, "b": 2} - typ = create_output_schema(returns_dict) + typ = infer_output_type(returns_dict) adapter = TypeAdapter(typ) assert adapter.json_schema() == { "type": "object", @@ -67,7 +67,7 @@ def test_output_schema_no_annotation(): def no_return_type(): return "something" - typ = create_output_schema(no_return_type) + typ = infer_output_type(no_return_type) adapter = TypeAdapter(typ) # Should use Any type - schema is empty dict for Any assert adapter.json_schema() == {} @@ -85,7 +85,7 @@ class UserOutput(BaseModel): def returns_model() -> UserOutput: return UserOutput(name="Alice", age=30) - typ = create_output_schema(returns_model) + typ = infer_output_type(returns_model) # Should be the same type (possibly wrapped in Annotated) assert _unwrap_annotated(typ) is UserOutput @@ -124,7 +124,7 @@ class Person(BaseModel): def returns_person() -> Person: return Person(name="Alice", address=Address(street="Main St", city="NYC")) - typ = create_output_schema(returns_person) + typ = infer_output_type(returns_person) # Should return Person as-is (possibly wrapped in Annotated) assert _unwrap_annotated(typ) is Person @@ -171,7 +171,7 @@ def test_output_schema_optional_type(): def returns_optional() -> str | None: return None - typ = create_output_schema(returns_optional) + typ = infer_output_type(returns_optional) adapter = TypeAdapter(typ) assert adapter.json_schema() == {"anyOf": [{"type": "string"}, {"type": "null"}]} @@ -182,7 +182,7 @@ def test_output_schema_union_type(): def returns_union() -> int | str: return 42 - typ = create_output_schema(returns_union) + typ = infer_output_type(returns_union) adapter = TypeAdapter(typ) assert adapter.json_schema() == {"anyOf": [{"type": "integer"}, {"type": "string"}]} @@ -193,7 +193,7 @@ def test_output_schema_bool_type(): def returns_bool() -> bool: return True - typ = create_output_schema(returns_bool) + typ = infer_output_type(returns_bool) adapter = TypeAdapter(typ) assert adapter.json_schema() == {"type": "boolean"} @@ -204,6 +204,6 @@ def test_output_schema_async_function(): async def async_returns_str() -> str: return "async result" - typ = create_output_schema(async_returns_str) + typ = infer_output_type(async_returns_str) adapter = TypeAdapter(typ) assert adapter.json_schema() == {"type": "string"} diff --git a/pctx-py/tests/test_tool_decorator.py b/pctx-py/tests/test_tool_decorator.py index 2dd52785..40a6232c 100644 --- a/pctx-py/tests/test_tool_decorator.py +++ b/pctx-py/tests/test_tool_decorator.py @@ -48,7 +48,7 @@ async def async_function() -> str: def test_registration_custom_name() -> None: """Test tool registration with custom name""" - @tool("custom_name") + @tool(name="custom_name") def my_function() -> str: """Function with custom name""" return "result" @@ -60,7 +60,7 @@ def my_function() -> str: def test_registration_custom_description() -> None: """Test tool registration with custom description""" - @tool("tool_name", description="Custom description here") + @tool(name="tool_name", description="Custom description here") def my_function() -> str: """Original docstring""" return "result" @@ -136,7 +136,7 @@ def no_doc() -> str: def test_registration_custom_description_overrides_docstring() -> None: """Test that custom description overrides docstring""" - @tool("func", description="Custom") + @tool(name="func", description="Custom") def with_docstring() -> str: """Original docstring""" return "result" @@ -165,37 +165,56 @@ def tool_two() -> str: assert tool_two.description == "Second tool" -def test_registration_error_too_many_arguments() -> None: - """Test that providing too many arguments raises ValueError""" +def test_direct_call_positional_fn_with_kwargs() -> None: + """``tool(fn, name=..., input_schema=...)`` — positional callable plus + keyword overrides — applies the overrides without going through the + decorator-factory branch.""" - with pytest.raises(ValueError, match="Too many arguments"): + class Args(BaseModel): + x: int + + def my_func(**kwargs) -> int: + return kwargs["x"] + 100 + + built = tool( + my_func, + name="custom", + namespace="explicit", + description="Adds 100 to x", + input_schema=Args, + output_schema={"type": "integer"}, + ) - @tool("name", "extra_arg") - def bad_function() -> str: - return "result" + assert isinstance(built, Tool) + assert built.name == "custom" + assert built.namespace == "explicit" + assert built.description == "Adds 100 to x" + assert built.input_schema is Args + assert built.output_json_schema() == {"type": "integer"} + assert built.invoke(x=5) == 105 -def test_registration_error_invalid_first_argument() -> None: - """Test that invalid first argument raises ValueError""" +def test_registration_error_too_many_positional_arguments() -> None: + """Passing more than one positional arg raises Python's own TypeError.""" - with pytest.raises( - ValueError, match="must be a string or a callable with a __name__" - ): - tool(123) # type: ignore + with pytest.raises(TypeError): + tool("name", "extra_arg") # type: ignore[call-overload] -def test_registration_error_callable_without_name() -> None: - """Test that callable without __name__ raises ValueError""" +def test_registration_error_name_as_positional_rejected() -> None: + """``name`` is keyword-only — passing it positionally is a TypeError.""" - class CallableWithoutName: - def __call__(self) -> str: - return "result" + with pytest.raises(TypeError): + # Strings used to be a valid first positional arg; under the new + # signature the only positional arg is the decorated callable. + tool("custom_name") # type: ignore[call-overload] + + +def test_registration_error_invalid_first_argument() -> None: + """A non-callable, non-None first positional argument raises TypeError.""" - obj = CallableWithoutName() - with pytest.raises( - ValueError, match="must be a string or a callable with a __name__" - ): - tool(obj) # type: ignore + with pytest.raises(TypeError, match="must be the decorated callable"): + tool(123) # type: ignore[call-overload] # ============================================================================ @@ -674,7 +693,7 @@ def test_input_schema_dict_skips_inference() -> None: "additionalProperties": False, } - @tool("from_dict", input_schema=schema) + @tool(name="from_dict", input_schema=schema) def from_dict(**kwargs) -> int: return kwargs["x"] + 1 @@ -688,7 +707,7 @@ def test_input_schema_dict_rejects_invalid_input() -> None: """Dict-defined input_schema validates via jsonschema and rejects bad input.""" @tool( - "add_one", + name="add_one", input_schema={ "type": "object", "properties": {"x": {"type": "integer"}}, @@ -710,8 +729,8 @@ def test_input_schema_dict_malformed_raises_at_construction() -> None: with pytest.raises(jsonschema.SchemaError): - @tool("bad", input_schema={"type": "not-a-real-type"}) - def bad(**kwargs) -> int: + @tool(name="bad", input_schema={"type": "not-a-real-type"}) + def bad(**kwargs) -> int: # pyright: ignore[reportUnusedFunction] return 0 @@ -721,7 +740,7 @@ def test_input_schema_pydantic_model_skips_inference() -> None: class Args(BaseModel): y: float - @tool("from_model", input_schema=Args) + @tool(name="from_model", input_schema=Args) def from_model(**kwargs) -> float: return kwargs["y"] * 2 @@ -738,7 +757,7 @@ def test_input_schema_pydantic_model_rejects_invalid_input() -> None: class Args(BaseModel): y: float - @tool("from_model", input_schema=Args) + @tool(name="from_model", input_schema=Args) def from_model(**kwargs) -> float: return kwargs["y"] * 2 @@ -757,7 +776,7 @@ def test_input_schema_explicit_overrides_signature() -> None: # Function signature says `n: int`, but we pass an unrelated schema. # The explicit schema wins; signature inference is skipped entirely. - @tool("search", input_schema=schema) + @tool(name="search", input_schema=schema) def search(n: int = 0, **kwargs) -> str: return kwargs.get("q", "") @@ -769,7 +788,7 @@ async def test_input_schema_dict_async() -> None: """Explicit JSON Schema dict works for async tools too.""" @tool( - "add_one_async", + name="add_one_async", input_schema={ "type": "object", "properties": {"x": {"type": "integer"}}, @@ -795,7 +814,7 @@ def test_input_schema_with_custom_name_and_namespace() -> None: "required": ["n"], } - @tool("add_one", namespace="math", input_schema=schema) + @tool(name="add_one", namespace="math", input_schema=schema) def named(**kwargs) -> int: return kwargs["n"] + 1 @@ -816,7 +835,7 @@ def test_output_schema_dict_skips_inference() -> None: schema: dict = {"type": "integer", "minimum": 0} # Function annotation says `str`, but explicit dict overrides it. - @tool("returns_int", output_schema=schema) + @tool(name="returns_int", output_schema=schema) def returns_int() -> str: return 42 # type: ignore[return-value] @@ -828,7 +847,7 @@ def test_output_schema_dict_rejects_invalid_output() -> None: """Dict-defined output_schema validates via jsonschema and rejects bad output.""" @tool( - "produces", + name="produces", output_schema={"type": "string", "minLength": 3}, ) def produces(value): @@ -848,7 +867,7 @@ def test_output_schema_dict_malformed_raises_at_construction() -> None: with pytest.raises(jsonschema.SchemaError): - @tool("bad_out", output_schema={"type": "not-a-real-type"}) + @tool(name="bad_out", output_schema={"type": "not-a-real-type"}) def _bad_out() -> int: # pyright: ignore[reportUnusedFunction] return 0 @@ -860,7 +879,7 @@ class Result(BaseModel): value: int label: str - @tool("returns_model", output_schema=Result) + @tool(name="returns_model", output_schema=Result) def returns_model() -> dict: return {"value": 1, "label": "one"} @@ -880,7 +899,7 @@ def test_output_schema_pydantic_model_rejects_invalid_output() -> None: class Result(BaseModel): value: int - @tool("returns_model", output_schema=Result) + @tool(name="returns_model", output_schema=Result) def returns_model(payload): return payload @@ -892,7 +911,7 @@ def test_output_schema_overrides_return_annotation() -> None: """When output_schema is provided, the function's return annotation is ignored.""" # Function annotated as -> int, but explicit schema says string. - @tool("override", output_schema={"type": "string"}) + @tool(name="override", output_schema={"type": "string"}) def override() -> int: return "hello" # type: ignore[return-value] @@ -903,7 +922,7 @@ def override() -> int: def test_output_schema_plain_python_type() -> None: """A plain python type (e.g. int) passed as output_schema works via TypeAdapter.""" - @tool("counter", output_schema=int) + @tool(name="counter", output_schema=int) def counter(n): return n @@ -917,7 +936,7 @@ def counter(n): async def test_output_schema_dict_async() -> None: """Explicit output JSON Schema works for async tools too.""" - @tool("async_int", output_schema={"type": "integer"}) + @tool(name="async_int", output_schema={"type": "integer"}) async def async_int(n): return n @@ -932,7 +951,7 @@ def test_input_and_output_schema_both_explicit() -> None: """Both schemas can be explicit dicts at once.""" @tool( - "echo", + name="echo", input_schema={ "type": "object", "properties": {"msg": {"type": "string"}}, From 3a9aa2604a2e8ab7fc1d450210ef848dc6812a98 Mon Sep 17 00:00:00 2001 From: Elias Posen Date: Thu, 7 May 2026 17:35:17 -0400 Subject: [PATCH 4/5] python-specific changelog with backfill & split makefile --- Makefile | 34 ++----- crates/pctx_session_server/openapi.json | 130 ++++++++++++++++++------ pctx-py/CHANGELOG.md | 96 +++++++++++++++++ pctx-py/Makefile | 39 +++++++ pctx-py/pyproject.toml | 2 +- pctx-py/uv.lock | 2 +- scripts/release.sh | 2 +- 7 files changed, 247 insertions(+), 58 deletions(-) create mode 100644 pctx-py/CHANGELOG.md create mode 100644 pctx-py/Makefile diff --git a/Makefile b/Makefile index f591ea1a..b53a9f15 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: help release publish-crates docs test-python test-python-integration format-python test-cli build-python +.PHONY: help release publish-crates docs test-cli # Default target - show help when running just 'make' .DEFAULT_GOAL := help @@ -7,14 +7,12 @@ help: @echo "pctx dev scripts" @echo "" @echo "Available targets:" - @echo " make docs - Generate CLI and Python documentation" - @echo " make test-python - Run Python client tests" - @echo " make test-python-integration - Run Python client tests with integration testing" - @echo " make format-python - Format and lint Python code with ruff" - @echo " make test-cli - Run CLI integration tests (pctx mcp start)" - @echo " make release - Interactive release script (bump version, update changelog)" - @echo " make publish-crates - Publish pctx_code_mode + dependencies to crates.io (runs locally)" - @echo " make build-python - Build Python package (resolves symlinks before build)" + @echo " make docs - Generate CLI, OpenAPI, and Python documentation" + @echo " make test-cli - Run CLI integration tests (pctx mcp start)" + @echo " make release - Interactive release script (bump version, update changelog)" + @echo " make publish-crates - Publish pctx_code_mode + dependencies to crates.io (runs locally)" + @echo "" + @echo "Python package targets live in pctx-py/Makefile (run 'make -C pctx-py help')." @echo "" # Generate CLI, OAS, and Python documentation @@ -24,21 +22,10 @@ docs: @./scripts/generate-openapi.sh @echo "" @echo "Building Python Sphinx documentation..." - @cd pctx-py && uv run sphinx-build -b html docs docs/_build/html + @$(MAKE) -C pctx-py docs @echo "" @echo "✓ Documentation built successfully!" -# Run Python client tests -test-python: - @cd pctx-py && uv run pytest tests/ -v - -# Run Python client tests with integration tests (expects pctx running on localhost on the default port) -test-python-integration: - @cd pctx-py && uv run pytest tests/ --integration -v - -format-python: - @cd pctx-py && uv run ruff format . && uv run ruff check . --fix - # Run CLI integration tests test-cli: @./scripts/test-mcp-cli.sh @@ -50,8 +37,3 @@ release: # Publish Rust crates to crates.io publish-crates: @./scripts/publish-crates.sh - -# Build Python package (resolves _tool_descriptions/data symlink before build, restores after) -build-python: - @./scripts/build-python.sh - diff --git a/crates/pctx_session_server/openapi.json b/crates/pctx_session_server/openapi.json index bde38fee..dc27a40d 100644 --- a/crates/pctx_session_server/openapi.json +++ b/crates/pctx_session_server/openapi.json @@ -12,7 +12,9 @@ "paths": { "/code-mode/execute-bash": { "post": { - "tags": ["CodeMode"], + "tags": [ + "CodeMode" + ], "summary": "Execute a bash command", "operationId": "execute_bash", "parameters": [ @@ -72,7 +74,9 @@ }, "/code-mode/functions/details": { "post": { - "tags": ["CodeMode"], + "tags": [ + "CodeMode" + ], "summary": "Get detailed information about a specific function", "operationId": "get_function_details", "parameters": [ @@ -132,7 +136,9 @@ }, "/code-mode/functions/list": { "post": { - "tags": ["CodeMode"], + "tags": [ + "CodeMode" + ], "summary": "List all available code mode functions from both server and tool registrations", "operationId": "list_functions", "parameters": [ @@ -172,7 +178,9 @@ }, "/code-mode/session/close": { "post": { - "tags": ["CodeMode"], + "tags": [ + "CodeMode" + ], "summary": "Close a `CodeMode` session", "operationId": "close_session", "parameters": [ @@ -222,7 +230,9 @@ }, "/code-mode/session/create": { "post": { - "tags": ["CodeMode"], + "tags": [ + "CodeMode" + ], "summary": "Create a new `CodeMode` session", "operationId": "create_session", "responses": { @@ -251,7 +261,9 @@ }, "/health": { "get": { - "tags": ["health"], + "tags": [ + "health" + ], "summary": "Health check endpoint", "operationId": "health", "responses": { @@ -270,7 +282,9 @@ }, "/register/servers": { "post": { - "tags": ["registration"], + "tags": [ + "registration" + ], "summary": "Register MCP servers dynamically at runtime", "operationId": "register_servers", "parameters": [ @@ -320,7 +334,9 @@ }, "/register/tools": { "post": { - "tags": ["registration"], + "tags": [ + "registration" + ], "summary": "Register tools that will be called via WebSocket callbacks", "operationId": "register_tools", "parameters": [ @@ -383,17 +399,25 @@ "schemas": { "CallbackConfig": { "type": "object", - "required": ["name"], + "required": [ + "name" + ], "properties": { "description": { - "type": ["string", "null"] + "type": [ + "string", + "null" + ] }, "input_schema": {}, "name": { "type": "string" }, "namespace": { - "type": ["string", "null"] + "type": [ + "string", + "null" + ] }, "output_schema": {} } @@ -401,7 +425,9 @@ "CloseSessionResponse": { "type": "object", "description": "Response after closing a `CodeMode` session", - "required": ["success"], + "required": [ + "success" + ], "properties": { "success": { "type": "boolean" @@ -411,7 +437,9 @@ "CreateSessionResponse": { "type": "object", "description": "Response after creating a new `CodeMode` session", - "required": ["session_id"], + "required": [ + "session_id" + ], "properties": { "session_id": { "type": "string" @@ -420,17 +448,27 @@ }, "ErrorCode": { "type": "string", - "enum": ["invalid_session", "internal", "execution"] + "enum": [ + "invalid_session", + "internal", + "execution" + ] }, "ErrorData": { "type": "object", - "required": ["code", "message"], + "required": [ + "code", + "message" + ], "properties": { "code": { "$ref": "#/components/schemas/ErrorCode" }, "details": { - "type": ["string", "null"] + "type": [ + "string", + "null" + ] }, "message": { "type": "string" @@ -449,7 +487,11 @@ }, "ExecuteBashOutput": { "type": "object", - "required": ["exit_code", "stdout", "stderr"], + "required": [ + "exit_code", + "stdout", + "stderr" + ], "properties": { "exit_code": { "type": "integer", @@ -473,7 +515,11 @@ }, { "type": "object", - "required": ["input_type", "output_type", "types"], + "required": [ + "input_type", + "output_type", + "types" + ], "properties": { "input_type": { "type": "string", @@ -493,7 +539,9 @@ }, "GetFunctionDetailsInput": { "type": "object", - "required": ["functions"], + "required": [ + "functions" + ], "properties": { "functions": { "type": "array", @@ -506,7 +554,10 @@ }, "GetFunctionDetailsOutput": { "type": "object", - "required": ["functions", "code"], + "required": [ + "functions", + "code" + ], "properties": { "code": { "type": "string" @@ -522,7 +573,10 @@ "HealthResponse": { "type": "object", "description": "Health check response", - "required": ["status", "version"], + "required": [ + "status", + "version" + ], "properties": { "status": { "type": "string" @@ -534,7 +588,10 @@ }, "ListFunctionsOutput": { "type": "object", - "required": ["functions", "code"], + "required": [ + "functions", + "code" + ], "properties": { "code": { "type": "string" @@ -550,10 +607,16 @@ }, "ListedFunction": { "type": "object", - "required": ["namespace", "name"], + "required": [ + "namespace", + "name" + ], "properties": { "description": { - "type": ["string", "null"], + "type": [ + "string", + "null" + ], "description": "Function description" }, "name": { @@ -569,7 +632,9 @@ "RegisterMcpServersRequest": { "type": "object", "description": "Request to register MCP servers", - "required": ["servers"], + "required": [ + "servers" + ], "properties": { "servers": { "type": "array", @@ -580,7 +645,10 @@ "RegisterMcpServersResponse": { "type": "object", "description": "Response after registering MCP servers", - "required": ["registered", "failed"], + "required": [ + "registered", + "failed" + ], "properties": { "failed": { "type": "array", @@ -597,7 +665,9 @@ "RegisterToolsRequest": { "type": "object", "description": "Request to register tools", - "required": ["tools"], + "required": [ + "tools" + ], "properties": { "tools": { "type": "array", @@ -610,7 +680,9 @@ "RegisterToolsResponse": { "type": "object", "description": "Response to registering tools", - "required": ["registered"], + "required": [ + "registered" + ], "properties": { "registered": { "type": "integer", @@ -630,4 +702,4 @@ "description": "Health check endpoints" } ] -} +} \ No newline at end of file diff --git a/pctx-py/CHANGELOG.md b/pctx-py/CHANGELOG.md new file mode 100644 index 00000000..f2bad05e --- /dev/null +++ b/pctx-py/CHANGELOG.md @@ -0,0 +1,96 @@ +# Changelog + +All notable changes to the `pctx-client` Python package will be documented in +this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this package adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). + +For changes to the underlying Rust crates and CLI, see the +[root CHANGELOG](../CHANGELOG.md). + +## [UNRELEASED] - YYYY-MM-DD + +### Added + +- `BaseTool.input_schema` now accepts a JSON Schema dict in addition to a + Pydantic `BaseModel` class. Dict-form schemas are validated via `jsonschema` + (Draft 2020-12); pydantic-form schemas continue to go through + `model_validate`. This lets integrators wrap tools defined externally + (MCP, OpenAPI, hand-written) without round-tripping through a synthetic + pydantic model. +- `BaseTool.output_schema` now accepts a JSON Schema dict in addition to + Python types / typing constructs. The `TypeAdapter` path is preserved for + rich Python output validation (`datetime`, `BaseModel` subclasses, etc.). +- `@tool` decorator: `input_schema=` and `output_schema=` keyword arguments to + override signature inference. Useful when wrapping a function whose + signature doesn't match the desired tool schema. +- `jsonschema>=4.26.0` added as a dependency. + +### Changed + +- **Breaking**: `@tool` decorator's `name` is now keyword-only. + `@tool("custom_name")` → `@tool(name="custom_name")`. The first positional + argument is reserved for the decorated callable so users can pass + `input_schema=...` / `output_schema=...` without also supplying a name. +- Dict-form schemas are validated against the JSON Schema metaschema and + compiled into a cached validator at tool construction. Malformed dict + schemas now raise `jsonschema.SchemaError` at definition time instead of on + first call. +- WebSocket client now surfaces `jsonschema.ValidationError` failures as + `INVALID_PARAMS` (alongside the existing `pydantic.ValidationError` path). + +## [v0.3.2] - 2026-04-06 + +### Changed + +- Documentation dependency group split out for ReadTheDocs builds. +- Internal version bump. + +## [v0.3.1] - 2026-03-25 + +### Added + +- `p.claude_agent_sdk_tools()` returns PCTX code mode tools as Claude Agent SDK tools, optional dependency requires `pctx[claude]` extra. + +## [v0.3.0] - 2026-03-12 + +### Added + +- `@tool` decorator now parses docstrings (Google, NumPy, reStructuredText, + and Epydoc formats) to extract parameter descriptions, return value + descriptions, and detailed function descriptions into tool schemas. +- Code mode config and all tools / descriptions easily configurable from + the Python client. +- `ToolDisclosure` support in the Python client and unified MCP via + `pctx mcp start`. + +### Changed + +- Improved code generation support for tools with no input schema or all + optional input schemas. + +## [v0.2.0] - 2026-01-12 + +### Added + +- Optional `search_functions` to allow the LLM to search for tools by + name/description before deciding which tool to call. + +## [v0.1.0] - 2025-12-16 + +### Added + +- Initial release of the `pctx-client` Python package. +- `@tool` decorator and `AsyncTool` / `Tool` base classes for registering + and interacting with the pctx session server. +- Convertors to export CodeMode tools to popular agent frameworks + (LangChain, CrewAI, OpenAI Agents, Pydantic AI). + +[UNRELEASED]: https://github.com/portofcontext/pctx/compare/pctx-py-v0.3.2...HEAD +[v0.3.2]: https://github.com/portofcontext/pctx/compare/pctx-py-v0.3.1...pctx-py-v0.3.2 +[v0.3.1]: https://github.com/portofcontext/pctx/compare/pctx-py-v0.3.0...pctx-py-v0.3.1 +[v0.3.0]: https://github.com/portofcontext/pctx/compare/pctx-py-v0.3.0b1...pctx-py-v0.3.0 +[v0.3.0b1]: https://github.com/portofcontext/pctx/compare/pctx-py-v0.2.0...pctx-py-v0.3.0b1 +[v0.2.0]: https://github.com/portofcontext/pctx/compare/pctx-py-v0.1.0...pctx-py-v0.2.0 +[v0.1.0]: https://github.com/portofcontext/pctx/releases/tag/pctx-py-v0.1.0 diff --git a/pctx-py/Makefile b/pctx-py/Makefile new file mode 100644 index 00000000..d355122f --- /dev/null +++ b/pctx-py/Makefile @@ -0,0 +1,39 @@ +.PHONY: help install test test-integration format build docs + +# Default target - show help when running just 'make' +.DEFAULT_GOAL := help + +help: + @echo "pctx-py dev scripts" + @echo "" + @echo "Available targets:" + @echo " make install - Install dev dependencies and all extras with uv" + @echo " make test - Run Python client tests" + @echo " make test-integration - Run Python client tests with integration testing" + @echo " make format - Format and lint Python code with ruff" + @echo " make build - Build Python package (resolves symlinks before build)" + @echo " make docs - Build Python Sphinx documentation" + @echo "" + +# Install with dev deps and all extras +install: + @uv sync --all-extras + +# Run Python client tests +test: + @uv run pytest tests/ -v + +# Run Python client tests with integration tests (expects pctx running on localhost on the default port) +test-integration: + @uv run pytest tests/ --integration -v + +format: + @uv run ruff format . && uv run ruff check . --fix + +# Build Python package (resolves _tool_descriptions/data symlink before build, restores after) +build: + @../scripts/build-python.sh + +# Build Python Sphinx documentation +docs: + @uv run --group docs sphinx-build -b html docs docs/_build/html diff --git a/pctx-py/pyproject.toml b/pctx-py/pyproject.toml index 0e6387bf..489d7836 100644 --- a/pctx-py/pyproject.toml +++ b/pctx-py/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "pctx-client" -version = "0.3.2" +version = "0.4.0" description = "Python client for using Code Mode via PCTX" readme = "README.md" authors = [ diff --git a/pctx-py/uv.lock b/pctx-py/uv.lock index b3c94f66..d22a6030 100644 --- a/pctx-py/uv.lock +++ b/pctx-py/uv.lock @@ -3339,7 +3339,7 @@ wheels = [ [[package]] name = "pctx-client" -version = "0.3.2" +version = "0.4.0" source = { editable = "." } dependencies = [ { name = "docstring-parser" }, diff --git a/scripts/release.sh b/scripts/release.sh index e9880d9b..322b9cc1 100755 --- a/scripts/release.sh +++ b/scripts/release.sh @@ -225,7 +225,7 @@ main() { echo " 2. Run 'cargo test' to ensure everything works" echo " 3. Commit and push the changes" echo " 3. Dispatch release from GitHub Actions" - echo " 4. If the Python SDK needs releasing, bump pyproject.toml, run make test-python (uv.lock) use the GH action manual dispatch" + echo " 4. If the Python SDK needs releasing, bump pyproject.toml, run make test (uv.lock) use the GH action manual dispatch" } From 396b07503175e7a31dfa21a7f24fe0887e861f91 Mon Sep 17 00:00:00 2001 From: Elias Posen Date: Thu, 7 May 2026 18:22:51 -0400 Subject: [PATCH 5/5] search functions input model --- pctx-py/src/pctx_client/_client.py | 6 ++---- pctx-py/src/pctx_client/models.py | 5 +++++ 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/pctx-py/src/pctx_client/_client.py b/pctx-py/src/pctx_client/_client.py index 7b3f8b5a..89393db3 100644 --- a/pctx-py/src/pctx_client/_client.py +++ b/pctx-py/src/pctx_client/_client.py @@ -26,6 +26,7 @@ GetFunctionDetailsOutput, ListedFunction, ListFunctionsOutput, + SearchFunctionsInput, ServerConfig, ToolConfig, ToolDisclosure, @@ -615,6 +616,7 @@ class SearchFunctionsTool(CrewAiBaseTool): description: str = get_tool_description( "search_functions", overrides=descriptions ) + args_schema: type[BaseModel] = SearchFunctionsInput def _run(_self, query: str, k: int = 10) -> str: return self._search_functions_result_to_string( @@ -911,10 +913,6 @@ async def get_function_details(args: dict[str, Any]) -> str: details = await self.get_function_details(tool_input.functions) return _text_content_block(details.code) - class SearchFunctionsInput(BaseModel): - query: str - k: int = 10 - @claude_tool( "search_functions", get_tool_description("search_functions", overrides=descriptions), diff --git a/pctx-py/src/pctx_client/models.py b/pctx-py/src/pctx_client/models.py index 1a6f0172..a3044dda 100644 --- a/pctx-py/src/pctx_client/models.py +++ b/pctx-py/src/pctx_client/models.py @@ -128,6 +128,11 @@ class GetFunctionDetailsInput(BaseModel): functions: list[str] +class SearchFunctionsInput(BaseModel): + query: str + k: int = 10 + + class GetFunctionDetailsOutput(BaseModel): """Output from getting detailed function information"""