From 1a87d9aa4cad3d7550be2c730e71efe48516a5d2 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 11:35:38 +0100 Subject: [PATCH 1/3] feat(providers): add celeste-anthropic Messages API package Add standalone provider package for Anthropic Messages API with mixin pattern for capability-agnostic reuse. ## Client (AnthropicMessagesClient mixin) - HTTP POST/streaming to /v1/messages endpoint - Usage parsing: input_tokens, output_tokens, cached_tokens, total_tokens - Content array extraction with finish reason (stop_reason) - Beta header injection system for feature flags - Default max_tokens=1024 (Anthropic requires this field) ## Parameters - TemperatureMapper: temperature float [0.0-1.0] - TopPMapper: nucleus sampling top_p - TopKMapper: top-k sampling - MaxTokensMapper: max output tokens - StopSequencesMapper: custom stop sequences - ThinkingMapper: extended thinking ("auto" or budget_tokens int) - OutputSchemaMapper: native structured outputs via output_format - Supports single BaseModel and list[BaseModel] - Uses StrictJsonSchemaGenerator for schema generation - Auto-injects structured-outputs beta header ## Streaming (AnthropicMessagesStream mixin) - SSE event parsing: content_block_delta, message_delta, message_stop - Text delta extraction from content_block_delta - Stop reason extraction from message_delta - Usage tracking in final message events ## Config - API version: 2023-06-01 - Endpoints: /v1/messages, /v1/messages/count_tokens - Beta features: prompt-caching, computer-use, pdfs, token-counting, max-tokens-sonnet-3.5, structured-outputs --- packages/providers/anthropic/pyproject.toml | 17 ++ .../src/celeste_anthropic/__init__.py | 0 .../celeste_anthropic/messages/__init__.py | 1 + .../src/celeste_anthropic/messages/client.py | 158 ++++++++++++++ .../src/celeste_anthropic/messages/config.py | 42 ++++ .../celeste_anthropic/messages/parameters.py | 205 ++++++++++++++++++ .../celeste_anthropic/messages/streaming.py | 76 +++++++ .../anthropic/src/celeste_anthropic/py.typed | 0 8 files changed, 499 insertions(+) create mode 100644 packages/providers/anthropic/pyproject.toml create mode 100644 packages/providers/anthropic/src/celeste_anthropic/__init__.py create mode 100644 packages/providers/anthropic/src/celeste_anthropic/messages/__init__.py create mode 100644 packages/providers/anthropic/src/celeste_anthropic/messages/client.py create mode 100644 packages/providers/anthropic/src/celeste_anthropic/messages/config.py create mode 100644 packages/providers/anthropic/src/celeste_anthropic/messages/parameters.py create mode 100644 packages/providers/anthropic/src/celeste_anthropic/messages/streaming.py create mode 100644 packages/providers/anthropic/src/celeste_anthropic/py.typed diff --git a/packages/providers/anthropic/pyproject.toml b/packages/providers/anthropic/pyproject.toml new file mode 100644 index 00000000..63e51e9d --- /dev/null +++ b/packages/providers/anthropic/pyproject.toml @@ -0,0 +1,17 @@ +[project] +name = "celeste-anthropic" +version = "0.3.0" +description = "Anthropic provider package for Celeste AI" +authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] +license = {text = "Apache-2.0"} +requires-python = ">=3.12" + +[tool.uv.sources] +celeste-ai = { workspace = true } + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/celeste_anthropic"] diff --git a/packages/providers/anthropic/src/celeste_anthropic/__init__.py b/packages/providers/anthropic/src/celeste_anthropic/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/providers/anthropic/src/celeste_anthropic/messages/__init__.py b/packages/providers/anthropic/src/celeste_anthropic/messages/__init__.py new file mode 100644 index 00000000..57184785 --- /dev/null +++ b/packages/providers/anthropic/src/celeste_anthropic/messages/__init__.py @@ -0,0 +1 @@ +"""Anthropic Messages API provider package.""" diff --git a/packages/providers/anthropic/src/celeste_anthropic/messages/client.py b/packages/providers/anthropic/src/celeste_anthropic/messages/client.py new file mode 100644 index 00000000..5fb1189e --- /dev/null +++ b/packages/providers/anthropic/src/celeste_anthropic/messages/client.py @@ -0,0 +1,158 @@ +"""Anthropic Messages API client with shared implementation.""" + +from collections.abc import AsyncIterator +from typing import Any + +import httpx + +from celeste.core import UsageField +from celeste.io import FinishReason +from celeste.mime_types import ApplicationMimeType + +from . import config + + +class AnthropicMessagesClient: + """Mixin for Anthropic Messages API capabilities. + + Provides shared implementation for all capabilities using the Messages API: + - _make_request() - HTTP POST to /v1/messages + - _make_stream_request() - HTTP streaming to /v1/messages + - _parse_usage() - Extract usage dict from response + - _parse_content() - Extract content array from response + - _parse_finish_reason() - Extract finish reason from response + - _build_metadata() - Filter content fields + + Usage: + class AnthropicTextGenerationClient(AnthropicMessagesClient, TextGenerationClient): + def _parse_content(self, response_data, **parameters): + content = super()._parse_content(response_data) # Raw content array + for block in content: + if block.get("type") == "text": + return self._transform_output(block.get("text") or "", **parameters) + return "" + """ + + def _build_request( + self, + inputs: Any, + **parameters: Any, + ) -> Any: + """Build request with Anthropic-specific defaults.""" + request = super()._build_request(inputs, **parameters) # type: ignore[misc] + request["model"] = self.model.id # type: ignore[attr-defined] + + # Apply max_tokens default if not set (Anthropic requires it) + if "max_tokens" not in request: + request["max_tokens"] = config.DEFAULT_MAX_TOKENS + + return request + + def _build_headers(self, request_body: dict[str, Any]) -> dict[str, str]: + """Build headers with beta features extracted from request.""" + beta_features: list[str] = request_body.pop("_beta_features", []) + + headers: dict[str, str] = { + **self.auth.get_headers(), # type: ignore[attr-defined] + config.HEADER_ANTHROPIC_VERSION: config.ANTHROPIC_VERSION, + "Content-Type": ApplicationMimeType.JSON, + } + + if beta_features: + beta_values = [ + getattr(config, f"BETA_{f.upper().replace('-', '_')}") + for f in beta_features + ] + headers[config.HEADER_ANTHROPIC_BETA] = ",".join(beta_values) + + return headers + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> httpx.Response: + """Make HTTP request to Anthropic Messages API endpoint.""" + headers = self._build_headers(request_body) + + return await self.http_client.post( # type: ignore[attr-defined,no-any-return] + f"{config.BASE_URL}{config.AnthropicMessagesEndpoint.CREATE_MESSAGE}", + headers=headers, + json_body=request_body, + ) + + def _make_stream_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> AsyncIterator[dict[str, Any]]: + """Make streaming request to Anthropic Messages API endpoint.""" + request_body["stream"] = True + headers = self._build_headers(request_body) + + return self.http_client.stream_post( # type: ignore[attr-defined,no-any-return] + f"{config.BASE_URL}{config.AnthropicMessagesEndpoint.CREATE_MESSAGE}", + headers=headers, + json_body=request_body, + ) + + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: + """Map Anthropic usage fields to unified names. + + Shared by client and streaming across all capabilities. + """ + input_tokens = usage_data.get("input_tokens") + output_tokens = usage_data.get("output_tokens") + cached_tokens = usage_data.get("cache_read_input_tokens") + total_tokens = ( + (input_tokens + output_tokens) + if (input_tokens is not None and output_tokens is not None) + else None + ) + return { + UsageField.INPUT_TOKENS: input_tokens, + UsageField.OUTPUT_TOKENS: output_tokens, + UsageField.TOTAL_TOKENS: total_tokens, + UsageField.CACHED_TOKENS: cached_tokens, + } + + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Extract usage data from Messages API response.""" + usage_data = response_data.get("usage", {}) + return self.map_usage_fields(usage_data) + + def _parse_content(self, response_data: dict[str, Any]) -> Any: + """Parse content array from Messages API. + + Returns raw content array that capability clients extract from. + """ + content = response_data.get("content", []) + if not content: + msg = "No content in response" + raise ValueError(msg) + return content + + def _parse_finish_reason( + self, response_data: dict[str, Any] + ) -> FinishReason | None: + """Extract finish reason from Messages API response. + + Returns FinishReason that capability clients wrap in their specific type. + """ + stop_reason = response_data.get("stop_reason") + if stop_reason is None: + return None + + return FinishReason(reason=stop_reason) + + def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Build metadata dictionary, filtering out content field.""" + content_fields = {"content"} + filtered_data = { + k: v for k, v in response_data.items() if k not in content_fields + } + return super()._build_metadata(filtered_data) # type: ignore[misc,no-any-return] + + +__all__ = ["AnthropicMessagesClient"] diff --git a/packages/providers/anthropic/src/celeste_anthropic/messages/config.py b/packages/providers/anthropic/src/celeste_anthropic/messages/config.py new file mode 100644 index 00000000..c2acbe6e --- /dev/null +++ b/packages/providers/anthropic/src/celeste_anthropic/messages/config.py @@ -0,0 +1,42 @@ +"""Configuration for Anthropic Messages API.""" + +from enum import StrEnum + + +class AnthropicMessagesEndpoint(StrEnum): + """Endpoints for Messages API.""" + + CREATE_MESSAGE = "/v1/messages" + COUNT_MESSAGE_TOKENS = "/v1/messages/count_tokens" + + +BASE_URL = "https://api.anthropic.com" + +# Required +ANTHROPIC_VERSION = "2023-06-01" +CONTENT_TYPE_JSON = "application/json" + +# Header +HEADER_ANTHROPIC_VERSION = "anthropic-version" +HEADER_ANTHROPIC_BETA = "anthropic-beta" + +# Beta +BETA_PROMPT_CACHING = "prompt-caching-2024-07-31" +BETA_COMPUTER_USE = "computer-use-2024-10-22" +BETA_PDFS = "pdfs-2024-09-25" +BETA_TOKEN_COUNTING = "token-counting-2024-11-01" # nosec B105 +BETA_MAX_TOKENS_SONNET_3_5 = "max-tokens-3-5-sonnet-2024-07-15" +BETA_STRUCTURED_OUTPUTS = "structured-outputs-2025-11-13" + +# Defaults +DEFAULT_MAX_TOKENS = 1024 + +# SSE +SSE_EVENT_MESSAGE_START = "message_start" +SSE_EVENT_CONTENT_BLOCK_START = "content_block_start" +SSE_EVENT_PING = "ping" +SSE_EVENT_CONTENT_BLOCK_DELTA = "content_block_delta" +SSE_EVENT_CONTENT_BLOCK_STOP = "content_block_stop" +SSE_EVENT_MESSAGE_DELTA = "message_delta" +SSE_EVENT_MESSAGE_STOP = "message_stop" +SSE_EVENT_ERROR = "error" diff --git a/packages/providers/anthropic/src/celeste_anthropic/messages/parameters.py b/packages/providers/anthropic/src/celeste_anthropic/messages/parameters.py new file mode 100644 index 00000000..e948254f --- /dev/null +++ b/packages/providers/anthropic/src/celeste_anthropic/messages/parameters.py @@ -0,0 +1,205 @@ +"""Anthropic Messages API parameter mappers.""" + +import json +from typing import Any, get_args, get_origin + +from pydantic import BaseModel, TypeAdapter + +from celeste.models import Model +from celeste.parameters import ParameterMapper +from celeste.structured_outputs import StrictJsonSchemaGenerator +from celeste.types import StructuredOutput + + +class TemperatureMapper(ParameterMapper): + """Map temperature to Anthropic temperature field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform temperature into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["temperature"] = validated_value + return request + + +class TopPMapper(ParameterMapper): + """Map top_p to Anthropic top_p field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform top_p into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["top_p"] = validated_value + return request + + +class TopKMapper(ParameterMapper): + """Map top_k to Anthropic top_k field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform top_k into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["top_k"] = validated_value + return request + + +class MaxTokensMapper(ParameterMapper): + """Map max_tokens to Anthropic max_tokens field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform max_tokens into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["max_tokens"] = validated_value + return request + + +class StopSequencesMapper(ParameterMapper): + """Map stop_sequences to Anthropic stop_sequences field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform stop_sequences into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["stop_sequences"] = validated_value + return request + + +class ThinkingMapper(ParameterMapper): + """Map thinking to Anthropic thinking field. + + Accepts provider-native values: + - "auto": Dynamic budget (API decides) + - int: Fixed budget with specified tokens + """ + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinking into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + if validated_value == "auto": + request["thinking"] = {"type": "auto"} + else: + request["thinking"] = {"type": "enabled", "budget_tokens": validated_value} + return request + + +class OutputSchemaMapper(ParameterMapper): + """Map output_schema to Anthropic native structured outputs (output_format). + + Handles both single BaseModel and list[BaseModel] types. + Anthropic supports top-level arrays, $ref, and $defs natively. + """ + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform output_schema into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + origin = get_origin(validated_value) + if origin is list: + # Anthropic supports top-level arrays directly + inner_type = get_args(validated_value)[0] + inner_schema = TypeAdapter(inner_type).json_schema( + schema_generator=StrictJsonSchemaGenerator, + mode="serialization", + ) + schema = {"type": "array", "items": inner_schema} + else: + schema = TypeAdapter(validated_value).json_schema( + schema_generator=StrictJsonSchemaGenerator, + mode="serialization", + ) + + request["output_format"] = { + "type": "json_schema", + "schema": schema, + } + + # Signal that structured outputs beta header is needed + request.setdefault("_beta_features", []).append("structured-outputs") + + return request + + def parse_output( + self, content: StructuredOutput, value: object | None + ) -> StructuredOutput: + """Parse JSON to BaseModel using Pydantic's TypeAdapter.""" + if value is None: + return content if isinstance(content, str) else json.dumps(content) + + # If content is already a BaseModel, return it unchanged + if isinstance(content, BaseModel): + return content + if isinstance(content, list) and content and isinstance(content[0], BaseModel): + return content + + parsed: object + if isinstance(content, dict): + parsed = content + elif isinstance(content, str): + parsed = json.loads(content) + else: + parsed = content + + return TypeAdapter(value).validate_python(parsed) + + +__all__ = [ + "MaxTokensMapper", + "OutputSchemaMapper", + "StopSequencesMapper", + "TemperatureMapper", + "ThinkingMapper", + "TopKMapper", + "TopPMapper", +] diff --git a/packages/providers/anthropic/src/celeste_anthropic/messages/streaming.py b/packages/providers/anthropic/src/celeste_anthropic/messages/streaming.py new file mode 100644 index 00000000..cca7eab4 --- /dev/null +++ b/packages/providers/anthropic/src/celeste_anthropic/messages/streaming.py @@ -0,0 +1,76 @@ +"""Anthropic Messages SSE parsing for streaming.""" + +from typing import Any + +from .client import AnthropicMessagesClient + + +class AnthropicMessagesStream: + """Mixin for Messages API SSE parsing. + + Provides shared implementation for all capabilities using Anthropic Messages API streaming: + - _parse_chunk() - Parse SSE event into raw chunk dict + + Capability streams extend via super() to wrap results in typed Chunks. + + Usage: + class AnthropicTextGenerationStream(AnthropicMessagesStream, TextGenerationStream): + def _parse_chunk(self, event): + raw = super()._parse_chunk(event) + if not raw: + return None + return TextGenerationChunk(...) + """ + + def _parse_chunk(self, event: dict[str, Any]) -> dict[str, Any] | None: + """Parse SSE event into raw chunk data.""" + event_type = event.get("type") + if not event_type: + return None + + if event_type == "content_block_delta": + delta = event.get("delta", {}) + if delta.get("type") == "text_delta": + text_delta = delta.get("text") + if text_delta is not None: + return { + "content": text_delta, + "finish_reason": None, + "usage": None, + "raw_event": event, + } + return None + + if event_type == "message_delta": + delta = event.get("delta", {}) + stop_reason = delta.get("stop_reason") + + usage = None + usage_data = event.get("usage") + if usage_data: + usage = AnthropicMessagesClient.map_usage_fields(usage_data) + + return { + "content": "", + "finish_reason": stop_reason, + "usage": usage, + "raw_event": event, + } + + if event_type == "message_stop": + usage = None + usage_data = event.get("usage") + if usage_data: + usage = AnthropicMessagesClient.map_usage_fields(usage_data) + + return { + "content": "", + "finish_reason": None, + "usage": usage, + "raw_event": event, + } + + return None + + +__all__ = ["AnthropicMessagesStream"] diff --git a/packages/providers/anthropic/src/celeste_anthropic/py.typed b/packages/providers/anthropic/src/celeste_anthropic/py.typed new file mode 100644 index 00000000..e69de29b From 8acc8ce05f690d0be047a0d9aa3011ffd291ce90 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 11:39:11 +0100 Subject: [PATCH 2/3] feat(core): add structured outputs and usage field support Add core modules required by Anthropic provider package: ## structured_outputs.py - StrictJsonSchemaGenerator: adds additionalProperties:false (OpenAI, Anthropic, xAI) - RefResolvingJsonSchemaGenerator: resolves $ref inline (Cohere) - StrictRefResolvingJsonSchemaGenerator: combines both (Mistral) ## types.py - JsonValue: recursive JSON type alias - StructuredOutput: union of str | JsonValue | BaseModel | list[BaseModel] ## core.py - UsageField: standard field names for usage mapping - INPUT_TOKENS, OUTPUT_TOKENS, TOTAL_TOKENS - CACHED_TOKENS, REASONING_TOKENS, BILLED_TOKENS - NUM_IMAGES, BILLED_UNITS, INPUT_MP, OUTPUT_MP - CACHE_CREATION_INPUT_TOKENS, CACHE_READ_INPUT_TOKENS ## io.py - FinishReason: add reason field (str | None) --- src/celeste/core.py | 22 +++++++- src/celeste/io.py | 2 +- src/celeste/structured_outputs.py | 90 +++++++++++++++++++++++++++++++ src/celeste/types.py | 11 ++++ 4 files changed, 123 insertions(+), 2 deletions(-) create mode 100644 src/celeste/structured_outputs.py create mode 100644 src/celeste/types.py diff --git a/src/celeste/core.py b/src/celeste/core.py index 8165702a..fadd94be 100644 --- a/src/celeste/core.py +++ b/src/celeste/core.py @@ -56,4 +56,24 @@ class Parameter(StrEnum): MAX_TOKENS = "max_tokens" -__all__ = ["Capability", "Parameter", "Provider"] +class UsageField(StrEnum): + """Standard usage field names across Celeste capabilities. + + Use these when mapping provider usage fields to unified names. + """ + + INPUT_TOKENS = "input_tokens" + OUTPUT_TOKENS = "output_tokens" + TOTAL_TOKENS = "total_tokens" + CACHED_TOKENS = "cached_tokens" + REASONING_TOKENS = "reasoning_tokens" + BILLED_TOKENS = "billed_tokens" + NUM_IMAGES = "num_images" + BILLED_UNITS = "billed_units" + INPUT_MP = "input_mp" + OUTPUT_MP = "output_mp" + CACHE_CREATION_INPUT_TOKENS = "cache_creation_input_tokens" + CACHE_READ_INPUT_TOKENS = "cache_read_input_tokens" + + +__all__ = ["Capability", "Parameter", "Provider", "UsageField"] diff --git a/src/celeste/io.py b/src/celeste/io.py index c129cd4d..e85f60a0 100644 --- a/src/celeste/io.py +++ b/src/celeste/io.py @@ -14,7 +14,7 @@ class Input(BaseModel): class FinishReason(BaseModel): """Base class for capability-specific finish reasons (used in streaming chunks and outputs).""" - pass + reason: str | None = None class Usage(BaseModel): diff --git a/src/celeste/structured_outputs.py b/src/celeste/structured_outputs.py new file mode 100644 index 00000000..e1f9a02d --- /dev/null +++ b/src/celeste/structured_outputs.py @@ -0,0 +1,90 @@ +"""JSON Schema generators for structured outputs.""" + +from typing import Any + +from pydantic.json_schema import GenerateJsonSchema, JsonSchemaMode, JsonSchemaValue +from pydantic_core import CoreSchema + + +class StrictJsonSchemaGenerator(GenerateJsonSchema): + """Adds additionalProperties: false to all objects (OpenAI, Anthropic, xAI).""" + + def generate( + self, schema: CoreSchema, mode: JsonSchemaMode = "validation" + ) -> JsonSchemaValue: + json_schema = super().generate(schema, mode=mode) + self._add_strict_props(json_schema) + return json_schema + + def _add_strict_props(self, schema: dict) -> None: + if not isinstance(schema, dict): + return + if schema.get("type") == "object" and "additionalProperties" not in schema: + schema["additionalProperties"] = False + for key in ("properties", "$defs", "items", "allOf", "anyOf", "oneOf"): + if key in schema: + if isinstance(schema[key], dict): + for v in schema[key].values(): + self._add_strict_props(v) + elif isinstance(schema[key], list): + for item in schema[key]: + self._add_strict_props(item) + + +class RefResolvingJsonSchemaGenerator(GenerateJsonSchema): + """Resolves $ref references inline (Cohere - doesn't support $defs).""" + + def generate( + self, schema: CoreSchema, mode: JsonSchemaMode = "validation" + ) -> JsonSchemaValue: + json_schema = super().generate(schema, mode=mode) + self._resolve_refs(json_schema) + return json_schema + + def _resolve_refs(self, schema: dict) -> None: + if not isinstance(schema, dict): + return + defs = schema.pop("$defs", {}) + if not defs: + return + + def resolve(obj: Any) -> Any: # noqa: ANN401 + if isinstance(obj, dict): + if "$ref" in obj: + ref_path = obj["$ref"] + if ref_path.startswith("#/$defs/"): + ref_name = ref_path.split("/")[-1] + if ref_name in defs: + resolved = defs[ref_name].copy() + resolved.update( + {k: v for k, v in obj.items() if k != "$ref"} + ) + return resolve(resolved) + return {k: resolve(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [resolve(item) for item in obj] + return obj + + for key in list(schema.keys()): + schema[key] = resolve(schema[key]) + + +class StrictRefResolvingJsonSchemaGenerator( + StrictJsonSchemaGenerator, RefResolvingJsonSchemaGenerator +): + """Adds strict props AND resolves refs (Mistral - $defs support unconfirmed).""" + + def generate( + self, schema: CoreSchema, mode: JsonSchemaMode = "validation" + ) -> JsonSchemaValue: + json_schema = GenerateJsonSchema.generate(self, schema, mode=mode) + self._add_strict_props(json_schema) + self._resolve_refs(json_schema) + return json_schema + + +__all__ = [ + "RefResolvingJsonSchemaGenerator", + "StrictJsonSchemaGenerator", + "StrictRefResolvingJsonSchemaGenerator", +] diff --git a/src/celeste/types.py b/src/celeste/types.py new file mode 100644 index 00000000..ab316763 --- /dev/null +++ b/src/celeste/types.py @@ -0,0 +1,11 @@ +"""Type definitions for Celeste.""" + +from pydantic import BaseModel + +type JsonValue = ( + str | int | float | bool | None | dict[str, JsonValue] | list[JsonValue] +) + +type StructuredOutput = str | JsonValue | BaseModel | list[BaseModel] + +__all__ = ["JsonValue", "StructuredOutput"] From a8c55a276daba890a3e4e21226b22d314d555097 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Mon, 15 Dec 2025 11:42:57 +0100 Subject: [PATCH 3/3] test(core): add comprehensive tests for structured output generators Tests cover: - StrictJsonSchemaGenerator: additionalProperties injection - RefResolvingJsonSchemaGenerator: $ref inline resolution - StrictRefResolvingJsonSchemaGenerator: combined transformations 14 test cases covering: - Simple object schemas - Nested model hierarchies - Array item schemas - Non-dict input handling - $defs removal after resolution --- tests/unit_tests/test_structured_outputs.py | 265 ++++++++++++++++++++ 1 file changed, 265 insertions(+) create mode 100644 tests/unit_tests/test_structured_outputs.py diff --git a/tests/unit_tests/test_structured_outputs.py b/tests/unit_tests/test_structured_outputs.py new file mode 100644 index 00000000..e2fbe0ea --- /dev/null +++ b/tests/unit_tests/test_structured_outputs.py @@ -0,0 +1,265 @@ +"""Tests for JSON Schema generators for structured outputs.""" + +from pydantic import BaseModel, TypeAdapter + +from celeste.structured_outputs import ( + RefResolvingJsonSchemaGenerator, + StrictJsonSchemaGenerator, + StrictRefResolvingJsonSchemaGenerator, +) + + +class TestStrictJsonSchemaGenerator: + """Test StrictJsonSchemaGenerator adds additionalProperties: false.""" + + def test_adds_additional_properties_false_to_simple_object(self) -> None: + """Object schemas get additionalProperties: false.""" + + class SimpleModel(BaseModel): + name: str + + schema = TypeAdapter(SimpleModel).json_schema( + schema_generator=StrictJsonSchemaGenerator + ) + + assert schema.get("additionalProperties") is False + + def test_preserves_existing_additional_properties(self) -> None: + """Doesn't override existing additionalProperties when already set.""" + generator = StrictJsonSchemaGenerator(ref_template="{model}") + schema: dict = {"type": "object", "additionalProperties": True} + + generator._add_strict_props(schema) + + assert schema["additionalProperties"] is True + + def test_recursively_applies_to_nested_objects(self) -> None: + """Nested objects also get additionalProperties: false.""" + + class InnerModel(BaseModel): + value: int + + class OuterModel(BaseModel): + inner: InnerModel + + schema = TypeAdapter(OuterModel).json_schema( + schema_generator=StrictJsonSchemaGenerator + ) + + # Root should have additionalProperties: false + assert schema.get("additionalProperties") is False + + # Nested object in $defs should also have it + inner_def = schema.get("$defs", {}).get("InnerModel", {}) + assert inner_def.get("additionalProperties") is False + + def test_handles_array_items(self) -> None: + """Objects in array items get additionalProperties: false.""" + + class ItemModel(BaseModel): + id: int + + class ContainerModel(BaseModel): + items: list[ItemModel] + + schema = TypeAdapter(ContainerModel).json_schema( + schema_generator=StrictJsonSchemaGenerator + ) + + assert schema.get("additionalProperties") is False + # ItemModel in $defs should also have it + item_def = schema.get("$defs", {}).get("ItemModel", {}) + assert item_def.get("additionalProperties") is False + + def test_ignores_non_object_types(self) -> None: + """Non-object schemas don't get additionalProperties.""" + + class StringModel(BaseModel): + name: str + + schema = TypeAdapter(StringModel).json_schema( + schema_generator=StrictJsonSchemaGenerator + ) + + # String field itself shouldn't have additionalProperties + name_prop = schema.get("properties", {}).get("name", {}) + assert "additionalProperties" not in name_prop + + def test_handles_non_dict_input_gracefully(self) -> None: + """_add_strict_props handles non-dict input without error.""" + generator = StrictJsonSchemaGenerator(ref_template="{model}") + + # Should not raise + generator._add_strict_props(None) # type: ignore[arg-type] + generator._add_strict_props([1, 2, 3]) # type: ignore[arg-type] + generator._add_strict_props("string") # type: ignore[arg-type] + + +class TestRefResolvingJsonSchemaGenerator: + """Test RefResolvingJsonSchemaGenerator resolves $ref inline.""" + + def test_resolves_simple_ref(self) -> None: + """$ref replaced with definition content.""" + + class ReferencedModel(BaseModel): + x: int + + class ContainerModel(BaseModel): + ref: ReferencedModel + + schema = TypeAdapter(ContainerModel).json_schema( + schema_generator=RefResolvingJsonSchemaGenerator + ) + + # After resolution, no $ref should remain in properties + ref_prop = schema.get("properties", {}).get("ref", {}) + assert "$ref" not in ref_prop + + # The type should be inlined + assert ref_prop.get("type") == "object" + assert "properties" in ref_prop + assert "x" in ref_prop["properties"] + + def test_no_defs_returns_unchanged(self) -> None: + """Schema without $defs passes through unchanged.""" + + class SimpleModel(BaseModel): + name: str + + schema = TypeAdapter(SimpleModel).json_schema( + schema_generator=RefResolvingJsonSchemaGenerator + ) + + # Simple model shouldn't have $defs + assert "$defs" not in schema + assert schema.get("type") == "object" + + def test_removes_defs_after_resolving(self) -> None: + """$defs section is removed after refs are resolved.""" + + class ReferencedModel(BaseModel): + value: int + + class ContainerModel(BaseModel): + ref: ReferencedModel + + schema = TypeAdapter(ContainerModel).json_schema( + schema_generator=RefResolvingJsonSchemaGenerator + ) + + # $defs should be gone after resolution + assert "$defs" not in schema + + def test_resolves_nested_refs(self) -> None: + """References within resolved definitions are also resolved.""" + + class DeepModel(BaseModel): + deep: str + + class MiddleModel(BaseModel): + middle: DeepModel + + class OuterModel(BaseModel): + outer: MiddleModel + + schema = TypeAdapter(OuterModel).json_schema( + schema_generator=RefResolvingJsonSchemaGenerator + ) + + # All refs should be resolved, no $ref or $defs remaining + schema_str = str(schema) + assert "$ref" not in schema_str + assert "$defs" not in schema + + def test_handles_list_of_refs(self) -> None: + """References in array items are resolved.""" + + class ItemModel(BaseModel): + id: int + + class ContainerModel(BaseModel): + items: list[ItemModel] + + schema = TypeAdapter(ContainerModel).json_schema( + schema_generator=RefResolvingJsonSchemaGenerator + ) + + # No $ref should remain + assert "$ref" not in str(schema) + + +class TestStrictRefResolvingJsonSchemaGenerator: + """Test StrictRefResolvingJsonSchemaGenerator combines both transformations.""" + + def test_combines_both_transformations(self) -> None: + """Both strict props and ref resolution are applied.""" + + class InnerModel(BaseModel): + value: int + + class OuterModel(BaseModel): + inner: InnerModel + + schema = TypeAdapter(OuterModel).json_schema( + schema_generator=StrictRefResolvingJsonSchemaGenerator + ) + + # Refs should be resolved (no $ref or $defs) + assert "$ref" not in str(schema) + assert "$defs" not in schema + + # Root should have additionalProperties: false + assert schema.get("additionalProperties") is False + + # Inlined inner object should also have additionalProperties: false + inner_prop = schema.get("properties", {}).get("inner", {}) + assert inner_prop.get("additionalProperties") is False + + def test_with_deeply_nested_models(self) -> None: + """Works correctly with multiple levels of nesting.""" + + class Level3(BaseModel): + c: str + + class Level2(BaseModel): + b: Level3 + + class Level1(BaseModel): + a: Level2 + + schema = TypeAdapter(Level1).json_schema( + schema_generator=StrictRefResolvingJsonSchemaGenerator + ) + + # All refs resolved + assert "$ref" not in str(schema) + assert "$defs" not in schema + + # All levels have additionalProperties: false + assert schema.get("additionalProperties") is False + + level2 = schema.get("properties", {}).get("a", {}) + assert level2.get("additionalProperties") is False + + level3 = level2.get("properties", {}).get("b", {}) + assert level3.get("additionalProperties") is False + + def test_with_list_types(self) -> None: + """Handles list types with nested models correctly.""" + + class ItemModel(BaseModel): + name: str + + class ContainerModel(BaseModel): + items: list[ItemModel] + + schema = TypeAdapter(ContainerModel).json_schema( + schema_generator=StrictRefResolvingJsonSchemaGenerator + ) + + # All refs resolved + assert "$ref" not in str(schema) + assert "$defs" not in schema + + # Root has additionalProperties: false + assert schema.get("additionalProperties") is False