| import enum |
| import sys |
| import typing |
|
|
| import pytest |
|
|
| from needle import tool, Field |
| from needle.agent.tools import build_schema, pydantic_schema, _is_pydantic_model |
|
|
| pep604 = pytest.mark.skipif(sys.version_info < (3, 10), |
| reason="PEP 604 unions need Python 3.10+") |
|
|
|
|
| def test_basic_schema_types_and_required(): |
| @tool |
| def add(a: int, b: int) -> int: |
| "Add two numbers." |
| return a + b |
|
|
| schema = add._needle_tool |
| assert schema["name"] == "add" |
| assert schema["description"] == "Add two numbers." |
| assert schema["parameters"]["properties"]["a"] == {"type": "integer"} |
| assert set(schema["parameters"]["required"]) == {"a", "b"} |
|
|
|
|
| def test_defaults_are_optional(): |
| def f(city: str, units: str = "metric"): |
| pass |
|
|
| schema = build_schema(f) |
| assert schema["parameters"]["required"] == ["city"] |
| assert "units" in schema["parameters"]["properties"] |
|
|
|
|
| def test_literal_list_and_dict_types(): |
| def f(mode: typing.Literal["fast", "slow"], items: list, meta: dict): |
| pass |
|
|
| props = build_schema(f)["parameters"]["properties"] |
| assert props["mode"] == {"type": "string", "enum": ["fast", "slow"]} |
| assert props["items"]["type"] == "array" |
| assert props["meta"] == {"type": "object"} |
|
|
|
|
| def test_typed_list_items(): |
| def f(tags: typing.List[int]): |
| pass |
|
|
| props = build_schema(f)["parameters"]["properties"] |
| assert props["tags"] == {"type": "array", "items": {"type": "integer"}} |
|
|
|
|
| def test_optional_annotation_not_required(): |
| def f(a: str, b: typing.Optional[int] = None): |
| pass |
|
|
| schema = build_schema(f) |
| assert schema["parameters"]["required"] == ["a"] |
|
|
|
|
| def test_optional_annotated_field_constraints(): |
| def f( |
| value: typing.Optional[typing.Annotated[ |
| str, Field(pattern="^[a-z]+$", min_length=2, format="email")]] = None, |
| count: typing.Annotated[typing.Optional[int], Field(ge=0)] = None, |
| ): |
| pass |
|
|
| schema = build_schema(f) |
| properties = schema["parameters"]["properties"] |
| assert properties["value"] == { |
| "type": "string", |
| "pattern": "^[a-z]+$", |
| "minLength": 2, |
| "format": "email", |
| } |
| assert properties["count"] == {"type": "integer", "minimum": 0} |
| assert not {"value", "count"} & set(schema["parameters"].get("required", [])) |
|
|
|
|
| @pep604 |
| def test_pep604_none_union_not_required(): |
| def f(a: str, b: str | None = None, c: int | None = None, d: str | int | None = None): |
| pass |
|
|
| schema = build_schema(f) |
| assert schema["parameters"]["required"] == ["a"] |
| props = schema["parameters"]["properties"] |
| assert props["b"] == {"type": "string"} |
| assert props["c"] == {"type": "integer"} |
| assert props["d"] == {"type": "string"} |
|
|
|
|
| @pep604 |
| def test_pep604_union_matches_typing_union(): |
| def f(value: int | str): |
| pass |
|
|
| def g(value: typing.Union[int, str]): |
| pass |
|
|
| assert build_schema(f)["parameters"] == build_schema(g)["parameters"] |
| assert build_schema(f)["parameters"]["properties"]["value"] == {"type": "integer"} |
|
|
|
|
| def test_field_constraints_and_docstring_args(): |
| def f(temp: int = Field(description="temperature", ge=0, le=100), |
| name: str = Field(default="x", pattern="^[a-z]+$", min_length=1)): |
| """Set the thing. |
| |
| Args: |
| temp: the temperature in C |
| """ |
|
|
| schema = build_schema(f) |
| props = schema["parameters"]["properties"] |
| assert props["temp"]["minimum"] == 0 and props["temp"]["maximum"] == 100 |
| assert props["temp"]["description"] == "temperature" |
| assert props["name"]["pattern"] == "^[a-z]+$" |
| assert props["name"]["minLength"] == 1 |
| assert "temp" in schema["parameters"]["required"] |
| assert "name" not in schema["parameters"].get("required", []) |
|
|
|
|
| def test_docstring_description_falls_back_to_args(): |
| def f(city: str): |
| """Look up weather. |
| |
| Args: |
| city: the city to look up |
| """ |
|
|
| props = build_schema(f)["parameters"]["properties"] |
| assert props["city"]["description"] == "the city to look up" |
|
|
|
|
| def test_docstring_arg_description_keeps_its_wrapped_lines(): |
| |
| |
| |
| def f(time: str, label: str = ""): |
| """Set an alarm. |
| |
| Args: |
| time: The alarm time in 24-hour HH:MM form, for example 07:30 |
| or 18:45. Never a relative time like "in ten minutes". |
| label (str): A short name |
| shown when it rings. |
| |
| Returns: |
| dict: The alarm that was set. |
| """ |
|
|
| props = build_schema(f)["parameters"]["properties"] |
| assert props["time"]["description"] == ( |
| 'The alarm time in 24-hour HH:MM form, for example 07:30 ' |
| 'or 18:45. Never a relative time like "in ten minutes".') |
| assert props["label"]["description"] == "A short name shown when it rings." |
|
|
| |
| |
| def g(count: int): |
| """Count things. |
| |
| Args: |
| count: How many to count. |
| |
| Returns: |
| count: How many were counted. |
| """ |
|
|
| props = build_schema(g)["parameters"]["properties"] |
| assert props["count"]["description"] == "How many to count." |
|
|
|
|
| def test_pydantic_model_schema(): |
| import pydantic |
|
|
| class Weather(pydantic.BaseModel): |
| "Weather query." |
| city: str |
| units: str = "metric" |
|
|
| assert _is_pydantic_model(Weather) |
| schema = pydantic_schema(Weather) |
| assert schema["name"] == "Weather" |
| assert "city" in schema["parameters"]["properties"] |
| assert schema["parameters"]["required"] == ["city"] |
| assert schema["description"] == "Weather query." |
|
|
|
|
| def test_tool_decorator_preserves_callable(): |
| @tool |
| def greet(name: str) -> str: |
| "Greet someone." |
| return "hi " + name |
|
|
| assert greet("bob") == "hi bob" |
| assert greet._needle_tool["name"] == "greet" |
|
|
|
|
| def test_numeric_literal_and_int_enum_schema(): |
| class Level(enum.IntEnum): |
| LOW = 1 |
| MEDIUM = 2 |
| HIGH = 3 |
|
|
| def f( |
| port: typing.Literal[80, 443, 8080], |
| level: Level, |
| mode: typing.Literal["fast", "slow"] = "fast" |
| ): |
| pass |
|
|
| props = build_schema(f)["parameters"]["properties"] |
| assert props["port"] == {"type": "integer", "enum": [80, 443, 8080]} |
| assert props["level"] == {"type": "integer", "enum": [1, 2, 3]} |
| assert props["mode"] == {"type": "string", "enum": ["fast", "slow"]} |
|
|
|
|
| def test_annotated_optional_is_not_required(): |
| def f(city: str, units: typing.Annotated[typing.Optional[str], Field(description="metric or imperial")]): |
| pass |
|
|
| def g(city: str, units: typing.Optional[str]): |
| pass |
|
|
| assert build_schema(f)["parameters"]["required"] == ["city"] |
| assert build_schema(f)["parameters"]["required"] == build_schema(g)["parameters"]["required"] |
| assert build_schema(f)["parameters"]["properties"]["units"]["description"] == "metric or imperial" |
|
|
|
|
| @pep604 |
| def test_annotated_pep604_optional_is_not_required(): |
| def f(city: str, units: typing.Annotated[str | None, Field(description="metric or imperial")]): |
| pass |
|
|
| assert build_schema(f)["parameters"]["required"] == ["city"] |
|
|
