Skip to content

Commit 3730ec0

Browse files
rustyconoverclaude
andcommitted
feat(arguments): declare per-argument constraints for agent discovery
Param and ConstParam gain choices/ge/le/gt/lt/pattern; scalar_function threads them through Arg, and argument_spec encodes them into Arrow field metadata (vgi_default/vgi_choices/vgi_range JSON+interval, vgi_pattern regex) surfaced by the C++ vgi_function_arguments() diagnostic. The fixture's format_number overloads now declare precision as ge=0/le=10 to exercise arg_range end to end. Metadata keys kept byte-for-byte in sync with vgi_protocol_constants.hpp on the C++ side; these are field-metadata keys, not RPC schema, so no codegen change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent e6782ce commit 3730ec0

5 files changed

Lines changed: 293 additions & 4 deletions

File tree

‎tests/test_argument_spec.py‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,13 +12,18 @@
1212
from vgi.argument_spec import (
1313
VGI_ARG_KEY,
1414
VGI_ARG_NAMED,
15+
VGI_CHOICES_KEY,
16+
VGI_DEFAULT_KEY,
1517
VGI_DOC_KEY,
18+
VGI_PATTERN_KEY,
19+
VGI_RANGE_KEY,
1620
VGI_TYPE_ANY,
1721
VGI_TYPE_KEY,
1822
VGI_TYPE_TABLE,
1923
VGI_VARARGS_KEY,
2024
VGI_VARARGS_TRUE,
2125
ArgumentSpec,
26+
_format_range,
2227
argument_specs_to_schema,
2328
extract_argument_specs,
2429
schema_to_argument_specs,
@@ -722,3 +727,99 @@ class FunctionWithDocs(TableInOutFunction): # type: ignore[type-arg]
722727
specs = {s.name: s for s in extract_argument_specs(FunctionWithDocs)}
723728
assert specs["unit"].doc == "Unit string, e.g. 'mi'"
724729
assert specs["scale"].doc == ""
730+
731+
732+
class TestArgumentConstraints:
733+
"""Per-argument constraint metadata (choices / range / pattern / default)."""
734+
735+
def test_format_range_notation(self) -> None:
736+
"""Interval notation covers inclusive, exclusive, open, and empty bounds."""
737+
assert _format_range(0, 100, None, None) == "[0, 100]"
738+
assert _format_range(None, None, 0, None) == "(0, +inf)"
739+
assert _format_range(1, None, None, 10) == "[1, 10)"
740+
assert _format_range(None, 5, None, None) == "(-inf, 5]"
741+
assert _format_range(None, None, None, None) is None
742+
743+
def test_writer_emits_constraint_keys(self) -> None:
744+
"""A constrained spec writes value-encoded constraint metadata keys."""
745+
spec = ArgumentSpec(
746+
name="unit",
747+
position="unit",
748+
arrow_type=pa.utf8(),
749+
default_json='"mm"',
750+
choices_json='["mm", "cm", "m"]',
751+
range_notation="[0, 100]",
752+
pattern="^[A-Z]{2}$",
753+
)
754+
meta = argument_specs_to_schema([spec])[0].metadata
755+
assert meta[VGI_DEFAULT_KEY] == b'"mm"'
756+
assert meta[VGI_CHOICES_KEY] == b'["mm", "cm", "m"]'
757+
assert meta[VGI_RANGE_KEY] == b"[0, 100]"
758+
assert meta[VGI_PATTERN_KEY] == b"^[A-Z]{2}$"
759+
760+
def test_unconstrained_arg_omits_keys(self) -> None:
761+
"""An arg without constraints yields none of the constraint keys."""
762+
meta = argument_specs_to_schema([ArgumentSpec(name="v", position=0, arrow_type=pa.int64())])[0].metadata
763+
for key in (VGI_DEFAULT_KEY, VGI_CHOICES_KEY, VGI_RANGE_KEY, VGI_PATTERN_KEY):
764+
assert meta is None or key not in meta
765+
766+
def test_constraints_round_trip(self) -> None:
767+
"""Constraint metadata survives the schema round-trip."""
768+
spec = ArgumentSpec(
769+
name="mode",
770+
position="mode",
771+
arrow_type=pa.utf8(),
772+
default_json='"add"',
773+
choices_json='["add", "mul"]',
774+
range_notation="[0, 10)",
775+
pattern="^x+$",
776+
)
777+
out = schema_to_argument_specs(argument_specs_to_schema([spec]))[0]
778+
assert out.default_json == '"add"'
779+
assert out.choices_json == '["add", "mul"]'
780+
assert out.range_notation == "[0, 10)"
781+
assert out.pattern == "^x+$"
782+
783+
def test_extract_constraints_from_arg(self) -> None:
784+
"""extract_argument_specs reads choices/default/bounds/pattern off Arg."""
785+
786+
class FunctionWithConstraints(TableInOutFunction): # type: ignore[type-arg]
787+
mode: str = Arg[str]( # type: ignore[assignment]
788+
"mode", doc="Blend mode", choices=["add", "mul"], default="add"
789+
)
790+
level: int = Arg[int](0, ge=1, le=10) # type: ignore[assignment]
791+
code: str = Arg[str](1, pattern="^[A-Z]{2}$") # type: ignore[assignment]
792+
793+
specs = {s.name: s for s in extract_argument_specs(FunctionWithConstraints)}
794+
assert specs["mode"].choices_json == '["add", "mul"]'
795+
assert specs["mode"].default_json == '"add"'
796+
assert specs["level"].range_notation == "[1, 10]"
797+
assert specs["code"].pattern == "^[A-Z]{2}$"
798+
799+
def test_const_param_constraints_surface(self) -> None:
800+
"""The modern Param/ConstParam API surfaces choices / range / pattern."""
801+
from vgi.arguments import ConstParam, Param, Returns
802+
from vgi.scalar_function import ScalarFunction
803+
804+
class ConstrainedScalar(ScalarFunction):
805+
class Meta:
806+
name = "constrained_scalar"
807+
description = "constraint demo"
808+
809+
@classmethod
810+
def compute(
811+
cls,
812+
unit: Annotated[str, ConstParam("Output unit", choices=["mm", "cm", "m"])],
813+
precision: Annotated[int, ConstParam("Decimals", ge=0, le=10)],
814+
code: Annotated[str, ConstParam("Label code", pattern="^[A-Z]{2}$")],
815+
value: Annotated[pa.DoubleArray, Param(doc="Value")],
816+
) -> Annotated[pa.StringArray, Returns()]:
817+
return value # never executed — the test only reads arg specs
818+
819+
specs = {s.name: s for s in extract_argument_specs(ConstrainedScalar)}
820+
assert specs["unit"].choices_json == '["mm", "cm", "m"]'
821+
assert specs["precision"].range_notation == "[0, 10]"
822+
assert specs["code"].pattern == "^[A-Z]{2}$"
823+
# The columnar Param carries no constraints.
824+
assert specs["value"].choices_json is None
825+
assert specs["value"].range_notation is None

‎vgi/_test_fixtures/scalar/formatting.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,7 @@ class Meta:
7777
@classmethod
7878
def compute(
7979
cls,
80-
precision: Annotated[int, ConstParam("Number of decimal places")],
80+
precision: Annotated[int, ConstParam("Number of decimal places", ge=0, le=10)],
8181
value: Annotated[pa.DoubleArray, Param(doc="Number to format")],
8282
) -> Annotated[pa.StringArray, Returns()]:
8383
"""Format each value with the specified precision."""
@@ -115,7 +115,7 @@ class Meta:
115115
@classmethod
116116
def compute(
117117
cls,
118-
precision: Annotated[int, ConstParam("Number of decimal places")],
118+
precision: Annotated[int, ConstParam("Number of decimal places", ge=0, le=10)],
119119
prefix: Annotated[str, ConstParam("Prefix string")],
120120
value: Annotated[pa.DoubleArray, Param(doc="Number to format")],
121121
) -> Annotated[pa.StringArray, Returns()]:

‎vgi/argument_spec.py‎

Lines changed: 146 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,15 @@
1414
- Special types ([`TableInput`][], [`AnyArrow`][], varargs) use field metadata markers
1515
"""
1616

17+
import json
1718
import warnings
1819
from collections.abc import Sequence
1920
from dataclasses import dataclass
2021
from typing import Annotated, Any, get_args, get_origin, get_type_hints
2122

2223
import pyarrow as pa
2324

24-
from vgi.arguments import PYTHON_TO_ARROW, AnyArrow, AnyArrowValue, Arg, TableInput
25+
from vgi.arguments import _MISSING, PYTHON_TO_ARROW, AnyArrow, AnyArrowValue, Arg, TableInput
2526

2627
__all__ = [
2728
"ArgumentSpec",
@@ -41,6 +42,10 @@
4142
"VGI_CONST_KEY",
4243
"VGI_CONST_TRUE",
4344
"VGI_DOC_KEY",
45+
"VGI_DEFAULT_KEY",
46+
"VGI_CHOICES_KEY",
47+
"VGI_RANGE_KEY",
48+
"VGI_PATTERN_KEY",
4449
]
4550

4651
# =============================================================================
@@ -69,6 +74,21 @@
6974
# ``vgi_doc_*`` prefix is reserved for future per-argument doc variants.
7075
VGI_DOC_KEY = b"vgi_doc"
7176

77+
# Keys carrying per-argument constraint metadata for agent discovery. All are
78+
# presence-only (omitted when the constraint is absent) and value-encoded as
79+
# UTF-8 so the column type stays uniform regardless of the arg's value type:
80+
# vgi_default — JSON scalar (the arg's default value; named/optional args only)
81+
# vgi_choices — JSON array (the closed set of allowed values)
82+
# vgi_range — interval notation built from ge/le/gt/lt (e.g. "[0, 100]",
83+
# "(0, +inf)", "[1, 10)"); a discovery surface, not the raw bounds
84+
# vgi_pattern — raw regex the value must match (open set)
85+
# Kept byte-for-byte in sync with the C++ reader in the vgi extension
86+
# (src/vgi_function_arguments_function.cpp / src/include/vgi_protocol_constants.hpp).
87+
VGI_DEFAULT_KEY = b"vgi_default"
88+
VGI_CHOICES_KEY = b"vgi_choices"
89+
VGI_RANGE_KEY = b"vgi_range"
90+
VGI_PATTERN_KEY = b"vgi_pattern"
91+
7292

7393
def _argument_spec_sort_key(spec: "ArgumentSpec") -> tuple[int, int | str]:
7494
"""Sort key: positional first (by index), then named (alphabetically)."""
@@ -109,6 +129,15 @@ class ArgumentSpec:
109129
doc: Optional human/agent-facing description of the argument. Surfaced
110130
through the catalog as the ``vgi_doc`` Arrow field metadata key
111131
(UTF-8); empty string means undocumented.
132+
default_json: Optional JSON-encoded default value (``vgi_default``);
133+
None when the argument is required.
134+
choices_json: Optional JSON array of allowed values (``vgi_choices``);
135+
None when the argument is not restricted to a closed set.
136+
range_notation: Optional interval-notation string built from the arg's
137+
numeric bounds (``vgi_range``, e.g. ``"[0, 100]"``); None when
138+
unbounded.
139+
pattern: Optional regex the value must match (``vgi_pattern``); None
140+
when there is no pattern constraint.
112141
113142
Note:
114143
For named arguments, the Python attribute name (``name``) and the SQL
@@ -131,6 +160,10 @@ class ArgumentSpec:
131160
is_varargs: bool = False
132161
is_const: bool = False
133162
doc: str = ""
163+
default_json: str | None = None
164+
choices_json: str | None = None
165+
range_notation: str | None = None
166+
pattern: str | None = None
134167

135168
def __repr__(self) -> str:
136169
"""Return concise repr showing key attributes."""
@@ -213,6 +246,16 @@ def argument_specs_to_schema(specs: Sequence[ArgumentSpec]) -> pa.Schema:
213246
if spec.doc:
214247
metadata[VGI_DOC_KEY] = spec.doc.encode("utf-8")
215248

249+
# Per-argument constraint metadata (presence-only; already value-encoded)
250+
if spec.default_json is not None:
251+
metadata[VGI_DEFAULT_KEY] = spec.default_json.encode("utf-8")
252+
if spec.choices_json is not None:
253+
metadata[VGI_CHOICES_KEY] = spec.choices_json.encode("utf-8")
254+
if spec.range_notation is not None:
255+
metadata[VGI_RANGE_KEY] = spec.range_notation.encode("utf-8")
256+
if spec.pattern is not None:
257+
metadata[VGI_PATTERN_KEY] = spec.pattern.encode("utf-8")
258+
216259
# Create field with or without metadata
217260
field = pa.field(
218261
spec.name,
@@ -266,6 +309,12 @@ def schema_to_argument_specs(schema: pa.Schema) -> list[ArgumentSpec]:
266309
doc_bytes = metadata.get(VGI_DOC_KEY)
267310
doc = doc_bytes.decode("utf-8") if doc_bytes else ""
268311

312+
# Per-argument constraint metadata (absent = unconstrained)
313+
default_bytes = metadata.get(VGI_DEFAULT_KEY)
314+
choices_bytes = metadata.get(VGI_CHOICES_KEY)
315+
range_bytes = metadata.get(VGI_RANGE_KEY)
316+
pattern_bytes = metadata.get(VGI_PATTERN_KEY)
317+
269318
specs.append(
270319
ArgumentSpec(
271320
name=field.name,
@@ -276,6 +325,10 @@ def schema_to_argument_specs(schema: pa.Schema) -> list[ArgumentSpec]:
276325
is_varargs=is_varargs,
277326
is_const=is_const,
278327
doc=doc,
328+
default_json=default_bytes.decode("utf-8") if default_bytes else None,
329+
choices_json=choices_bytes.decode("utf-8") if choices_bytes else None,
330+
range_notation=range_bytes.decode("utf-8") if range_bytes else None,
331+
pattern=pattern_bytes.decode("utf-8") if pattern_bytes else None,
279332
)
280333
)
281334

@@ -367,6 +420,94 @@ def macro_parameter_docs_from_schema(schema: pa.Schema) -> dict[str, str]:
367420
# =============================================================================
368421

369422

423+
def _format_range(
424+
ge: float | int | None,
425+
le: float | int | None,
426+
gt: float | int | None,
427+
lt: float | int | None,
428+
) -> str | None:
429+
"""Build interval notation from an argument's numeric bounds.
430+
431+
Inclusive bounds (``ge``/``le``) render as square brackets, exclusive bounds
432+
(``gt``/``lt``) as parentheses, and an open side as ``-inf``/``+inf``. Returns
433+
None when the argument has no numeric bound at all.
434+
435+
Args:
436+
ge: Inclusive lower bound (value >= ge).
437+
le: Inclusive upper bound (value <= le).
438+
gt: Exclusive lower bound (value > gt).
439+
lt: Exclusive upper bound (value < lt).
440+
441+
Returns:
442+
Interval string such as ``"[0, 100]"``, ``"(0, +inf)"``, ``"[1, 10)"``,
443+
or None when unbounded.
444+
445+
"""
446+
if ge is None and le is None and gt is None and lt is None:
447+
return None
448+
if gt is not None:
449+
low = f"({gt}"
450+
elif ge is not None:
451+
low = f"[{ge}"
452+
else:
453+
low = "(-inf"
454+
if lt is not None:
455+
high = f"{lt})"
456+
elif le is not None:
457+
high = f"{le}]"
458+
else:
459+
high = "+inf)"
460+
return f"{low}, {high}"
461+
462+
463+
def _constraint_kwargs(arg: Arg[Any]) -> dict[str, Any]:
464+
"""Extract discovery-facing constraint metadata from an [`Arg`][] descriptor.
465+
466+
Returns keyword arguments for [`ArgumentSpec`][] carrying only the constraints
467+
that are present, value-encoded for uniform surfacing: ``default``/``choices``
468+
as JSON, numeric bounds collapsed into a single ``range_notation`` interval
469+
string, ``pattern`` as-is. A defensively unserializable default falls back to
470+
its ``repr`` rather than dropping the whole registration.
471+
472+
Args:
473+
arg: The `Arg` descriptor to read constraints from.
474+
475+
Returns:
476+
A dict of `ArgumentSpec` keyword arguments (possibly empty).
477+
478+
"""
479+
kwargs: dict[str, Any] = {}
480+
481+
default = getattr(arg, "default", _MISSING)
482+
if default is not _MISSING:
483+
try:
484+
kwargs["default_json"] = json.dumps(default)
485+
except (TypeError, ValueError):
486+
kwargs["default_json"] = json.dumps(repr(default))
487+
488+
choices = getattr(arg, "choices", None)
489+
if choices is not None:
490+
try:
491+
kwargs["choices_json"] = json.dumps(list(choices))
492+
except (TypeError, ValueError):
493+
kwargs["choices_json"] = json.dumps([repr(c) for c in choices])
494+
495+
range_notation = _format_range(
496+
getattr(arg, "ge", None),
497+
getattr(arg, "le", None),
498+
getattr(arg, "gt", None),
499+
getattr(arg, "lt", None),
500+
)
501+
if range_notation is not None:
502+
kwargs["range_notation"] = range_notation
503+
504+
pattern = getattr(arg, "pattern", None)
505+
if pattern is not None:
506+
kwargs["pattern"] = pattern
507+
508+
return kwargs
509+
510+
370511
def extract_argument_specs(
371512
cls: type,
372513
) -> list[ArgumentSpec]:
@@ -414,6 +555,7 @@ def extract_argument_specs(
414555
is_varargs=param_arg.varargs,
415556
is_const=False,
416557
doc=param_arg.doc or "",
558+
**_constraint_kwargs(param_arg),
417559
)
418560
)
419561

@@ -431,6 +573,7 @@ def extract_argument_specs(
431573
is_varargs=const_arg.varargs,
432574
is_const=True,
433575
doc=const_arg.doc or "",
576+
**_constraint_kwargs(const_arg),
434577
)
435578
)
436579

@@ -496,6 +639,7 @@ def extract_argument_specs(
496639
is_varargs=arg_instance.varargs,
497640
is_const=getattr(arg_instance, "const", False),
498641
doc=getattr(arg_instance, "doc", "") or "",
642+
**_constraint_kwargs(arg_instance),
499643
)
500644
)
501645

@@ -571,6 +715,7 @@ def extract_argument_specs(
571715
is_varargs=is_varargs,
572716
is_const=is_const,
573717
doc=getattr(arg_legacy, "doc", "") or "",
718+
**_constraint_kwargs(arg_legacy),
574719
)
575720
)
576721

0 commit comments

Comments
 (0)