diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 26e9aa5e..86093d71 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -105,7 +105,7 @@ jobs: - uses: ./.github/actions/setup-python-uv with: python-version: ${{ matrix.python-version }} - - run: uv run pytest tests/ -v --cov=celeste --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=90 + - run: uv run pytest tests/unit_tests -v --cov=celeste --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=90 - uses: codecov/codecov-action@v4 if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.12' with: diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index c3291acf..beb57d8a 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -41,9 +41,26 @@ jobs: uses: ./.github/workflows/ci.yml secrets: inherit - build: + integration-test: needs: [validate-release, run-ci] runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 1 + - uses: ./.github/actions/setup-python-uv + - name: Run integration tests + env: + OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} + ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }} + GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }} + MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }} + COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }} + run: make integration-test + + build: + needs: [validate-release, run-ci, integration-test] + runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 with: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 4906f409..d372c061 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -53,7 +53,7 @@ repos: hooks: - id: pytest name: "๐Ÿงช Run tests with coverage" - entry: uv run pytest tests/ --cov=celeste --cov-report=term-missing + entry: uv run pytest tests/unit_tests --cov=celeste --cov-report=term-missing language: system types: [python] pass_filenames: false diff --git a/Makefile b/Makefile index 94c95e57..b4e30dfd 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: help sync lint lint-fix format typecheck test security ci clean +.PHONY: help sync lint lint-fix format typecheck test integration-test security ci clean # Default target help: @@ -8,6 +8,7 @@ help: @echo " make format - Apply Ruff formatting" @echo " make typecheck - Run mypy type checking" @echo " make test - Run pytest with coverage" + @echo " make integration-test - Run integration tests (requires API keys)" @echo " make security - Run Bandit security scan" @echo " make ci - Run full CI/CD pipeline" @echo " make clean - Clean cache directories" @@ -36,7 +37,11 @@ typecheck: # Testing test: - uv run pytest tests/ --cov=celeste --cov-report=term-missing --cov-fail-under=90 + uv run pytest tests/unit_tests --cov=celeste --cov-report=term-missing --cov-fail-under=90 + +# Integration testing (requires API keys) +integration-test: + uv run pytest tests/integration_tests/ -m integration -v --dist=worksteal -n auto # Security scanning (config reads from pyproject.toml) security: diff --git a/packages/text-generation/README.md b/packages/text-generation/README.md new file mode 100644 index 00000000..4f33a4d9 --- /dev/null +++ b/packages/text-generation/README.md @@ -0,0 +1,81 @@ +
+ +# Celeste Logo Celeste Text Generation + +**Text Generation capability for Celeste AI** + +[![Python](https://img.shields.io/badge/Python-3.12+-blue?style=for-the-badge)](https://www.python.org/) +[![License](https://img.shields.io/badge/License-Apache_2.0-red?style=for-the-badge)](../../LICENSE) + +[Quick Start](#-quick-start) โ€ข [Documentation](https://withceleste.ai/docs) โ€ข [Request Provider](https://github.com/withceleste/celeste-python/issues/new) + +
+ +--- + +## ๐Ÿš€ Quick Start + +```python +from celeste import create_client, Capability, Provider + +client = create_client( + capability=Capability.TEXT_GENERATION, + provider=Provider.OPENAI, +) + +response = await client.generate(prompt="Hello, world!") +print(response.content) +``` + +**Install:** +```bash +uv add "celeste-ai[text-generation]" +``` + +--- + +## Supported Providers + + +
+ +Google +Anthropic +OpenAI +Mistral +Cohere + + +**Missing a provider?** [Request it](https://github.com/withceleste/celeste-python/issues/new) โ€“ โšก **we ship fast**. + +
+ +--- + +**Streaming**: โœ… Supported + +**Parameters**: See [API Documentation](https://withceleste.ai/docs/api) for full parameter reference. + +--- + +## ๐Ÿค Contributing + +See [CONTRIBUTING.md](../../CONTRIBUTING.md) for guidelines. + +**Request a provider:** [GitHub Issues](https://github.com/withceleste/celeste-python/issues/new) + +--- + +## ๐Ÿ“„ License + +Apache 2.0 License โ€“ see [LICENSE](../../LICENSE) for details. + +--- + +
+ +**[Get Started](https://withceleste.ai/docs/quickstart)** โ€ข **[Documentation](https://withceleste.ai/docs)** โ€ข **[GitHub](https://github.com/withceleste/celeste-python)** + +Made with โค๏ธ by developers tired of framework lock-in + +
diff --git a/packages/text-generation/pyproject.toml b/packages/text-generation/pyproject.toml new file mode 100644 index 00000000..71a81002 --- /dev/null +++ b/packages/text-generation/pyproject.toml @@ -0,0 +1,39 @@ +[project] +name = "text-generation" +version = "0.0.3" +description = "Type-safe text generation for Celeste AI. Unified interface for OpenAI, Anthropic, Google, Mistral, Cohere, and more" +authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] +readme = "README.md" +license = {text = "Apache-2.0"} +requires-python = ">=3.12" +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Operating System :: OS Independent", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Typing :: Typed", +] +keywords = ["ai", "text-generation", "llm", "openai", "anthropic", "claude", "gemini", "mistral", "cohere"] + +[project.urls] +Homepage = "https://withceleste.ai" +Documentation = "https://withceleste.ai/docs" +Repository = "https://github.com/withceleste/celeste-python" +Issues = "https://github.com/withceleste/celeste-python/issues" + +[tool.uv.sources] +celeste-ai = { workspace = true } + +[project.entry-points."celeste.packages"] +text_generation = "celeste_text_generation:register_package" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/celeste_text_generation"] diff --git a/packages/text-generation/src/celeste_text_generation/__init__.py b/packages/text-generation/src/celeste_text_generation/__init__.py new file mode 100644 index 00000000..70cdb9e5 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/__init__.py @@ -0,0 +1,33 @@ +"""Celeste text generation capability.""" + + +def register_package() -> None: + """Register text generation package (client and models).""" + from celeste.client import register_client + from celeste.core import Capability + from celeste.models import register_models + from celeste_text_generation.models import MODELS + from celeste_text_generation.providers import PROVIDERS + + # Register provider-specific clients + for provider, client_class in PROVIDERS: + register_client(Capability.TEXT_GENERATION, provider, client_class) + + register_models(MODELS, capability=Capability.TEXT_GENERATION) + + +# Import after register_package is defined to avoid circular imports +from celeste_text_generation.io import ( # noqa: E402 + TextGenerationInput, + TextGenerationOutput, + TextGenerationUsage, +) +from celeste_text_generation.streaming import TextGenerationStream # noqa: E402 + +__all__ = [ + "TextGenerationInput", + "TextGenerationOutput", + "TextGenerationStream", + "TextGenerationUsage", + "register_package", +] diff --git a/packages/text-generation/src/celeste_text_generation/client.py b/packages/text-generation/src/celeste_text_generation/client.py new file mode 100644 index 00000000..144da2af --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/client.py @@ -0,0 +1,82 @@ +"""Base client for text generation.""" + +from abc import abstractmethod +from typing import Any, Unpack + +import httpx +from pydantic import BaseModel + +from celeste.client import Client +from celeste_text_generation.io import ( + TextGenerationFinishReason, + TextGenerationInput, + TextGenerationOutput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters + + +class TextGenerationClient( + Client[TextGenerationInput, TextGenerationOutput, TextGenerationParameters] +): + """Client for text generation operations.""" + + @abstractmethod + def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]: + """Initialize provider-specific request structure.""" + ... + + @abstractmethod + def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: + """Parse usage information from provider response.""" + ... + + @abstractmethod + def _parse_content( + self, + response_data: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> str | BaseModel: + """Parse content from provider response.""" + ... + + @abstractmethod + def _parse_finish_reason( + self, response_data: dict[str, Any] + ) -> TextGenerationFinishReason | None: + """Parse finish reason from provider response.""" + ... + + def _create_inputs( + self, *args: str, **parameters: Unpack[TextGenerationParameters] + ) -> TextGenerationInput: + """Map positional arguments to Input type.""" + if args: + return TextGenerationInput(prompt=args[0]) + prompt = parameters.get("prompt") + if prompt is None: + msg = ( + "prompt is required (either as positional argument or keyword argument)" + ) + raise TypeError(msg) + return TextGenerationInput(prompt=prompt) + + @classmethod + def _output_class(cls) -> type[TextGenerationOutput]: + """Return the Output class for this client.""" + return TextGenerationOutput + + def _build_metadata(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Build metadata dictionary from response data.""" + metadata = super()._build_metadata(response_data) + metadata["finish_reason"] = self._parse_finish_reason(response_data) + return metadata + + @abstractmethod + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> httpx.Response: + """Make HTTP request(s) and return response object.""" + ... diff --git a/packages/text-generation/src/celeste_text_generation/io.py b/packages/text-generation/src/celeste_text_generation/io.py new file mode 100644 index 00000000..9db36c39 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/io.py @@ -0,0 +1,57 @@ +"""Input and output types for text generation.""" + +from celeste.io import Chunk, FinishReason, Input, Output, Usage + + +class TextGenerationInput(Input): + """Input for text generation requests.""" + + prompt: str + + +class TextGenerationFinishReason(FinishReason): + """Text generation finish reason. + + Stores raw provider reason. Providers map their values in implementation. + """ + + reason: str # Raw provider string (e.g., "stop", "end_turn", "STOP", "COMPLETE") + + +class TextGenerationUsage(Usage): + """Text generation usage metrics. + + All fields optional since providers vary. + """ + + input_tokens: int | None = None + output_tokens: int | None = None + total_tokens: int | None = None + billed_tokens: int | None = None + cached_tokens: int | None = None + reasoning_tokens: int | None = None + + +class TextGenerationOutput[Content](Output[Content]): + """Output with text or structured content.""" + + pass + + +class TextGenerationChunk(Chunk[str]): + """Typed chunk for text generation streaming. + + Content is incremental text delta. + """ + + finish_reason: TextGenerationFinishReason | None = None + usage: TextGenerationUsage | None = None + + +__all__ = [ + "TextGenerationChunk", + "TextGenerationFinishReason", + "TextGenerationInput", + "TextGenerationOutput", + "TextGenerationUsage", +] diff --git a/packages/text-generation/src/celeste_text_generation/models.py b/packages/text-generation/src/celeste_text_generation/models.py new file mode 100644 index 00000000..899a4260 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/models.py @@ -0,0 +1,18 @@ +"""Model definitions for text generation.""" + +from celeste import Model +from celeste_text_generation.providers.anthropic.models import ( + MODELS as ANTHROPIC_MODELS, +) +from celeste_text_generation.providers.cohere.models import MODELS as COHERE_MODELS +from celeste_text_generation.providers.google.models import MODELS as GOOGLE_MODELS +from celeste_text_generation.providers.mistral.models import MODELS as MISTRAL_MODELS +from celeste_text_generation.providers.openai.models import MODELS as OPENAI_MODELS + +MODELS: list[Model] = [ + *ANTHROPIC_MODELS, + *COHERE_MODELS, + *GOOGLE_MODELS, + *MISTRAL_MODELS, + *OPENAI_MODELS, +] diff --git a/packages/text-generation/src/celeste_text_generation/parameters.py b/packages/text-generation/src/celeste_text_generation/parameters.py new file mode 100644 index 00000000..d4cb37eb --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/parameters.py @@ -0,0 +1,23 @@ +"""Parameters for text generation.""" + +from enum import StrEnum + +from pydantic import BaseModel + +from celeste.parameters import Parameters + + +class TextGenerationParameter(StrEnum): + """Unified parameter names for text generation capability.""" + + THINKING_BUDGET = "thinking_budget" + OUTPUT_SCHEMA = "output_schema" + + +class TextGenerationParameters(Parameters): + """Parameters for text generation.""" + + temperature: float | None + max_tokens: int | None + thinking_budget: int | None + output_schema: type[BaseModel] | None diff --git a/packages/text-generation/src/celeste_text_generation/providers/__init__.py b/packages/text-generation/src/celeste_text_generation/providers/__init__.py new file mode 100644 index 00000000..d884fc41 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/__init__.py @@ -0,0 +1,32 @@ +"""Provider implementations for text generation.""" + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from celeste.client import Client + from celeste.core import Provider + +__all__ = ["PROVIDERS"] + + +def _get_providers() -> list[tuple["Provider", type["Client"]]]: + """Lazy-load providers.""" + from celeste.core import Provider + from celeste_text_generation.providers.anthropic import ( + AnthropicTextGenerationClient, + ) + from celeste_text_generation.providers.cohere import CohereTextGenerationClient + from celeste_text_generation.providers.google import GoogleTextGenerationClient + from celeste_text_generation.providers.mistral import MistralTextGenerationClient + from celeste_text_generation.providers.openai import OpenAITextGenerationClient + + return [ + (Provider.ANTHROPIC, AnthropicTextGenerationClient), + (Provider.COHERE, CohereTextGenerationClient), + (Provider.GOOGLE, GoogleTextGenerationClient), + (Provider.MISTRAL, MistralTextGenerationClient), + (Provider.OPENAI, OpenAITextGenerationClient), + ] + + +PROVIDERS: list[tuple["Provider", type["Client"]]] = _get_providers() diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/__init__.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/__init__.py new file mode 100644 index 00000000..62028733 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/__init__.py @@ -0,0 +1,7 @@ +"""Anthropic provider.""" + +from .client import AnthropicTextGenerationClient +from .models import MODELS +from .streaming import AnthropicTextGenerationStream + +__all__ = ["MODELS", "AnthropicTextGenerationClient", "AnthropicTextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py new file mode 100644 index 00000000..3bd90a5e --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/client.py @@ -0,0 +1,138 @@ +"""Anthropic client implementation.""" + +from collections.abc import AsyncIterator +from typing import Any, Unpack + +import httpx +from pydantic import BaseModel + +from celeste.parameters import ParameterMapper +from celeste_text_generation.client import TextGenerationClient +from celeste_text_generation.io import ( + TextGenerationFinishReason, + TextGenerationInput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters + +from . import config +from .parameters import ANTHROPIC_PARAMETER_MAPPERS +from .streaming import AnthropicTextGenerationStream + + +class AnthropicTextGenerationClient(TextGenerationClient): + """Anthropic client.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + return ANTHROPIC_PARAMETER_MAPPERS + + def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]: + """Initialize request from Anthropic Messages API format.""" + return {"messages": [{"role": "user", "content": inputs.prompt}]} + + def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: + """Parse usage from response.""" + usage_data = response_data.get("usage", {}) + input_tokens = usage_data.get("input_tokens") + output_tokens = usage_data.get("output_tokens") + + total_tokens = None + if input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens + + return TextGenerationUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + billed_tokens=None, + cached_tokens=None, + reasoning_tokens=None, + ) + + def _parse_content( + self, + response_data: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> str | BaseModel: + """Parse content from response.""" + content = response_data.get("content", []) + if not content: + msg = "No content blocks in response" + raise ValueError(msg) + + output_schema = parameters.get("output_schema") + if output_schema is not None: + for content_block in content: + if content_block.get("type") == "tool_use": + tool_input = content_block.get("input") + if tool_input is not None: + return self._transform_output(tool_input, **parameters) + + text_content = "" + for content_block in content: + if content_block.get("type") == "text": + text_content = content_block.get("text") or "" + break + + return self._transform_output(text_content, **parameters) + + def _parse_finish_reason( + self, response_data: dict[str, Any] + ) -> TextGenerationFinishReason | None: + """Parse finish reason from response.""" + stop_reason = response_data.get("stop_reason") + if stop_reason is None: + return None + + return TextGenerationFinishReason(reason=stop_reason) + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> httpx.Response: + """Make HTTP request(s) and return response object.""" + request_body["model"] = self.model.id + request_body["max_tokens"] = parameters.get("max_tokens") or 1024 + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + config.ANTHROPIC_VERSION_HEADER: config.ANTHROPIC_VERSION, + "Content-Type": "application/json", + } + + return await self.http_client.post( + f"{config.BASE_URL}{config.ENDPOINT}", + headers=headers, + json_body=request_body, + ) + + def _stream_class(self) -> type[AnthropicTextGenerationStream]: + """Return the Stream class for this client.""" + return AnthropicTextGenerationStream + + def _make_stream_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> AsyncIterator[dict[str, Any]]: + """Make HTTP streaming request and return async iterator of events.""" + request_body["model"] = self.model.id + request_body["max_tokens"] = parameters.get("max_tokens") or 1024 + request_body["stream"] = True + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + config.ANTHROPIC_VERSION_HEADER: config.ANTHROPIC_VERSION, + "Content-Type": "application/json", + } + + return self.http_client.stream_post( + f"{config.BASE_URL}{config.STREAM_ENDPOINT}", + headers=headers, + json_body=request_body, + ) + + +__all__ = ["AnthropicTextGenerationClient"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/config.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/config.py new file mode 100644 index 00000000..7c1e8971 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/config.py @@ -0,0 +1,14 @@ +"""Anthropic provider configuration.""" + +# HTTP Configuration +BASE_URL = "https://api.anthropic.com" +ENDPOINT = "/v1/messages" +STREAM_ENDPOINT = ENDPOINT + +# Authentication +AUTH_HEADER_NAME = "x-api-key" +AUTH_HEADER_PREFIX = "" + +# API Version Header (required by Anthropic) +ANTHROPIC_VERSION_HEADER = "anthropic-version" +ANTHROPIC_VERSION = "2023-06-01" diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/models.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/models.py new file mode 100644 index 00000000..d87b6e1a --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/models.py @@ -0,0 +1,67 @@ +"""Anthropic models.""" + +from celeste import Model, Provider +from celeste.constraints import Range, Schema +from celeste_text_generation.parameters import TextGenerationParameter + +MODELS: list[Model] = [ + Model( + id="claude-sonnet-4-5", + provider=Provider.ANTHROPIC, + display_name="Claude Sonnet 4.5", + streaming=True, + parameter_constraints={ + TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=64000), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="claude-haiku-4-5", + provider=Provider.ANTHROPIC, + display_name="Claude Haiku 4.5", + streaming=True, + parameter_constraints={ + TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=32000), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="claude-opus-4-1", + provider=Provider.ANTHROPIC, + display_name="Claude Opus 4.1", + streaming=True, + parameter_constraints={ + TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=32000), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="claude-sonnet-4-20250514", + provider=Provider.ANTHROPIC, + display_name="Claude Sonnet 4", + streaming=True, + parameter_constraints={ + TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=64000), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="claude-sonnet-3-7", + provider=Provider.ANTHROPIC, + display_name="Claude Sonnet 3.7", + streaming=True, + parameter_constraints={ + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="claude-opus-4-20250514", + provider=Provider.ANTHROPIC, + display_name="Claude Opus 4", + streaming=True, + parameter_constraints={ + TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=32000), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), +] diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py new file mode 100644 index 00000000..f263eab7 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/parameters.py @@ -0,0 +1,288 @@ +"""Anthropic parameter mappers.""" + +import json +from enum import StrEnum +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_text_generation.parameters import TextGenerationParameter + + +class ThinkingBudgetMapper(ParameterMapper): + """Map thinking_budget parameter to Anthropic thinking.budget_tokens.""" + + name: StrEnum = TextGenerationParameter.THINKING_BUDGET + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinking_budget into provider request. + + Maps unified thinking_budget to Anthropic thinking parameter: + - thinking_budget=None โ†’ No thinking parameter (thinking disabled) + - thinking_budget=-1 โ†’ {"type": "auto"} (dynamic budget, automatic) + - thinking_budget=N (where N >= 1024) โ†’ {"type": "enabled", "budget_tokens": N} (fixed budget) + + Args: + request: Provider request dict. + value: thinking_budget value (int | None). + model: Model instance containing parameter_constraints for validation. + model: Model instance with parameter constraints for validation. + + Returns: + Updated request dict with thinking parameter if value provided. + + Raises: + ValueError: If value is not -1 and is less than 1024 (minimum required). + """ + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + # Build thinking parameter object + if validated_value == -1: + # Dynamic thinking: use "auto" type (no budget_tokens) + thinking_config: dict[str, Any] = {"type": "auto"} + else: + # Fixed budget: validate minimum is 1024 + if validated_value < 1024: + msg = f"thinking_budget must be -1 (dynamic) or >= 1024 for {model.id}, got {validated_value}" + raise ValueError(msg) + thinking_config = {"type": "enabled", "budget_tokens": validated_value} + + request["thinking"] = thinking_config + return request + + +class OutputSchemaMapper(ParameterMapper): + """Map output_schema parameter to Anthropic tools parameter (tool-based structured output).""" + + name: StrEnum = TextGenerationParameter.OUTPUT_SCHEMA + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform output_schema into provider request. + + Converts unified output_schema to Anthropic tools parameter: + - Creates a single tool definition with input_schema matching the output schema + - Sets tool_choice to force tool use + - Handles both BaseModel and list[BaseModel] types + + Args: + request: Provider request dict. + value: output_schema value (type[BaseModel] | None). + model: Model instance containing parameter_constraints for validation. + model: Model instance with parameter constraints for validation. + + Returns: + Updated request dict with tools and tool_choice if value provided. + """ + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + # Convert Pydantic model to JSON Schema + schema = self._convert_to_anthropic_schema(validated_value) + tool_name = self._get_tool_name(validated_value) + + # Create tool definition with input_schema matching output schema + tool_def = { + "name": tool_name, + "description": f"Extract structured data conforming to {self._get_schema_description(validated_value)}", + "input_schema": schema, + } + + # Add tools array to request + request["tools"] = [tool_def] + + # Force tool use by setting tool_choice + request["tool_choice"] = {"type": "tool", "name": tool_name} + + return request + + def parse_output( + self, content: str | dict[str, Any], value: object | None + ) -> str | BaseModel: + """Parse tool_use blocks from response to BaseModel instance. + + Extracts structured data from tool_use.input field and converts to BaseModel. + For list[BaseModel], extracts the "items" array from the wrapped object. + + Args: + content: Either tool_use.input dict (from tool_use block) or JSON string. + value: Original output_schema parameter value. + + Returns: + BaseModel instance if value provided, otherwise str unchanged. + """ + if value is None: + return content if isinstance(content, str) else json.dumps(content) + + # If content is already a dict (from tool_use.input), use it directly + if isinstance(content, dict): + parsed_json = content + else: + # Otherwise parse as JSON string + parsed_json = json.loads(content) + + # Check if value is list[BaseModel] and content is wrapped in object + origin = get_origin(value) + if origin is list: + # Handle empty dict case FIRST - convert to empty array before checking for "items" + if isinstance(parsed_json, dict) and not parsed_json: + # Empty dict when expecting list - convert to empty array + parsed_json = [] + elif isinstance(parsed_json, dict) and "items" in parsed_json: + # Extract items array from wrapped format + parsed_json = parsed_json["items"] + # If it's already an array (backward compatibility), use it directly + # parsed_json is now the array, ready for TypeAdapter + elif isinstance(parsed_json, dict) and not parsed_json: + # Empty dict for BaseModel (not list) - this is invalid, raise error + msg = "Empty tool_use input dict cannot be converted to BaseModel" + raise ValueError(msg) + + # Parse to BaseModel instance using TypeAdapter + # TypeAdapter handles both BaseModel and list[BaseModel] + return TypeAdapter(value).validate_json(json.dumps(parsed_json)) + + def _convert_to_anthropic_schema(self, output_schema: Any) -> dict[str, Any]: # noqa: ANN401 + """Convert Pydantic BaseModel or list[BaseModel] to Anthropic JSON Schema format. + + Anthropic requires input_schema to always be an object type. + For list[T], wraps array schema in an object with "items" property. + + Args: + output_schema: Pydantic BaseModel class or list[BaseModel] type. + + Returns: + JSON Schema dictionary compatible with Anthropic (always object type). + """ + origin = get_origin(output_schema) + if origin is list: + # For list[T], wrap array schema in an object (Anthropic requirement) + inner_type = get_args(output_schema)[0] + items_schema = inner_type.model_json_schema() + # Resolve refs in items schema first + items_schema = self._resolve_refs(items_schema) + # Wrap in object with "items" property + json_schema = { + "type": "object", + "properties": { + "items": { + "type": "array", + "items": items_schema, + }, + }, + "required": ["items"], + } + else: + # For BaseModel, use model_json_schema directly + json_schema = output_schema.model_json_schema() + # Resolve $ref references inline (Anthropic may not support $ref) + json_schema = self._resolve_refs(json_schema) + + return json_schema + + def _resolve_refs(self, schema: dict[str, Any]) -> dict[str, Any]: + """Resolve all $ref references and inline definitions. + + Args: + schema: JSON Schema dictionary potentially containing $ref. + + Returns: + Schema with $ref references resolved inline. + """ + defs: dict[str, Any] = {} + + def collect_defs(value: Any) -> None: # noqa: ANN401 + """Recursively collect all $defs dictionaries.""" + if isinstance(value, dict): + if "$defs" in value: + defs.update(value["$defs"]) + for v in value.values(): + collect_defs(v) + elif isinstance(value, list): + for item in value: + collect_defs(item) + + collect_defs(schema) + + def remove_defs(value: Any) -> Any: # noqa: ANN401 + """Recursively remove all $defs keys.""" + if isinstance(value, dict): + result = {k: remove_defs(v) for k, v in value.items() if k != "$defs"} + return result + elif isinstance(value, list): + return [remove_defs(item) for item in value] + return value + + schema = remove_defs(schema) + + def resolve(value: Any) -> Any: # noqa: ANN401 + """Recursively resolve $ref references in schema.""" + if isinstance(value, dict): + if "$ref" in value: + ref_path = value["$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 value.items() if k != "$ref"} + ) + return resolve(resolved) + return {k: resolve(v) for k, v in value.items()} + elif isinstance(value, list): + return [resolve(item) for item in value] + return value + + return resolve(schema) + + def _get_tool_name(self, output_schema: Any) -> str: # noqa: ANN401 + """Derive tool name from model class name. + + Args: + output_schema: Pydantic BaseModel class or list[BaseModel] type. + + Returns: + Tool name (lowercase class name or "extract_data" as fallback). + """ + origin = get_origin(output_schema) + if origin is list: + inner_type = get_args(output_schema)[0] + return inner_type.__name__.lower() or "extract_data" + return output_schema.__name__.lower() or "extract_data" + + def _get_schema_description(self, output_schema: Any) -> str: # noqa: ANN401 + """Get description for tool definition. + + Args: + output_schema: Pydantic BaseModel class or list[BaseModel] type. + + Returns: + Schema description string. + """ + origin = get_origin(output_schema) + if origin is list: + inner_type = get_args(output_schema)[0] + return f"array of {inner_type.__name__}" + return output_schema.__name__ + + +ANTHROPIC_PARAMETER_MAPPERS: list[ParameterMapper] = [ + ThinkingBudgetMapper(), + OutputSchemaMapper(), +] + +__all__ = ["ANTHROPIC_PARAMETER_MAPPERS"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py new file mode 100644 index 00000000..3603353a --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/anthropic/streaming.py @@ -0,0 +1,342 @@ +"""Anthropic streaming for text generation.""" + +import json +from collections.abc import Callable +from typing import Any, Unpack + +from celeste.io import Chunk +from celeste_text_generation.io import ( + TextGenerationChunk, + TextGenerationFinishReason, + TextGenerationOutput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters +from celeste_text_generation.streaming import TextGenerationStream + + +class AnthropicTextGenerationStream(TextGenerationStream): + """Anthropic streaming for text generation.""" + + def __init__( + self, + sse_iterator: Any, # noqa: ANN401 + transform_output: Callable[[object, Any], object], + **parameters: Unpack[TextGenerationParameters], + ) -> None: + """Initialize stream with output transformation support. + + Args: + sse_iterator: Server-Sent Events iterator. + transform_output: Function to transform accumulated content (e.g., JSON โ†’ BaseModel). + **parameters: Parameters passed to stream() for output transformation. + """ + super().__init__(sse_iterator, **parameters) + self._transform_output = transform_output + # Track tool_use blocks for structured output + self._tool_use_blocks: list[dict[str, Any]] = [] + self._current_tool_use: dict[str, Any] | None = None + self._current_tool_use_partial_json: str = "" + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Parse SSE event into Chunk.""" + event_type = event.get("type") + if not event_type: + return None + + # Parse content_block_start for tool_use blocks + if event_type == "content_block_start": + content_block = event.get("content_block", {}) + if content_block.get("type") == "tool_use": + # Initialize new tool_use block + self._current_tool_use = { + "type": "tool_use", + "id": content_block.get("id"), + "name": content_block.get("name"), + "input": {}, + } + return None # No chunk yet, waiting for deltas + + # Parse content_block_delta for tool_use and text + if event_type == "content_block_delta": + delta = event.get("delta", {}) + + # Handle input_json_delta for structured output (Anthropic sends input_json_delta, not tool_use_delta) + if ( + delta.get("type") == "input_json_delta" + and self._current_tool_use is not None + ): + partial_json = delta.get("partial_json") + if partial_json is not None: + # Accumulate partial JSON string fragments + self._current_tool_use_partial_json += partial_json + # Emit chunk with accumulated JSON for UI live rendering (only when output_schema is provided) + output_schema = self._parameters.get("output_schema") + if ( + output_schema is not None + and self._current_tool_use_partial_json + ): + return TextGenerationChunk( + content=self._current_tool_use_partial_json, + finish_reason=None, + usage=None, + ) + return None + + # Handle tool_use_delta for backward compatibility (older API versions) + if ( + delta.get("type") == "tool_use_delta" + and self._current_tool_use is not None + ): + partial_json = delta.get("partial_json") + if partial_json is not None: + # Accumulate partial JSON string fragments + self._current_tool_use_partial_json += partial_json + # Emit chunk with accumulated JSON for UI live rendering (only when output_schema is provided) + output_schema = self._parameters.get("output_schema") + if ( + output_schema is not None + and self._current_tool_use_partial_json + ): + return TextGenerationChunk( + content=self._current_tool_use_partial_json, + finish_reason=None, + usage=None, + ) + return None + + # Handle text_delta for regular text content + if delta.get("type") == "text_delta": + text_delta = delta.get("text") + if text_delta is not None: + return TextGenerationChunk( + content=text_delta, + finish_reason=None, + usage=None, + ) + + # Parse content_block_stop to finalize tool_use blocks + if event_type == "content_block_stop": + if self._current_tool_use is not None: + # Tool use block completed - parse accumulated JSON + tool_id = self._current_tool_use.get("id") + # Check if we already have this tool_use block from message_start + existing_block = None + for block in self._tool_use_blocks: + if block.get("id") == tool_id: + existing_block = block + break + + # Emit final chunk with complete JSON for UI (only when output_schema is provided) + output_schema = self._parameters.get("output_schema") + emit_final_chunk = False + final_json_content = "" + + if self._current_tool_use_partial_json: + final_json_content = self._current_tool_use_partial_json + try: + parsed_input = json.loads(self._current_tool_use_partial_json) + if existing_block: + # Update existing block from message_start + existing_block["input"] = parsed_input + else: + # New block from content_block_start + self._current_tool_use["input"] = parsed_input + self._tool_use_blocks.append(self._current_tool_use) + emit_final_chunk = output_schema is not None + except json.JSONDecodeError: + # If JSON parsing fails, only update if we have existing block + if existing_block: + existing_block["input"] = {} + else: + self._current_tool_use["input"] = {} + self._tool_use_blocks.append(self._current_tool_use) + emit_final_chunk = output_schema is not None + else: + # No partial_json - only add if we don't have this block already + if not existing_block: + self._current_tool_use["input"] = {} + self._tool_use_blocks.append(self._current_tool_use) + + # Emit final chunk with complete JSON before clearing + if emit_final_chunk and final_json_content: + chunk = TextGenerationChunk( + content=final_json_content, + finish_reason=None, + usage=None, + ) + else: + chunk = None + + self._current_tool_use = None + self._current_tool_use_partial_json = "" + + return chunk + return None + + # Parse message_start to capture initial content blocks (includes tool_use) + # Note: In streaming, message_start may contain tool_use blocks with complete input + # or empty input (which will be filled by content_block_delta events) + if event_type == "message_start": + message = event.get("message", {}) + content_blocks = message.get("content", []) + # Extract tool_use blocks from initial message + # If input is already populated, use it; otherwise it will be filled by deltas + for block in content_blocks: + if block.get("type") == "tool_use": + tool_input = block.get("input") + # If input is already complete (not empty), use it directly + # Otherwise, content_block_start/content_block_stop will fill it + if tool_input and isinstance(tool_input, dict) and tool_input: + # Complete tool_use block from message_start + self._tool_use_blocks.append( + { + "type": "tool_use", + "id": block.get("id"), + "name": block.get("name"), + "input": tool_input, + } + ) + return None + + # Parse message delta event for finish reason and usage + if event_type == "message_delta": + delta = event.get("delta", {}) + stop_reason = delta.get("stop_reason") + + finish_reason: TextGenerationFinishReason | None = None + if stop_reason is not None: + finish_reason = TextGenerationFinishReason(reason=stop_reason) + + usage = self._parse_usage_from_event(event) + + return TextGenerationChunk( + content="", + finish_reason=finish_reason, + usage=usage, + ) + + # Parse message stop event (final event) + if event_type == "message_stop": + usage = self._parse_usage_from_event(event) + + return TextGenerationChunk( + content="", + finish_reason=None, + usage=usage, + ) + + # Ignore other event types + return None + + def _parse_usage_from_event( + self, event: dict[str, Any] + ) -> TextGenerationUsage | None: + """Parse usage from SSE event data. + + Args: + event: SSE event dictionary containing usage data. + + Returns: + TextGenerationUsage object if usage data present, None otherwise. + """ + usage_data = event.get("usage") + if not usage_data: + return None + + input_tokens = usage_data.get("input_tokens") + output_tokens = usage_data.get("output_tokens") + total_tokens = None + if input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens + + return TextGenerationUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + billed_tokens=None, + cached_tokens=None, + reasoning_tokens=None, + ) + + def _parse_usage(self, chunks: list[TextGenerationChunk]) -> TextGenerationUsage: + """Parse usage from chunks.""" + if not chunks: + return TextGenerationUsage() + + # Usage typically appears in message_delta or message_stop events + # Search backwards for the most recent usage + for chunk in reversed(chunks): + if chunk.usage: + return chunk.usage + + return TextGenerationUsage() + + def _parse_output( + self, + chunks: list[TextGenerationChunk], + **parameters: Unpack[TextGenerationParameters], + ) -> TextGenerationOutput: + """Assemble chunks into final output with structured output support. + + Checks for tool_use blocks first (structured output), then falls back + to concatenated text chunks. + """ + # Check if output_schema is provided (tool-based structured output) + output_schema = self._parameters.get("output_schema") + + if output_schema is not None and self._tool_use_blocks: + # Extract structured data from tool_use blocks + # Use the first tool_use block's input + # For list[BaseModel], tool_input will be wrapped format {"items": [...]} + # _transform_output will call OutputSchemaMapper.parse_output which handles empty dicts + tool_input = self._tool_use_blocks[0].get("input") + # Check if tool_input is valid (not None and not empty dict for BaseModel) + # Empty dict is OK for list[BaseModel] (converts to []), but invalid for BaseModel + if tool_input is not None: + # For BaseModel (not list), empty dict is invalid - try to find text chunks as fallback + if isinstance(tool_input, dict) and not tool_input: + from typing import get_origin + + origin = get_origin(output_schema) + if origin is not list: + # Empty dict for BaseModel - try text chunks, but if none, raise error + text_content = "".join(chunk.content for chunk in chunks) + if text_content: + content = self._transform_output( + text_content, **self._parameters + ) + else: + msg = "Empty tool_use input dict and no text chunks available for BaseModel" + raise ValueError(msg) + else: + # Empty dict for list[BaseModel] - OK, parse_output will convert to [] + content = self._transform_output(tool_input, **self._parameters) + else: + # Valid tool_input - transform to BaseModel + content = self._transform_output(tool_input, **self._parameters) + else: + # Fallback: concatenate text chunks + text_content = "".join(chunk.content for chunk in chunks) + if text_content: + content = self._transform_output(text_content, **self._parameters) + else: + msg = "No tool_use input and no text chunks available" + raise ValueError(msg) + else: + # No tool_use blocks or no output_schema: concatenate text chunks + content = "".join(chunk.content for chunk in chunks) + # Apply parameter transformations (e.g., JSON โ†’ BaseModel if output_schema provided) + content = self._transform_output(content, **self._parameters) + + usage = self._parse_usage(chunks) + finish_reason = chunks[-1].finish_reason if chunks else None + + return TextGenerationOutput( + content=content, + usage=usage, + metadata={"finish_reason": finish_reason}, + ) + + +__all__ = ["AnthropicTextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/__init__.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/__init__.py new file mode 100644 index 00000000..1711928e --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/__init__.py @@ -0,0 +1,7 @@ +"""Cohere provider.""" + +from .client import CohereTextGenerationClient +from .models import MODELS +from .streaming import CohereTextGenerationStream + +__all__ = ["MODELS", "CohereTextGenerationClient", "CohereTextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py new file mode 100644 index 00000000..6c5adc2a --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/client.py @@ -0,0 +1,133 @@ +"""Cohere client implementation.""" + +from collections.abc import AsyncIterator +from typing import Any, Unpack + +import httpx +from pydantic import BaseModel + +from celeste.parameters import ParameterMapper +from celeste_text_generation.client import TextGenerationClient +from celeste_text_generation.io import ( + TextGenerationFinishReason, + TextGenerationInput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters + +from . import config +from .parameters import COHERE_PARAMETER_MAPPERS +from .streaming import CohereTextGenerationStream + + +class CohereTextGenerationClient(TextGenerationClient): + """Cohere client.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + return COHERE_PARAMETER_MAPPERS + + def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]: + """Initialize request from Cohere v2 Chat API messages array format.""" + messages = [ + { + "role": "user", + "content": inputs.prompt, + } + ] + + return {"messages": messages} + + def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: + """Parse usage from response.""" + meta = response_data.get("meta", {}) + + billed_units = meta.get("billed_units", {}) + tokens = meta.get("tokens", {}) + + input_tokens = billed_units.get("input_tokens") + output_tokens = billed_units.get("output_tokens") + + if input_tokens is not None or output_tokens is not None: + return TextGenerationUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=tokens.get("total_tokens") if tokens else None, + cached_tokens=meta.get("cached_tokens"), + ) + + return TextGenerationUsage() + + def _parse_content( + self, + response_data: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> str | BaseModel: + """Parse content from response.""" + message = response_data.get("message", {}) + content_array = message.get("content", []) + if not content_array: + msg = "No content in response message" + raise ValueError(msg) + + first_content = content_array[0] + text = first_content.get("text") or "" + + return self._transform_output(text, **parameters) + + def _parse_finish_reason( + self, response_data: dict[str, Any] + ) -> TextGenerationFinishReason | None: + """Parse finish reason from response.""" + finish_reason_str = response_data.get("finish_reason") + return ( + TextGenerationFinishReason(reason=finish_reason_str) + if finish_reason_str + else None + ) + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> httpx.Response: + """Make HTTP request(s) and return response object.""" + request_body["model"] = self.model.id + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + "Content-Type": "application/json", + } + + return await self.http_client.post( + f"{config.BASE_URL}{config.ENDPOINT}", + headers=headers, + json_body=request_body, + ) + + def _stream_class(self) -> type[CohereTextGenerationStream]: + """Return the Stream class for this client.""" + return CohereTextGenerationStream + + def _make_stream_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> AsyncIterator[dict[str, Any]]: + """Make HTTP streaming request and return async iterator of events.""" + request_body["model"] = self.model.id + request_body["stream"] = True + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + "Content-Type": "application/json", + } + + return self.http_client.stream_post( + f"{config.BASE_URL}{config.STREAM_ENDPOINT}", + headers=headers, + json_body=request_body, + ) + + +__all__ = ["CohereTextGenerationClient"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/config.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/config.py new file mode 100644 index 00000000..ea0000f0 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/config.py @@ -0,0 +1,11 @@ +"""Cohere provider configuration.""" + +# HTTP Configuration +BASE_URL = "https://api.cohere.com" +ENDPOINT = "/v2/chat" +STREAM_ENDPOINT = ENDPOINT + +# Authentication +AUTH_HEADER_NAME = "Authorization" +AUTH_HEADER_PREFIX = "Bearer " +CLIENT_NAME_HEADER = "X-Client-Name" diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/models.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/models.py new file mode 100644 index 00000000..f3f9cba9 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/models.py @@ -0,0 +1,58 @@ +"""Cohere models.""" + +from celeste import Model, Provider +from celeste.constraints import Range, Schema +from celeste.core import Parameter +from celeste_text_generation.parameters import TextGenerationParameter + +MODELS: list[Model] = [ + Model( + id="command-a-03-2025", + provider=Provider.COHERE, + display_name="Command A 03-2025", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=1.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=4096, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + # thinking_budget: Not confirmed for this model, omit constraint + }, + ), + Model( + id="command-r-plus", + provider=Provider.COHERE, + display_name="Command R+", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=1.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=4096, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + # command-r-plus supports reasoning (optimized for complex reasoning) + TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=31000, step=1), + }, + ), + Model( + id="command-r", + provider=Provider.COHERE, + display_name="Command R", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=1.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=4096, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + # thinking_budget: Support unclear, omit constraint for now + }, + ), + Model( + id="command-r7b-12-2024", + provider=Provider.COHERE, + display_name="Command R7B", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=1.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=4096, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + # thinking_budget: Support unclear, omit constraint for now + }, + ), +] diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/parameters.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/parameters.py new file mode 100644 index 00000000..e2221526 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/parameters.py @@ -0,0 +1,238 @@ +"""Cohere parameter mappers.""" + +import json +from enum import StrEnum +from typing import Any, get_args, get_origin + +from pydantic import BaseModel, TypeAdapter + +from celeste.core import Parameter +from celeste.models import Model +from celeste.parameters import ParameterMapper +from celeste_text_generation.parameters import TextGenerationParameter + + +class TemperatureMapper(ParameterMapper): + """Map temperature parameter to Cohere temperature field.""" + + name: StrEnum = Parameter.TEMPERATURE + + 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 MaxTokensMapper(ParameterMapper): + """Map max_tokens parameter to Cohere max_tokens field.""" + + name: StrEnum = Parameter.MAX_TOKENS + + 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 ThinkingBudgetMapper(ParameterMapper): + """Map thinking_budget parameter to Cohere thinking parameter.""" + + name: StrEnum = TextGenerationParameter.THINKING_BUDGET + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinking_budget into provider request. + + Maps unified thinking_budget to Cohere thinking parameter: + - -1: Unlimited thinking ({"type": "enabled"}) + - 0: Disable thinking ({"type": "disabled"}) + - > 0: Set token budget ({"token_budget": value}) + """ + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + # Map to Cohere thinking parameter format + if validated_value == -1: + # Unlimited thinking (default) + request["thinking"] = {"type": "enabled"} + elif validated_value == 0: + # Disable thinking + request["thinking"] = {"type": "disabled"} + else: + # Set token budget + request["thinking"] = {"token_budget": validated_value} + + return request + + +class OutputSchemaMapper(ParameterMapper): + """Map output_schema parameter to Cohere response_format.""" + + name: StrEnum = TextGenerationParameter.OUTPUT_SCHEMA + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform output_schema into provider request. + + Converts Pydantic BaseModel or list[BaseModel] to Cohere JSON Schema format. + Sets request["response_format"] = {"type": "json_object", "schema": {...}}. + """ + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + schema = self._convert_to_cohere_schema(validated_value) + request["response_format"] = { + "type": "json_object", + "schema": schema, + } + + return request + + def parse_output(self, content: str, value: object | None) -> str | BaseModel: + """Parse JSON string to BaseModel instance if output_schema provided. + + Args: + content: Raw text content (JSON string when output_schema is set). + value: Original output_schema parameter value. + + Returns: + BaseModel instance if value provided, otherwise str unchanged. + """ + if value is None: + return content + + # Parse JSON string first + parsed_json = json.loads(content) + + # For list[T] models, unwrap the items wrapper (Cohere wraps arrays in {"items": [...]}) + origin = get_origin(value) + if origin is list and isinstance(parsed_json, dict) and "items" in parsed_json: + parsed_json = parsed_json["items"] + + # Parse to BaseModel instance using TypeAdapter + # TypeAdapter handles both Person and list[Person] + return TypeAdapter(value).validate_json(json.dumps(parsed_json)) + + def _convert_to_cohere_schema(self, output_schema: Any) -> dict[str, Any]: # noqa: ANN401 + """Convert Pydantic BaseModel or list[BaseModel] to Cohere JSON Schema format. + + Cohere requires flattened schemas without $ref/$defs. + For list[T] models, wraps array schema in object with items property. + """ + origin = get_origin(output_schema) + if origin is list: + # For list[T], wrap array schema in object wrapper + inner_type = get_args(output_schema)[0] + items_schema = inner_type.model_json_schema() + json_schema = { + "type": "object", + "properties": { + "items": { + "type": "array", + "items": items_schema, + } + }, + "required": ["items"], + } + else: + # For BaseModel, use model_json_schema directly + json_schema = output_schema.model_json_schema() + + # Resolve $ref references inline (Cohere requires flattened schemas) + json_schema = self._resolve_refs(json_schema) + + return json_schema + + def _resolve_refs(self, schema: dict[str, Any]) -> dict[str, Any]: + """Resolve all $ref references and inline definitions (Cohere requires flattened schemas). + + This method: + 1. Collects all $defs dictionaries from the schema tree + 2. Removes $defs keys from the schema + 3. Replaces $ref references with inlined definitions + 4. Recursively processes nested objects/arrays + """ + defs: dict[str, Any] = {} + + def collect_defs(value: object) -> None: + """Recursively collect all $defs dictionaries.""" + if isinstance(value, dict): + if "$defs" in value: + defs.update(value["$defs"]) + for v in value.values(): + collect_defs(v) + elif isinstance(value, list): + for item in value: + collect_defs(item) + + collect_defs(schema) + + def remove_defs(value: object) -> object: + """Recursively remove all $defs keys.""" + if isinstance(value, dict): + result = {k: remove_defs(v) for k, v in value.items() if k != "$defs"} + return result + elif isinstance(value, list): + return [remove_defs(item) for item in value] + return value + + schema = remove_defs(schema) + + def resolve(value: object) -> object: + """Recursively resolve $ref references in schema.""" + if isinstance(value, dict): + if "$ref" in value: + ref_path = value["$ref"] + if ref_path.startswith("#/$defs/"): + ref_name = ref_path.split("/")[-1] + if ref_name in defs: + resolved = defs[ref_name].copy() + # Merge any additional properties from the $ref object + resolved.update( + {k: v for k, v in value.items() if k != "$ref"} + ) + return resolve(resolved) + return {k: resolve(v) for k, v in value.items()} + elif isinstance(value, list): + return [resolve(item) for item in value] + return value + + return resolve(schema) + + +COHERE_PARAMETER_MAPPERS: list[ParameterMapper] = [ + TemperatureMapper(), + MaxTokensMapper(), + ThinkingBudgetMapper(), + OutputSchemaMapper(), +] + +__all__ = ["COHERE_PARAMETER_MAPPERS"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py new file mode 100644 index 00000000..f36d36e9 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/cohere/streaming.py @@ -0,0 +1,154 @@ +"""Cohere streaming for text generation.""" + +import logging +from collections.abc import Callable +from typing import Any, Unpack + +from celeste.io import Chunk +from celeste_text_generation.io import ( + TextGenerationChunk, + TextGenerationFinishReason, + TextGenerationOutput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters +from celeste_text_generation.streaming import TextGenerationStream + +logger = logging.getLogger(__name__) + + +class CohereTextGenerationStream(TextGenerationStream): + """Cohere streaming for text generation.""" + + def __init__( + self, + sse_iterator: Any, # noqa: ANN401 + transform_output: Callable[[object, Any], object], + **parameters: Unpack[TextGenerationParameters], + ) -> None: + """Initialize stream with output transformation support. + + Args: + sse_iterator: Server-Sent Events iterator. + transform_output: Function to transform accumulated content (e.g., JSON โ†’ BaseModel). + **parameters: Parameters passed to stream() for output transformation. + """ + super().__init__(sse_iterator, **parameters) + self._transform_output = transform_output + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Parse SSE event into Chunk, extracting text deltas and metadata.""" + event_type = event.get("type") + + if event_type == "content-delta": + delta = event.get("delta", {}) + message = delta.get("message", {}) + content = message.get("content", {}) + text_delta = content.get("text") + + if not text_delta: + return None + + return TextGenerationChunk( + content=text_delta, + finish_reason=None, + usage=None, + ) + + if event_type == "message-end": + delta = event.get("delta", {}) + finish_reason_str = delta.get("finish_reason") + finish_reason = ( + TextGenerationFinishReason(reason=finish_reason_str) + if finish_reason_str + else None + ) + + usage_dict = delta.get("usage", {}) + usage = None + if isinstance(usage_dict, dict): + billed_units = usage_dict.get("billed_units", {}) + tokens = usage_dict.get("tokens", {}) + + input_tokens = billed_units.get("input_tokens") + output_tokens = billed_units.get("output_tokens") + + if input_tokens is not None or output_tokens is not None: + usage = TextGenerationUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=tokens.get("total_tokens") if tokens else None, + cached_tokens=usage_dict.get("cached_tokens"), + ) + + return TextGenerationChunk( + content="", + finish_reason=finish_reason, + usage=usage, + ) + + if event_type == "stream-end": + finish_reason_str = event.get("finish_reason") + finish_reason = ( + TextGenerationFinishReason(reason=finish_reason_str) + if finish_reason_str + else None + ) + + meta = event.get("meta", {}) + usage = None + if isinstance(meta, dict): + billed_units = meta.get("billed_units", {}) + tokens = meta.get("tokens", {}) + + input_tokens = billed_units.get("input_tokens") + output_tokens = billed_units.get("output_tokens") + + if input_tokens is not None or output_tokens is not None: + usage = TextGenerationUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=tokens.get("total_tokens") if tokens else None, + cached_tokens=meta.get("cached_tokens"), + ) + + return TextGenerationChunk( + content="", + finish_reason=finish_reason, + usage=usage, + ) + + return None + + def _parse_usage(self, chunks: list[TextGenerationChunk]) -> TextGenerationUsage: + """Parse usage from chunks, using the last chunk with usage metadata.""" + if not chunks: + return TextGenerationUsage() + + for chunk in reversed(chunks): + if chunk.usage: + return chunk.usage + + return TextGenerationUsage() + + def _parse_output( + self, + chunks: list[TextGenerationChunk], + **parameters: Unpack[TextGenerationParameters], + ) -> TextGenerationOutput: + """Assemble chunks into final output, applying parameter transformations.""" + content_chunks = [chunk for chunk in chunks if chunk.content] + content = "".join(chunk.content for chunk in content_chunks) + content = self._transform_output(content, **self._parameters) + + usage = self._parse_usage(chunks) + finish_reason = chunks[-1].finish_reason if chunks else None + + return TextGenerationOutput( + content=content, + usage=usage, + metadata={"finish_reason": finish_reason}, + ) + + +__all__ = ["CohereTextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/__init__.py b/packages/text-generation/src/celeste_text_generation/providers/google/__init__.py new file mode 100644 index 00000000..1b760d52 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/google/__init__.py @@ -0,0 +1,7 @@ +"""Google provider.""" + +from .client import GoogleTextGenerationClient +from .models import MODELS +from .streaming import GoogleTextGenerationStream + +__all__ = ["MODELS", "GoogleTextGenerationClient", "GoogleTextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/client.py b/packages/text-generation/src/celeste_text_generation/providers/google/client.py new file mode 100644 index 00000000..e0862dab --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/google/client.py @@ -0,0 +1,137 @@ +"""Google client implementation.""" + +from collections.abc import AsyncIterator +from typing import Any, Unpack + +import httpx +from pydantic import BaseModel + +from celeste.parameters import ParameterMapper +from celeste_text_generation.client import TextGenerationClient +from celeste_text_generation.io import ( + TextGenerationFinishReason, + TextGenerationInput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters + +from . import config +from .parameters import GOOGLE_PARAMETER_MAPPERS +from .streaming import GoogleTextGenerationStream + + +class GoogleTextGenerationClient(TextGenerationClient): + """Google client.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + return GOOGLE_PARAMETER_MAPPERS + + def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]: + """Initialize request from Google contents array format.""" + contents = [ + { + "role": "user", + "parts": [{"text": inputs.prompt}], + } + ] + + return {"contents": contents} + + def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: + """Parse usage from response.""" + usage_metadata = response_data.get("usageMetadata", {}) + + return TextGenerationUsage( + input_tokens=usage_metadata.get("promptTokenCount"), + output_tokens=usage_metadata.get("candidatesTokenCount"), + total_tokens=usage_metadata.get("totalTokenCount"), + reasoning_tokens=usage_metadata.get("thoughtsTokenCount"), + billed_tokens=None, + cached_tokens=None, + ) + + def _parse_content( + self, + response_data: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> str | BaseModel: + """Parse content from response.""" + candidates = response_data.get("candidates", []) + if not candidates: + msg = "No candidates in response" + raise ValueError(msg) + + candidate = candidates[0] + content = candidate.get("content", {}) + parts = content.get("parts", []) + + if not parts: + msg = "No parts in candidate content" + raise ValueError(msg) + + text_part = parts[0] + text = text_part.get("text") or "" + + return self._transform_output(text, **parameters) + + def _parse_finish_reason( + self, response_data: dict[str, Any] + ) -> TextGenerationFinishReason | None: + """Parse finish reason from response.""" + candidates = response_data.get("candidates", []) + if not candidates: + return None + + candidate = candidates[0] + finish_reason_str = candidate.get("finishReason") + + if not finish_reason_str: + return None + + return TextGenerationFinishReason(reason=finish_reason_str) + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> httpx.Response: + """Make HTTP request(s) and return response object.""" + endpoint = config.ENDPOINT.format(model_id=self.model.id) + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + "Content-Type": "application/json", + } + + return await self.http_client.post( + f"{config.BASE_URL}{endpoint}", + headers=headers, + json_body=request_body, + ) + + def _stream_class(self) -> type[GoogleTextGenerationStream]: + """Return the Stream class for this client.""" + return GoogleTextGenerationStream + + def _make_stream_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> AsyncIterator[dict[str, Any]]: + """Make HTTP streaming request and return async iterator of events.""" + stream_endpoint = config.STREAM_ENDPOINT.format(model_id=self.model.id) + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + "Content-Type": "application/json", + } + + return self.http_client.stream_post( + f"{config.BASE_URL}{stream_endpoint}", + headers=headers, + json_body=request_body, + ) + + +__all__ = ["GoogleTextGenerationClient"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/config.py b/packages/text-generation/src/celeste_text_generation/providers/google/config.py new file mode 100644 index 00000000..86057e2b --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/google/config.py @@ -0,0 +1,10 @@ +"""Google provider configuration.""" + +# HTTP Configuration +BASE_URL = "https://generativelanguage.googleapis.com" +ENDPOINT = "/v1beta/models/{model_id}:generateContent" +STREAM_ENDPOINT = "/v1beta/models/{model_id}:streamGenerateContent?alt=sse" + +# Authentication +AUTH_HEADER_NAME = "x-goog-api-key" +AUTH_HEADER_PREFIX = "" # Empty string for plain key diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/models.py b/packages/text-generation/src/celeste_text_generation/providers/google/models.py new file mode 100644 index 00000000..89f8c0de --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/google/models.py @@ -0,0 +1,52 @@ +"""Google models.""" + +from celeste import Model, Provider +from celeste.constraints import Range, Schema +from celeste.core import Parameter +from celeste_text_generation.parameters import TextGenerationParameter + +MODELS: list[Model] = [ + Model( + id="gemini-2.5-flash", + provider=Provider.GOOGLE, + display_name="Gemini 2.5 Flash", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0), + Parameter.MAX_TOKENS: Range(min=1, max=8192), + # Flash: allows -1 (dynamic), 0 (disable), or >= 0 + TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=24576), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="gemini-2.5-flash-lite", + provider=Provider.GOOGLE, + display_name="Gemini 2.5 Flash Lite", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0), + Parameter.MAX_TOKENS: Range(min=1, max=8192), + # Flash Lite: allows -1 (dynamic), 0 (disable), or >= 512 + TextGenerationParameter.THINKING_BUDGET: Range( + min=512, max=24576, special_values=[-1, 0] + ), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="gemini-2.5-pro", + provider=Provider.GOOGLE, + display_name="Gemini 2.5 Pro", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0), + Parameter.MAX_TOKENS: Range(min=1, max=8192), + # Pro: allows -1 (dynamic) or >= 128 (cannot use 0) + TextGenerationParameter.THINKING_BUDGET: Range( + min=128, max=32768, special_values=[-1] + ), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), +] diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/parameters.py b/packages/text-generation/src/celeste_text_generation/providers/google/parameters.py new file mode 100644 index 00000000..19e1a491 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/google/parameters.py @@ -0,0 +1,226 @@ +"""Google parameter mappers.""" + +from enum import StrEnum +from typing import Any, get_args, get_origin + +from pydantic import BaseModel, TypeAdapter + +from celeste.core import Parameter +from celeste.models import Model +from celeste.parameters import ParameterMapper +from celeste_text_generation.parameters import TextGenerationParameter + + +class TemperatureMapper(ParameterMapper): + """Map temperature parameter to Google generationConfig.""" + + name: StrEnum = Parameter.TEMPERATURE + + 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.setdefault("generationConfig", {})["temperature"] = validated_value + return request + + +class MaxTokensMapper(ParameterMapper): + """Map max_tokens parameter to Google generationConfig.maxOutputTokens.""" + + name: StrEnum = Parameter.MAX_TOKENS + + 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.setdefault("generationConfig", {})["maxOutputTokens"] = validated_value + return request + + +class ThinkingBudgetMapper(ParameterMapper): + """Map thinking_budget parameter to Google generationConfig.thinkingConfig.thinkingBudget.""" + + name: StrEnum = TextGenerationParameter.THINKING_BUDGET + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinking_budget into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request.setdefault("generationConfig", {}).setdefault("thinkingConfig", {})[ + "thinkingBudget" + ] = validated_value + return request + + +class OutputSchemaMapper(ParameterMapper): + """Map output_schema parameter to Google generationConfig.responseSchema.""" + + name: StrEnum = TextGenerationParameter.OUTPUT_SCHEMA + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform response_model into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + schema = self._convert_to_google_schema(validated_value) + + config = request.setdefault("generationConfig", {}) + config["responseSchema"] = schema + config["responseMimeType"] = "application/json" + + return request + + def parse_output(self, content: str, value: object | None) -> str | BaseModel: + """Parse JSON string to BaseModel instance if output_schema provided. + + Args: + content: Raw text content (JSON string when output_schema is set). + value: Original output_schema parameter value. + + Returns: + BaseModel instance if value provided, otherwise str unchanged. + """ + if value is None: + return content + + return TypeAdapter(value).validate_json(content) + + def _convert_to_google_schema(self, output_schema: Any) -> dict[str, Any]: # noqa: ANN401 + """Convert Pydantic BaseModel or list[BaseModel] to Google OpenAPI 3.0 format.""" + origin = get_origin(output_schema) + if origin is list: + inner_type = get_args(output_schema)[0] + items_schema = inner_type.model_json_schema() + json_schema = {"type": "array", "items": items_schema} + else: + json_schema = output_schema.model_json_schema() + + json_schema = self._resolve_refs(json_schema) + json_schema = self._remove_unsupported_fields(json_schema) + return self._uppercase_types(json_schema) + + def _uppercase_types(self, schema: dict[str, Any]) -> dict[str, Any]: + """Recursively uppercase all 'type' field values in schema.""" + result: dict[str, Any] = {} + + for key, value in schema.items(): + if key == "type" and isinstance(value, str): + result[key] = value.upper() + elif isinstance(value, dict): + result[key] = self._uppercase_types(value) + elif isinstance(value, list): + result[key] = [ + self._uppercase_types(item) if isinstance(item, dict) else item + for item in value + ] + else: + result[key] = value + + return result + + def _resolve_refs(self, schema: dict[str, Any]) -> dict[str, Any]: + """Resolve all $ref references and inline definitions (Google API doesn't support $ref).""" + defs: dict[str, Any] = {} + + def collect_defs(value: object) -> None: + """Recursively collect all $defs dictionaries.""" + if isinstance(value, dict): + if "$defs" in value: + defs.update(value["$defs"]) + for v in value.values(): + collect_defs(v) + elif isinstance(value, list): + for item in value: + collect_defs(item) + + collect_defs(schema) + + def remove_defs(value: object) -> object: + """Recursively remove all $defs keys.""" + if isinstance(value, dict): + result = {k: remove_defs(v) for k, v in value.items() if k != "$defs"} + return result + elif isinstance(value, list): + return [remove_defs(item) for item in value] + return value + + schema = remove_defs(schema) + + def resolve(value: object) -> object: + """Recursively resolve $ref references in schema.""" + if isinstance(value, dict): + if "$ref" in value: + ref_path = value["$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 value.items() if k != "$ref"} + ) + return resolve(resolved) + return {k: resolve(v) for k, v in value.items()} + elif isinstance(value, list): + return [resolve(item) for item in value] + return value + + return resolve(schema) + + def _remove_unsupported_fields(self, schema: dict[str, Any]) -> dict[str, Any]: + """Remove unsupported fields from schema (e.g., 'title' that Google API doesn't accept).""" + result: dict[str, Any] = {} + + for key, value in schema.items(): + if key == "title": + continue + + if isinstance(value, dict): + result[key] = self._remove_unsupported_fields(value) + elif isinstance(value, list): + result[key] = [ + self._remove_unsupported_fields(item) + if isinstance(item, dict) + else item + for item in value + ] + else: + result[key] = value + + return result + + +GOOGLE_PARAMETER_MAPPERS: list[ParameterMapper] = [ + TemperatureMapper(), + MaxTokensMapper(), + ThinkingBudgetMapper(), + OutputSchemaMapper(), +] + +__all__ = ["GOOGLE_PARAMETER_MAPPERS"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/google/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/google/streaming.py new file mode 100644 index 00000000..aa3f6587 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/google/streaming.py @@ -0,0 +1,164 @@ +"""Google streaming for text generation.""" + +from collections.abc import Callable +from typing import Any, Unpack + +from celeste.io import Chunk +from celeste_text_generation.io import ( + TextGenerationChunk, + TextGenerationFinishReason, + TextGenerationOutput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters +from celeste_text_generation.streaming import TextGenerationStream + + +class GoogleTextGenerationStream(TextGenerationStream): + """Google streaming for text generation.""" + + def __init__( + self, + sse_iterator: Any, # noqa: ANN401 + transform_output: Callable[[object, Any], object], + **parameters: Unpack[TextGenerationParameters], + ) -> None: + """Initialize stream with output transformation support. + + Args: + sse_iterator: Server-Sent Events iterator. + transform_output: Function to transform accumulated content (e.g., JSON โ†’ BaseModel). + **parameters: Parameters passed to stream() for output transformation. + """ + super().__init__(sse_iterator, **parameters) + self._transform_output = transform_output + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Parse SSE event into Chunk. + + Extract text delta from candidates[0].content.parts[0].text. + Extract finishReason and usageMetadata if present. + Return None if no text delta (filter lifecycle events). + """ + # Extract candidates array + candidates = event.get("candidates", []) + if not candidates: + return None + + candidate = candidates[0] + content = candidate.get("content", {}) + parts = content.get("parts", []) + + # Extract text delta + text_delta = None + if parts: + text_part = parts[0] + text_delta = text_part.get("text") + + # If no text delta, this is likely a lifecycle event - filter it + if not text_delta: + return None + + # Extract finish reason if present + finish_reason: TextGenerationFinishReason | None = None + finish_reason_str = candidate.get("finishReason") + if finish_reason_str: + finish_reason = TextGenerationFinishReason(reason=finish_reason_str) + + # Extract usage metadata if present (store in chunk metadata for later) + usage: TextGenerationUsage | None = None + usage_metadata = event.get("usageMetadata") + if usage_metadata: + usage = TextGenerationUsage( + input_tokens=usage_metadata.get("promptTokenCount"), + output_tokens=usage_metadata.get("candidatesTokenCount"), + total_tokens=usage_metadata.get("totalTokenCount"), + reasoning_tokens=usage_metadata.get("thoughtsTokenCount"), + ) + + return TextGenerationChunk( + content=text_delta, + finish_reason=finish_reason, + usage=usage, + ) + + def _parse_usage(self, chunks: list[TextGenerationChunk]) -> TextGenerationUsage: + """Parse usage from chunks. + + Google provides usageMetadata in the final chunk(s). + Accumulate usage from all chunks, prioritizing later chunks for totals. + """ + if not chunks: + return TextGenerationUsage() + + # Usage metadata is typically in the final chunk + final_chunk = chunks[-1] + if final_chunk.usage: + # Return final chunk usage directly (contains complete usageMetadata) + return final_chunk.usage + + # Fallback: check metadata if stored there + usage_metadata = final_chunk.metadata.get("usageMetadata") + if usage_metadata: + return TextGenerationUsage( + input_tokens=usage_metadata.get("promptTokenCount"), + output_tokens=usage_metadata.get("candidatesTokenCount"), + total_tokens=usage_metadata.get("totalTokenCount"), + reasoning_tokens=usage_metadata.get("thoughtsTokenCount"), + ) + + # Accumulate usage from chunks that have usage metadata + total_input_tokens = 0 + total_output_tokens = 0 + total_reasoning_tokens = 0 + total_total_tokens = 0 + + for chunk in chunks: + if chunk.usage: + if chunk.usage.input_tokens is not None: + total_input_tokens = chunk.usage.input_tokens # Use latest value + if chunk.usage.output_tokens is not None: + total_output_tokens = chunk.usage.output_tokens # Use latest value + if chunk.usage.reasoning_tokens is not None: + total_reasoning_tokens = ( + chunk.usage.reasoning_tokens + ) # Use latest value + if chunk.usage.total_tokens is not None: + total_total_tokens = chunk.usage.total_tokens # Use latest value + + return TextGenerationUsage( + input_tokens=total_input_tokens if total_input_tokens > 0 else None, + output_tokens=total_output_tokens if total_output_tokens > 0 else None, + total_tokens=total_total_tokens if total_total_tokens > 0 else None, + reasoning_tokens=total_reasoning_tokens + if total_reasoning_tokens > 0 + else None, + ) + + def _parse_output( + self, + chunks: list[TextGenerationChunk], + **parameters: Unpack[TextGenerationParameters], + ) -> TextGenerationOutput: + """Assemble chunks into final output with structured output support. + + Concatenates text chunks, then applies parameter transformations + (e.g., JSON โ†’ BaseModel if output_schema provided). + """ + # Concatenate text chunks + content = "".join(chunk.content for chunk in chunks) + + # Apply parameter transformations (e.g., JSON โ†’ BaseModel if output_schema provided) + content = self._transform_output(content, **self._parameters) + + usage = self._parse_usage(chunks) + finish_reason = chunks[-1].finish_reason if chunks else None + + return TextGenerationOutput( + content=content, + usage=usage, + metadata={"finish_reason": finish_reason}, + ) + + +__all__ = ["GoogleTextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/__init__.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/__init__.py new file mode 100644 index 00000000..d77eaf55 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/__init__.py @@ -0,0 +1,7 @@ +"""Mistral provider.""" + +from .client import MistralTextGenerationClient +from .models import MODELS +from .streaming import MistralTextGenerationStream + +__all__ = ["MODELS", "MistralTextGenerationClient", "MistralTextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py new file mode 100644 index 00000000..689e295d --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/client.py @@ -0,0 +1,128 @@ +"""Mistral client implementation.""" + +from collections.abc import AsyncIterator +from typing import Any, Unpack + +import httpx +from pydantic import BaseModel + +from celeste.parameters import ParameterMapper +from celeste_text_generation.client import TextGenerationClient +from celeste_text_generation.io import ( + TextGenerationFinishReason, + TextGenerationInput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters + +from . import config +from .parameters import MISTRAL_PARAMETER_MAPPERS +from .streaming import MistralTextGenerationStream + + +class MistralTextGenerationClient(TextGenerationClient): + """Mistral client.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + return MISTRAL_PARAMETER_MAPPERS + + def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]: + """Initialize request from Mistral messages array format.""" + messages = [ + { + "role": "user", + "content": inputs.prompt, + } + ] + + return {"messages": messages} + + def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: + """Parse usage from response.""" + usage_dict = response_data.get("usage", {}) + + return TextGenerationUsage( + input_tokens=usage_dict.get("prompt_tokens"), + output_tokens=usage_dict.get("completion_tokens"), + total_tokens=usage_dict.get("total_tokens"), + ) + + def _parse_content( + self, + response_data: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> str | BaseModel: + """Parse content from response.""" + choices = response_data.get("choices", []) + if not choices: + msg = "No choices in response" + raise ValueError(msg) + + first_choice = choices[0] + message = first_choice.get("message", {}) + content = message.get("content") or "" + + return self._transform_output(content, **parameters) + + def _parse_finish_reason( + self, response_data: dict[str, Any] + ) -> TextGenerationFinishReason | None: + """Parse finish reason from response.""" + choices = response_data.get("choices", []) + if not choices: + return None + + first_choice = choices[0] + finish_reason_str = first_choice.get("finish_reason") + return ( + TextGenerationFinishReason(reason=finish_reason_str) + if finish_reason_str + else None + ) + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> httpx.Response: + """Make HTTP request(s) and return response object.""" + request_body["model"] = self.model.id + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + "Content-Type": "application/json", + } + + return await self.http_client.post( + f"{config.BASE_URL}{config.ENDPOINT}", + headers=headers, + json_body=request_body, + ) + + def _stream_class(self) -> type[MistralTextGenerationStream]: + """Return the Stream class for this client.""" + return MistralTextGenerationStream + + def _make_stream_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> AsyncIterator[dict[str, Any]]: + """Make HTTP streaming request and return async iterator of events.""" + request_body["model"] = self.model.id + request_body["stream"] = True + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + "Content-Type": "application/json", + } + + return self.http_client.stream_post( + f"{config.BASE_URL}{config.STREAM_ENDPOINT}", + headers=headers, + json_body=request_body, + ) + + +__all__ = ["MistralTextGenerationClient"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/config.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/config.py new file mode 100644 index 00000000..aa9d24d5 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/config.py @@ -0,0 +1,10 @@ +"""Mistral provider configuration.""" + +# HTTP Configuration +BASE_URL = "https://api.mistral.ai" +ENDPOINT = "/v1/chat/completions" +STREAM_ENDPOINT = ENDPOINT # Same endpoint + +# Authentication +AUTH_HEADER_NAME = "Authorization" +AUTH_HEADER_PREFIX = "Bearer " diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/models.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/models.py new file mode 100644 index 00000000..10fffc73 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/models.py @@ -0,0 +1,165 @@ +"""Mistral models.""" + +from celeste import Model, Provider +from celeste.constraints import Range, Schema +from celeste.core import Parameter +from celeste_text_generation.parameters import TextGenerationParameter + +MODELS: list[Model] = [ + Model( + id="mistral-large-latest", + provider=Provider.MISTRAL, + display_name="Mistral Large", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="mistral-medium-latest", + provider=Provider.MISTRAL, + display_name="Mistral Medium", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="mistral-small-latest", + provider=Provider.MISTRAL, + display_name="Mistral Small", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="mistral-tiny", + provider=Provider.MISTRAL, + display_name="Mistral Tiny", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="open-mistral-7b", + provider=Provider.MISTRAL, + display_name="Open Mistral 7B", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="open-mixtral-8x7b", + provider=Provider.MISTRAL, + display_name="Open Mixtral 8x7B", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="open-mixtral-8x22b", + provider=Provider.MISTRAL, + display_name="Open Mixtral 8x22B", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="codestral-latest", + provider=Provider.MISTRAL, + display_name="Codestral", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="devstral-medium-latest", + provider=Provider.MISTRAL, + display_name="Devstral Medium", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="pixtral-12b-2409", + provider=Provider.MISTRAL, + display_name="Pixtral 12B", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="ministral-8b-latest", + provider=Provider.MISTRAL, + display_name="Ministral 8B", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="voxtral-mini-2507", + provider=Provider.MISTRAL, + display_name="Voxtral Mini", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="magistral-small-latest", + provider=Provider.MISTRAL, + display_name="Magistral Small", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="magistral-medium-latest", + provider=Provider.MISTRAL, + display_name="Magistral Medium", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0, step=0.01), + Parameter.MAX_TOKENS: Range(min=1, max=32768, step=1), + TextGenerationParameter.THINKING_BUDGET: Range(min=-1, max=32768, step=1), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), +] diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/parameters.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/parameters.py new file mode 100644 index 00000000..725f4fd4 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/parameters.py @@ -0,0 +1,218 @@ +"""Mistral parameter mappers.""" + +from enum import StrEnum +from typing import Any, get_args, get_origin + +from pydantic import BaseModel, TypeAdapter + +from celeste.core import Parameter +from celeste.models import Model +from celeste.parameters import ParameterMapper +from celeste_text_generation.parameters import TextGenerationParameter + + +class TemperatureMapper(ParameterMapper): + """Map temperature parameter to Mistral temperature field.""" + + name = Parameter.TEMPERATURE + + 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 MaxTokensMapper(ParameterMapper): + """Map max_tokens parameter to Mistral max_tokens field.""" + + name = Parameter.MAX_TOKENS + + 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 ThinkingBudgetMapper(ParameterMapper): + """Map thinking_budget parameter to Mistral prompt_mode. + + Maps unified thinking_budget to Mistral's prompt_mode parameter for reasoning models: + - -1: Enable reasoning (prompt_mode: "reasoning") + - 0: Disable reasoning (prompt_mode: null) + - >0: Enable reasoning (prompt_mode: "reasoning") - Note: Mistral doesn't support budget control + """ + + name = TextGenerationParameter.THINKING_BUDGET + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinking_budget into provider request. + + Only applies to magistral reasoning models. For other models, silently ignores. + """ + if not model.id.startswith("magistral"): + return request + + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + if validated_value == -1: + request["prompt_mode"] = "reasoning" + elif validated_value == 0: + request["prompt_mode"] = None + else: + request["prompt_mode"] = "reasoning" + + return request + + +class OutputSchemaMapper(ParameterMapper): + """Map output_schema parameter to Mistral response_format.""" + + name: StrEnum = TextGenerationParameter.OUTPUT_SCHEMA + + 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 + + schema = self._convert_to_mistral_schema(validated_value) + schema_name = self._get_schema_name(validated_value) + + request["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": schema_name, + "description": validated_value.__doc__ + if hasattr(validated_value, "__doc__") + else "", + "schema": schema, + "strict": True, + }, + } + + return request + + def parse_output(self, content: str, value: object | None) -> str | BaseModel: + """Parse JSON string to BaseModel instance if output_schema provided. + + Args: + content: Raw text content (JSON string when output_schema is set). + value: Original output_schema parameter value. + + Returns: + BaseModel instance if value provided, otherwise str unchanged. + """ + if value is None: + return content + + return TypeAdapter(value).validate_json(content) + + def _convert_to_mistral_schema(self, output_schema: Any) -> dict[str, Any]: # noqa: ANN401 + """Convert Pydantic BaseModel or list[BaseModel] to Mistral JSON Schema format.""" + origin = get_origin(output_schema) + if origin is list: + inner_type = get_args(output_schema)[0] + items_schema = inner_type.model_json_schema() + json_schema = {"type": "array", "items": items_schema} + else: + json_schema = output_schema.model_json_schema() + + json_schema = self._resolve_refs(json_schema) + return json_schema + + def _resolve_refs(self, schema: dict[str, Any]) -> dict[str, Any]: + """Resolve all $ref references and inline definitions for reliability.""" + defs: dict[str, Any] = {} + + def collect_defs(value: Any) -> None: # noqa: ANN401 + """Recursively collect all $defs dictionaries.""" + if isinstance(value, dict): + if "$defs" in value: + defs.update(value["$defs"]) + for v in value.values(): + collect_defs(v) + elif isinstance(value, list): + for item in value: + collect_defs(item) + + collect_defs(schema) + + def remove_defs(value: Any) -> Any: # noqa: ANN401 + """Recursively remove all $defs keys.""" + if isinstance(value, dict): + result = {k: remove_defs(v) for k, v in value.items() if k != "$defs"} + return result + elif isinstance(value, list): + return [remove_defs(item) for item in value] + return value + + schema = remove_defs(schema) + + def resolve(value: Any) -> Any: # noqa: ANN401 + """Recursively resolve $ref references in schema.""" + if isinstance(value, dict): + if "$ref" in value: + ref_path = value["$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 value.items() if k != "$ref"} + ) + return resolve(resolved) + return {k: resolve(v) for k, v in value.items()} + elif isinstance(value, list): + return [resolve(item) for item in value] + return value + + return resolve(schema) + + def _get_schema_name(self, output_schema: Any) -> str: # noqa: ANN401 + """Derive schema name from model class name.""" + origin = get_origin(output_schema) + if origin is list: + inner_type = get_args(output_schema)[0] + return f"{inner_type.__name__.lower()}_list" + else: + return output_schema.__name__.lower() + + +MISTRAL_PARAMETER_MAPPERS: list[ParameterMapper] = [ + TemperatureMapper(), + MaxTokensMapper(), + ThinkingBudgetMapper(), + OutputSchemaMapper(), +] + +__all__ = ["MISTRAL_PARAMETER_MAPPERS"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py new file mode 100644 index 00000000..981be553 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/mistral/streaming.py @@ -0,0 +1,134 @@ +"""Mistral streaming for text generation.""" + +import logging +from collections.abc import Callable +from typing import Any, Unpack + +from celeste.io import Chunk +from celeste_text_generation.io import ( + TextGenerationChunk, + TextGenerationFinishReason, + TextGenerationOutput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters +from celeste_text_generation.streaming import TextGenerationStream + +logger = logging.getLogger(__name__) + + +class MistralTextGenerationStream(TextGenerationStream): + """Mistral streaming for text generation.""" + + def __init__( + self, + sse_iterator: Any, # noqa: ANN401 + transform_output: Callable[[object, Any], object], + **parameters: Unpack[TextGenerationParameters], + ) -> None: + """Initialize stream with output transformation support. + + Args: + sse_iterator: Server-Sent Events iterator. + transform_output: Function to transform accumulated content (e.g., JSON โ†’ BaseModel). + **parameters: Parameters passed to stream() for output transformation. + """ + super().__init__(sse_iterator, **parameters) + self._transform_output = transform_output + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Parse chunk from SSE event. + + Extract from choices[0].delta.content (content delta events). + Extract finish_reason and usage from final event when finish_reason is not null. + Return None if no text delta (filter lifecycle events). + """ + choices = event.get("choices", []) + if not choices: + return None + + first_choice = choices[0] + if not isinstance(first_choice, dict): + return None + + delta = first_choice.get("delta", {}) + if not isinstance(delta, dict): + return None + + # Extract content delta + content_delta = delta.get("content") + finish_reason_str = first_choice.get("finish_reason") + + # Extract usage from event if present (in final event) + usage = None + usage_dict = event.get("usage") + if isinstance(usage_dict, dict): + usage = TextGenerationUsage( + input_tokens=usage_dict.get("prompt_tokens"), + output_tokens=usage_dict.get("completion_tokens"), + total_tokens=usage_dict.get("total_tokens"), + ) + + # Create finish reason if present + finish_reason = ( + TextGenerationFinishReason(reason=finish_reason_str) + if finish_reason_str + else None + ) + + # If no content delta and no finish reason, filter this event + if not content_delta and not finish_reason: + return None + + return TextGenerationChunk( + content=content_delta or "", # Empty string if no content (final event) + finish_reason=finish_reason, + usage=usage, + ) + + def _parse_usage(self, chunks: list[TextGenerationChunk]) -> TextGenerationUsage: + """Parse usage from chunks. + + Mistral provides usage metadata in the final event (when finish_reason is not null). + Use the last chunk that has usage metadata. + """ + if not chunks: + return TextGenerationUsage() + + # Usage metadata is typically in the final chunk (when finish_reason is set) + for chunk in reversed(chunks): + if chunk.usage: + return chunk.usage + + return TextGenerationUsage() + + def _parse_output( + self, + chunks: list[TextGenerationChunk], + **parameters: Unpack[TextGenerationParameters], + ) -> TextGenerationOutput: + """Assemble chunks into final output with structured output support. + + Concatenates text chunks, then applies parameter transformations + (e.g., JSON โ†’ BaseModel if output_schema provided). + """ + # Filter out empty chunks (from final events) + content_chunks = [chunk for chunk in chunks if chunk.content] + + # Concatenate text chunks + content = "".join(chunk.content for chunk in content_chunks) + + # Apply parameter transformations (e.g., JSON โ†’ BaseModel if output_schema provided) + content = self._transform_output(content, **self._parameters) + + usage = self._parse_usage(chunks) + finish_reason = chunks[-1].finish_reason if chunks else None + + return TextGenerationOutput( + content=content, + usage=usage, + metadata={"finish_reason": finish_reason}, + ) + + +__all__ = ["MistralTextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/__init__.py b/packages/text-generation/src/celeste_text_generation/providers/openai/__init__.py new file mode 100644 index 00000000..be7016aa --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/__init__.py @@ -0,0 +1,7 @@ +"""OpenAI provider.""" + +from .client import OpenAITextGenerationClient +from .models import MODELS +from .streaming import OpenAITextGenerationStream + +__all__ = ["MODELS", "OpenAITextGenerationClient", "OpenAITextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/client.py b/packages/text-generation/src/celeste_text_generation/providers/openai/client.py new file mode 100644 index 00000000..15c89419 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/client.py @@ -0,0 +1,144 @@ +"""OpenAI client implementation.""" + +from collections.abc import AsyncIterator +from typing import Any, Unpack + +import httpx +from pydantic import BaseModel + +from celeste.parameters import ParameterMapper +from celeste_text_generation.client import TextGenerationClient +from celeste_text_generation.io import ( + TextGenerationFinishReason, + TextGenerationInput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters + +from . import config +from .parameters import OPENAI_PARAMETER_MAPPERS +from .streaming import OpenAITextGenerationStream + + +class OpenAITextGenerationClient(TextGenerationClient): + """OpenAI client.""" + + @classmethod + def parameter_mappers(cls) -> list[ParameterMapper]: + return OPENAI_PARAMETER_MAPPERS + + def _init_request(self, inputs: TextGenerationInput) -> dict[str, Any]: + """Initialize request from OpenAI Responses API format.""" + return {"input": inputs.prompt} + + def _parse_usage(self, response_data: dict[str, Any]) -> TextGenerationUsage: + """Parse usage from response.""" + usage_data = response_data.get("usage", {}) + input_tokens_details = usage_data.get("input_tokens_details", {}) + output_tokens_details = usage_data.get("output_tokens_details", {}) + + return TextGenerationUsage( + input_tokens=usage_data.get("input_tokens"), + output_tokens=usage_data.get("output_tokens"), + total_tokens=usage_data.get("total_tokens"), + cached_tokens=input_tokens_details.get("cached_tokens"), + reasoning_tokens=output_tokens_details.get("reasoning_tokens"), + billed_tokens=None, + ) + + def _parse_content( + self, + response_data: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> str | BaseModel: + """Parse content from response.""" + output_items = response_data.get("output", []) + if not output_items: + msg = "No output items in response" + raise ValueError(msg) + + message_item = None + for item in output_items: + if item.get("type") == "message": + message_item = item + break + + if not message_item: + msg = "No message item found in output array" + raise ValueError(msg) + + content_parts = message_item.get("content", []) + if not content_parts: + msg = "No content parts in message item" + raise ValueError(msg) + + text_content = "" + for content_part in content_parts: + if content_part.get("type") == "output_text": + text_content = content_part.get("text") or "" + break + + return self._transform_output(text_content, **parameters) + + def _parse_finish_reason( + self, response_data: dict[str, Any] + ) -> TextGenerationFinishReason | None: + """Parse finish reason from response.""" + status = response_data.get("status") + if status != "completed": + return None + + output_items = response_data.get("output", []) + for item in output_items: + if item.get("type") == "message": + item_status = item.get("status") + if item_status == "completed": + return TextGenerationFinishReason(reason="completed") + + return None + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> httpx.Response: + """Make HTTP request(s) and return response object.""" + request_body["model"] = self.model.id + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + "Content-Type": "application/json", + } + + return await self.http_client.post( + f"{config.BASE_URL}{config.ENDPOINT}", + headers=headers, + json_body=request_body, + ) + + def _stream_class(self) -> type[OpenAITextGenerationStream]: + """Return the Stream class for this client.""" + return OpenAITextGenerationStream + + def _make_stream_request( + self, + request_body: dict[str, Any], + **parameters: Unpack[TextGenerationParameters], + ) -> AsyncIterator[dict[str, Any]]: + """Make HTTP streaming request and return async iterator of events.""" + request_body["model"] = self.model.id + request_body["stream"] = True + + headers = { + config.AUTH_HEADER_NAME: f"{config.AUTH_HEADER_PREFIX}{self.api_key.get_secret_value()}", + "Content-Type": "application/json", + } + + return self.http_client.stream_post( + f"{config.BASE_URL}{config.STREAM_ENDPOINT}", + headers=headers, + json_body=request_body, + ) + + +__all__ = ["OpenAITextGenerationClient"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/config.py b/packages/text-generation/src/celeste_text_generation/providers/openai/config.py new file mode 100644 index 00000000..f6f2a625 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/config.py @@ -0,0 +1,10 @@ +"""OpenAI provider configuration.""" + +# HTTP Configuration +BASE_URL = "https://api.openai.com" +ENDPOINT = "/v1/responses" +STREAM_ENDPOINT = ENDPOINT + +# Authentication +AUTH_HEADER_NAME = "Authorization" +AUTH_HEADER_PREFIX = "Bearer " diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/models.py b/packages/text-generation/src/celeste_text_generation/providers/openai/models.py new file mode 100644 index 00000000..64a7afa6 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/models.py @@ -0,0 +1,85 @@ +"""OpenAI models.""" + +from celeste import Model, Provider +from celeste.constraints import Choice, Range, Schema +from celeste.core import Parameter +from celeste_text_generation.parameters import TextGenerationParameter + +MODELS: list[Model] = [ + Model( + id="gpt-4o", + provider=Provider.OPENAI, + display_name="GPT-4o", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0), + Parameter.MAX_TOKENS: Range(min=1, max=16384), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="gpt-4o-mini", + provider=Provider.OPENAI, + display_name="GPT-4o Mini", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0), + Parameter.MAX_TOKENS: Range(min=1, max=16384), + }, + ), + Model( + id="gpt-4-turbo", + provider=Provider.OPENAI, + display_name="GPT-4 Turbo", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0), + Parameter.MAX_TOKENS: Range(min=1, max=4096), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="gpt-4", + provider=Provider.OPENAI, + display_name="GPT-4", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0), + Parameter.MAX_TOKENS: Range(min=1, max=8192), + }, + ), + Model( + id="gpt-3.5-turbo", + provider=Provider.OPENAI, + display_name="GPT-3.5 Turbo", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0), + Parameter.MAX_TOKENS: Range(min=1, max=4096), + }, + ), + Model( + id="gpt-5", + provider=Provider.OPENAI, + display_name="GPT-5", + streaming=True, + parameter_constraints={ + Parameter.MAX_TOKENS: Range(min=1, max=128000), + TextGenerationParameter.THINKING_BUDGET: Choice( + options=["minimal", "low", "medium", "high"] + ), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), + Model( + id="gpt-4.1", + provider=Provider.OPENAI, + display_name="GPT-4.1", + streaming=True, + parameter_constraints={ + Parameter.TEMPERATURE: Range(min=0.0, max=2.0), + Parameter.MAX_TOKENS: Range(min=1, max=32768), + TextGenerationParameter.OUTPUT_SCHEMA: Schema(), + }, + ), +] diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/parameters.py b/packages/text-generation/src/celeste_text_generation/providers/openai/parameters.py new file mode 100644 index 00000000..bb750b0f --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/parameters.py @@ -0,0 +1,222 @@ +"""OpenAI parameter mappers.""" + +import json +from enum import StrEnum +from typing import Any, get_args, get_origin + +from pydantic import BaseModel, TypeAdapter + +from celeste.core import Parameter +from celeste.models import Model +from celeste.parameters import ParameterMapper +from celeste_text_generation.parameters import TextGenerationParameter + + +class OutputSchemaMapper(ParameterMapper): + """Map output_schema parameter to OpenAI text.format.""" + + name: StrEnum = TextGenerationParameter.OUTPUT_SCHEMA + + 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 + + schema = self._convert_to_openai_schema(validated_value) + schema_name = self._get_schema_name(validated_value) + + request.setdefault("text", {})["format"] = { + "type": "json_schema", + "name": schema_name, + "strict": True, + "schema": schema, + } + + return request + + def parse_output(self, content: str, value: object | None) -> str | BaseModel: + """Parse JSON string to BaseModel instance if output_schema provided.""" + if value is None: + return content + + parsed_json = json.loads(content) + origin = get_origin(value) + if origin is list and isinstance(parsed_json, dict) and "items" in parsed_json: + parsed_json = parsed_json["items"] + + return TypeAdapter(value).validate_json(json.dumps(parsed_json)) + + def _convert_to_openai_schema(self, output_schema: Any) -> dict[str, Any]: # noqa: ANN401 + """Convert Pydantic BaseModel or list[BaseModel] to OpenAI JSON Schema format.""" + origin = get_origin(output_schema) + if origin is list: + inner_type = get_args(output_schema)[0] + items_schema = inner_type.model_json_schema() + json_schema = { + "type": "object", + "properties": { + "items": { + "type": "array", + "items": items_schema, + } + }, + "required": ["items"], + } + else: + json_schema = output_schema.model_json_schema() + + json_schema = self._transform_schema_for_openai(json_schema) + return json_schema + + def _transform_schema_for_openai( + self, schema: dict[str, Any], defs: dict[str, Any] | None = None + ) -> dict[str, Any]: + """Recursively transform schema for OpenAI Responses API.""" + if not isinstance(schema, dict): + return schema + + if defs is None: + defs = self._collect_all_defs(schema) + + if "$ref" in schema: + ref_path = schema["$ref"] + if ref_path.startswith("#/$defs/"): + def_name = ref_path.split("/")[-1] + if def_name in defs: + expanded = defs[def_name].copy() + expanded.pop("description", None) + return self._transform_schema_for_openai(expanded, defs) + return schema + + result: dict[str, Any] = {} + for key, value in schema.items(): + if key == "$defs": + continue + elif isinstance(value, dict): + result[key] = self._transform_schema_for_openai(value, defs) + elif isinstance(value, list): + result[key] = [ + self._transform_schema_for_openai(item, defs) + if isinstance(item, dict) + else item + for item in value + ] + else: + result[key] = value + + if result.get("type") == "object": + result["additionalProperties"] = False + + return result + + def _collect_all_defs(self, schema: Any) -> dict[str, Any]: # noqa: ANN401 + """Recursively collect all $defs dictionaries from schema tree.""" + defs: dict[str, Any] = {} + + def collect(value: Any) -> None: # noqa: ANN401 + if isinstance(value, dict): + if "$defs" in value: + defs.update(value["$defs"]) + for v in value.values(): + collect(v) + elif isinstance(value, list): + for item in value: + collect(item) + + collect(schema) + return defs + + def _get_schema_name(self, output_schema: Any) -> str: # noqa: ANN401 + """Derive schema name from model class name.""" + origin = get_origin(output_schema) + if origin is list: + inner_type = get_args(output_schema)[0] + class_name = inner_type.__name__ + return f"{class_name.lower()}_list" + else: + return output_schema.__name__.lower() + + +class TemperatureMapper(ParameterMapper): + """Map temperature parameter to OpenAI temperature field.""" + + name: StrEnum = Parameter.TEMPERATURE + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform temperature into provider request.""" + # Skip temperature for gpt-5 (uses reasoning.effort instead) + if model.id == "gpt-5": + return request + + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["temperature"] = validated_value + return request + + +class MaxTokensMapper(ParameterMapper): + """Map max_tokens parameter to OpenAI max_output_tokens field.""" + + name: StrEnum = Parameter.MAX_TOKENS + + 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_output_tokens"] = validated_value + return request + + +class ThinkingBudgetMapper(ParameterMapper): + """Map thinking_budget parameter to OpenAI reasoning.effort field.""" + + name: StrEnum = TextGenerationParameter.THINKING_BUDGET + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform thinking_budget into provider request.""" + # Only supported for GPT-5 models + if model.id != "gpt-5": + return request + + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + # Map to reasoning.effort nested structure + request.setdefault("reasoning", {})["effort"] = validated_value + return request + + +OPENAI_PARAMETER_MAPPERS: list[ParameterMapper] = [ + TemperatureMapper(), + MaxTokensMapper(), + ThinkingBudgetMapper(), + OutputSchemaMapper(), +] + +__all__ = ["OPENAI_PARAMETER_MAPPERS"] diff --git a/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py b/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py new file mode 100644 index 00000000..bde9c630 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/providers/openai/streaming.py @@ -0,0 +1,108 @@ +"""OpenAI streaming for text generation.""" + +from collections.abc import Callable +from typing import Any, Unpack + +from celeste.io import Chunk +from celeste_text_generation.io import ( + TextGenerationChunk, + TextGenerationFinishReason, + TextGenerationOutput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters +from celeste_text_generation.streaming import TextGenerationStream + + +class OpenAITextGenerationStream(TextGenerationStream): + """OpenAI streaming for text generation.""" + + def __init__( + self, + sse_iterator: Any, # noqa: ANN401 + transform_output: Callable[[object, Any], object], + **parameters: Unpack[TextGenerationParameters], + ) -> None: + """Initialize stream.""" + super().__init__(sse_iterator, **parameters) + self._transform_output = transform_output + + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Parse SSE event into Chunk.""" + event_type = event.get("type") + if not event_type: + return None + + if event_type == "response.output_text.delta": + delta = event.get("delta") + if delta is None: + return None + return TextGenerationChunk( + content=delta, + finish_reason=None, + usage=None, + ) + + if event_type == "response.output_text.done": + return None + + if event_type == "response.completed": + response_data = event.get("response", {}) + usage_data = response_data.get("usage") + + usage: TextGenerationUsage | None = None + if usage_data: + input_tokens_details = usage_data.get("input_tokens_details", {}) + output_tokens_details = usage_data.get("output_tokens_details", {}) + usage = TextGenerationUsage( + input_tokens=usage_data.get("input_tokens"), + output_tokens=usage_data.get("output_tokens"), + total_tokens=usage_data.get("total_tokens"), + cached_tokens=input_tokens_details.get("cached_tokens"), + reasoning_tokens=output_tokens_details.get("reasoning_tokens"), + billed_tokens=None, + ) + + finish_reason: TextGenerationFinishReason | None = None + status = response_data.get("status") + if status == "completed": + finish_reason = TextGenerationFinishReason(reason="completed") + + return TextGenerationChunk( + content="", + finish_reason=finish_reason, + usage=usage, + ) + + return None + + def _parse_usage(self, chunks: list[TextGenerationChunk]) -> TextGenerationUsage: + """Parse usage from chunks.""" + if not chunks: + return TextGenerationUsage() + + for chunk in reversed(chunks): + if chunk.usage: + return chunk.usage + + return TextGenerationUsage() + + def _parse_output( + self, + chunks: list[TextGenerationChunk], + **parameters: Unpack[TextGenerationParameters], + ) -> TextGenerationOutput: + """Assemble chunks into final output.""" + content = "".join(chunk.content for chunk in chunks) + content = self._transform_output(content, **self._parameters) + usage = self._parse_usage(chunks) + finish_reason = chunks[-1].finish_reason if chunks else None + + return TextGenerationOutput( + content=content, + usage=usage, + metadata={"finish_reason": finish_reason}, + ) + + +__all__ = ["OpenAITextGenerationStream"] diff --git a/packages/text-generation/src/celeste_text_generation/py.typed b/packages/text-generation/src/celeste_text_generation/py.typed new file mode 100644 index 00000000..321d0ae1 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/py.typed @@ -0,0 +1 @@ +# Marker file for PEP 561 - this package supports type checking diff --git a/packages/text-generation/src/celeste_text_generation/streaming.py b/packages/text-generation/src/celeste_text_generation/streaming.py new file mode 100644 index 00000000..daeff747 --- /dev/null +++ b/packages/text-generation/src/celeste_text_generation/streaming.py @@ -0,0 +1,45 @@ +"""Streaming for text generation.""" + +from abc import abstractmethod +from typing import Any, Unpack + +from celeste.io import Chunk +from celeste.streaming import Stream +from celeste_text_generation.io import ( + TextGenerationChunk, + TextGenerationOutput, + TextGenerationUsage, +) +from celeste_text_generation.parameters import TextGenerationParameters + + +class TextGenerationStream(Stream[TextGenerationOutput, TextGenerationParameters]): + """Streaming for text generation.""" + + @abstractmethod + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: + """Parse SSE event into Chunk (provider-specific).""" + ... + + def _parse_output( + self, + chunks: list[TextGenerationChunk], + **parameters: Unpack[TextGenerationParameters], + ) -> TextGenerationOutput: + """Assemble chunks into final output.""" + content = "".join(chunk.content for chunk in chunks) + usage = self._parse_usage(chunks) + finish_reason = chunks[-1].finish_reason if chunks else None + + return TextGenerationOutput( + content=content, + usage=usage, + metadata={"finish_reason": finish_reason}, + ) + + @abstractmethod + def _parse_usage(self, chunks: list[TextGenerationChunk]) -> TextGenerationUsage: + """Parse usage from chunks (provider-specific).""" + + +__all__ = ["TextGenerationStream"] diff --git a/pyproject.toml b/pyproject.toml index 5c4f79fd..4ac5dd7b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dev = [ "pytest-cov>=7.0", "pytest-randomly>=4.0", "pytest-asyncio>=1.2.0", + "pytest-xdist>=3.0.0", "ruff>=0.8.0", "mypy>=1.13.0", "types-requests>=2.31.0", @@ -63,6 +64,7 @@ asyncio_mode = "auto" markers = [ "slow: marks tests as slow (deselect with '-m \"not slow\"')", "smoke: quick checks for critical paths", + "integration: marks tests as integration tests (require API keys)", ] [tool.coverage.run] diff --git a/tests/integration_tests/__init__.py b/tests/integration_tests/__init__.py new file mode 100644 index 00000000..3bcf811d --- /dev/null +++ b/tests/integration_tests/__init__.py @@ -0,0 +1 @@ +"""Integration tests for Celeste AI.""" diff --git a/tests/integration_tests/test_text_generation/__init__.py b/tests/integration_tests/test_text_generation/__init__.py new file mode 100644 index 00000000..f264f567 --- /dev/null +++ b/tests/integration_tests/test_text_generation/__init__.py @@ -0,0 +1 @@ +"""Integration tests for text generation capability.""" diff --git a/tests/integration_tests/test_text_generation/test_generate.py b/tests/integration_tests/test_text_generation/test_generate.py new file mode 100644 index 00000000..028a436e --- /dev/null +++ b/tests/integration_tests/test_text_generation/test_generate.py @@ -0,0 +1,58 @@ +"""Integration tests for text generation across all providers.""" + +import pytest +from celeste_text_generation import TextGenerationOutput, TextGenerationUsage +from celeste_text_generation.parameters import TextGenerationParameters + +from celeste import Capability, Provider, create_client + + +@pytest.mark.parametrize( + ("provider", "model", "parameters"), + [ + (Provider.OPENAI, "gpt-3.5-turbo", {}), + (Provider.ANTHROPIC, "claude-haiku-4-5", {}), + (Provider.GOOGLE, "gemini-2.5-flash-lite", {"thinking_budget": 0}), + (Provider.MISTRAL, "mistral-tiny", {}), + (Provider.COHERE, "command-a-03-2025", {}), + ], +) +@pytest.mark.integration +@pytest.mark.asyncio +async def test_generate( + provider: Provider, model: str, parameters: TextGenerationParameters +) -> None: + """Test text generation with max_tokens parameter across all providers. + + This test demonstrates that the unified API works identically across + all providers using the same code - proving the abstraction value. + """ + # Arrange + client = create_client( + capability=Capability.TEXT_GENERATION, + provider=provider, + ) + prompt = "Hi" + max_tokens = 30 + + # Act + response = await client.generate( + prompt=prompt, + model=model, + max_tokens=max_tokens, + **parameters, + ) + + # Assert + assert isinstance(response, TextGenerationOutput), ( + f"Expected TextGenerationOutput, got {type(response)}" + ) + assert isinstance(response.content, str), ( + f"Expected str content, got {type(response.content)}" + ) + assert len(response.content) > 0, f"Content is empty: {response.content!r}" + + # Validate usage metrics + assert isinstance(response.usage, TextGenerationUsage), ( + f"Expected TextGenerationUsage, got {type(response.usage)}" + )