diff --git a/src/celeste/client.py b/src/celeste/client.py index 375d701..2efd8e6 100644 --- a/src/celeste/client.py +++ b/src/celeste/client.py @@ -24,7 +24,7 @@ from celeste.models import Model from celeste.parameters import ParameterMapper, Parameters from celeste.streaming import Stream, enrich_stream_errors -from celeste.tools import ToolCall +from celeste.tools import ToolCall, validate_tool_calls from celeste.types import RawUsage @@ -232,7 +232,10 @@ async def _predict( ) content = self._parse_content(response_data) content = self._transform_output(content, **parameters) - tool_calls = self._parse_tool_calls(response_data) + tool_calls = validate_tool_calls( + self._parse_tool_calls(response_data), + parameters.get("tools"), + ) reasoning, signature = self._parse_reasoning(response_data) kwargs: dict[str, Any] = {} if reasoning is not None: diff --git a/src/celeste/streaming.py b/src/celeste/streaming.py index 983dbb7..59a267b 100644 --- a/src/celeste/streaming.py +++ b/src/celeste/streaming.py @@ -13,7 +13,7 @@ from celeste.io import Chunk as ChunkBase from celeste.io import FinishReason, Output, Usage from celeste.parameters import Parameters -from celeste.tools import ToolCall +from celeste.tools import ToolCall, validate_tool_calls from celeste.types import RawUsage @@ -175,12 +175,16 @@ def _parse_output(self, chunks: list[Chunk], **parameters: Unpack[Params]) -> Ou kwargs["reasoning"] = reasoning if signature: kwargs["signature"] = signature + tool_calls = validate_tool_calls( + self._aggregate_tool_calls(chunks, raw_events), + parameters.get("tools"), + ) output = self._output_class( content=content, usage=self._aggregate_usage(chunks), finish_reason=self._aggregate_finish_reason(chunks), metadata=self._build_stream_metadata(raw_events), - tool_calls=self._aggregate_tool_calls(chunks, raw_events), + tool_calls=tool_calls, **kwargs, ) return output # type: ignore[return-value] diff --git a/src/celeste/tools.py b/src/celeste/tools.py index bcd0c75..6dad22c 100644 --- a/src/celeste/tools.py +++ b/src/celeste/tools.py @@ -5,7 +5,9 @@ from typing import Any, ClassVar from pydantic import BaseModel, ConfigDict, Field +from pydantic import ValidationError as PydanticValidationError +from celeste.exceptions import ValidationError from celeste.types import Message, Role, ToolCall @@ -55,6 +57,61 @@ def map_tool(self, tool: Tool) -> dict[str, Any]: ... type ToolDefinition = Tool | dict[str, Any] +def _tool_parameter_models(tools: object | None) -> dict[str, type[BaseModel]]: + if not isinstance(tools, list): + return {} + + models: dict[str, type[BaseModel]] = {} + for tool in tools: + if not isinstance(tool, dict): + continue + name = tool.get("name") + parameters = tool.get("parameters") + if ( + isinstance(name, str) + and isinstance(parameters, type) + and issubclass(parameters, BaseModel) + ): + models[name] = parameters + return models + + +def validate_tool_calls( + tool_calls: list[ToolCall], + tools: object | None, +) -> list[ToolCall]: + """Validate returned tool calls against local Pydantic tool parameter models.""" + parameter_models = _tool_parameter_models(tools) + if not parameter_models: + return tool_calls + + validated_calls: list[ToolCall] = [] + for tool_call in tool_calls: + parameters_model = parameter_models.get(tool_call.name) + if parameters_model is None: + validated_calls.append(tool_call) + continue + + try: + validated_arguments = parameters_model.model_validate(tool_call.arguments) + except PydanticValidationError as exc: + raise ValidationError( + f"Tool call '{tool_call.name}' arguments failed validation: {exc}" + ) from exc + + validated_calls.append( + tool_call.model_copy( + update={ + "arguments": validated_arguments.model_dump( + mode="json", + exclude_unset=True, + ) + } + ) + ) + return validated_calls + + class ToolChoice(StrEnum): """Controls whether the model must use tools.""" @@ -102,4 +159,5 @@ class ToolError[Content](BaseModel): "ToolResult", "WebSearch", "XSearch", + "validate_tool_calls", ] diff --git a/tests/unit_tests/test_tool_call_validation.py b/tests/unit_tests/test_tool_call_validation.py new file mode 100644 index 0000000..b608cb0 --- /dev/null +++ b/tests/unit_tests/test_tool_call_validation.py @@ -0,0 +1,199 @@ +"""Tool-call argument validation tests.""" + +from collections.abc import AsyncIterator +from enum import StrEnum +from typing import Any, Unpack + +import pytest +from pydantic import BaseModel, SecretStr + +from celeste.auth import APIKey +from celeste.client import ModalityClient +from celeste.core import Modality, Provider +from celeste.exceptions import ValidationError +from celeste.io import Chunk, Input, Output +from celeste.modalities.text.parameters import TextParameter +from celeste.models import Model, Operation +from celeste.parameters import ParameterMapper, Parameters +from celeste.streaming import Stream +from celeste.tools import CodeExecution, ToolCall, validate_tool_calls + + +class ImageId(StrEnum): + IMG_1 = "img-1" + + +class AnalyzeImageParams(BaseModel): + image_id: ImageId + + +class NoFirstFrameParams(BaseModel): + first_frame: None = None + + +def _tool(parameters: type[BaseModel]) -> dict[str, object]: + return {"name": "analyze_image", "parameters": parameters} + + +def _model() -> Model: + return Model( + id="test-model", + provider=Provider.OPENAI, + display_name="Test Model", + operations={Modality.TEXT: {Operation.GENERATE}}, + ) + + +class _TextInput(Input): + prompt: str + + +class _ToolsMapper(ParameterMapper[str]): + name = TextParameter.TOOLS + + def map( + self, + request: dict[str, Any], + value: Any, # noqa: ANN401 + model: Model, + ) -> dict[str, Any]: + return request + + +class _Client(ModalityClient[_TextInput, Output[str], Parameters, str, Chunk[str]]): + returned_tool_calls: list[ToolCall] + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper[str]]: + return [_ToolsMapper()] + + def _init_request(self, inputs: _TextInput) -> dict[str, Any]: + return {"prompt": inputs.prompt} + + def _parse_usage( + self, response_data: dict[str, Any] + ) -> dict[str, int | float | None]: + return {} + + def _parse_content(self, response_data: dict[str, Any]) -> str: + return "ok" + + @classmethod + def _output_class(cls) -> type[Output[str]]: + return Output[str] + + async def _make_request( + self, + request_body: dict[str, Any], + *, + endpoint: str | None = None, + extra_headers: dict[str, str] | None = None, + **parameters: Unpack[Parameters], + ) -> dict[str, Any]: + return {} + + def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]: + return self.returned_tool_calls + + +class _Stream(Stream[Output[str], Parameters, Chunk[str]]): + _chunk_class = Chunk[str] + _output_class = Output[str] + _empty_content = "" + returned_tool_calls: list[ToolCall] + + def _aggregate_content(self, chunks: list[Chunk[str]]) -> str: + return "".join(chunk.content for chunk in chunks) + + def _aggregate_tool_calls( + self, + chunks: list[Chunk[str]], + raw_events: list[dict[str, Any]], + ) -> list[ToolCall]: + return self.returned_tool_calls + + +async def _empty_events() -> AsyncIterator[dict[str, Any]]: + if False: + yield {} + + +def test_validate_tool_calls_accepts_enum_and_preserves_extra_fields() -> None: + call = ToolCall( + id="call-1", + name="analyze_image", + arguments={"image_id": "img-1"}, + thoughtSignature="sig-1", + ) + + validated = validate_tool_calls([call], [_tool(AnalyzeImageParams)]) + + assert validated[0].arguments == {"image_id": "img-1"} + assert validated[0].thoughtSignature == "sig-1" + + +def test_validate_tool_calls_rejects_invalid_enum_argument() -> None: + call = ToolCall(id="call-1", name="analyze_image", arguments={"image_id": "fake"}) + + with pytest.raises(ValidationError, match=r"fake"): + validate_tool_calls([call], [_tool(AnalyzeImageParams)]) + + +@pytest.mark.parametrize("arguments", [{}, {"first_frame": None}]) +def test_validate_tool_calls_allows_omitted_or_null_null_only_argument( + arguments: dict[str, object], +) -> None: + call = ToolCall(id="call-1", name="analyze_image", arguments=arguments) + + validated = validate_tool_calls([call], [_tool(NoFirstFrameParams)]) + + assert validated[0].arguments == arguments + + +def test_validate_tool_calls_rejects_string_for_null_only_argument() -> None: + call = ToolCall( + id="call-1", + name="analyze_image", + arguments={"first_frame": "img-1"}, + ) + + with pytest.raises(ValidationError, match="first_frame"): + validate_tool_calls([call], [_tool(NoFirstFrameParams)]) + + +def test_validate_tool_calls_ignores_raw_schema_and_builtin_tools() -> None: + call = ToolCall(id="call-1", name="raw_tool", arguments={"anything": "ok"}) + + assert validate_tool_calls( + [call], + [{"name": "raw_tool", "parameters": {"type": "object"}}, CodeExecution()], + ) == [call] + + +async def test_predict_rejects_invalid_returned_tool_call() -> None: + client = _Client( + modality=Modality.TEXT, + model=_model(), + provider=Provider.OPENAI, + auth=APIKey(secret=SecretStr("test")), + returned_tool_calls=[ + ToolCall(id="call-1", name="analyze_image", arguments={"image_id": "fake"}) + ], + ) + + with pytest.raises(ValidationError, match="fake"): + await client._predict( + _TextInput(prompt="test"), tools=[_tool(AnalyzeImageParams)] + ) + + +def test_stream_output_rejects_invalid_returned_tool_call() -> None: + stream = _Stream(_empty_events(), tools=[_tool(AnalyzeImageParams)]) + stream.returned_tool_calls = [ + ToolCall(id="call-1", name="analyze_image", arguments={"image_id": "fake"}) + ] + + with pytest.raises(ValidationError, match="fake"): + stream._parse_output( + [Chunk[str](content="ok")], tools=[_tool(AnalyzeImageParams)] + )