refactor(text-generation): use StructuredOutput and improve API

packages/capabilities/text-generation/src/celeste_text_generation/client.py

@@ -4,10 +4,10 @@
 from typing import Any, Unpack
 
 import httpx
-from pydantic import BaseModel
 
 from celeste.client import Client
 from celeste.exceptions import ValidationError
+from celeste.types import StructuredOutput
 from celeste_text_generation.io import (
     TextGenerationFinishReason,
     TextGenerationInput,
@@ -35,24 +35,22 @@ def _parse_content(
         self,
         response_data: dict[str, Any],
         **parameters: Unpack[TextGenerationParameters],
-    ) -> str | BaseModel:
+    ) -> StructuredOutput:
         """Parse content from provider response."""
 
     @abstractmethod
     def _parse_finish_reason(
         self, response_data: dict[str, Any]
-    ) -> TextGenerationFinishReason | None:
+    ) -> TextGenerationFinishReason:
         """Parse finish reason from provider response."""
 
     def _create_inputs(
-        self,
-        *args: str,
-        prompt: str | None = None,
-        **parameters: Unpack[TextGenerationParameters],
+        self, *args: str, **parameters: Unpack[TextGenerationParameters]
     ) -> TextGenerationInput:
         """Map positional arguments to Input type."""
         if args:
             return TextGenerationInput(prompt=args[0])
+        prompt: str | None = parameters.get("prompt")
         if prompt is None:
             msg = (
                 "prompt is required (either as positional argument or keyword argument)"

packages/capabilities/text-generation/tests/integration_tests/test_text_generation/__init__.py

@@ -1 +1 @@
-"""Integration tests for text generation capability."""
+"""Text generation integration test module."""

packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py

@@ -31,14 +31,14 @@ async def test_generate(provider: Provider, model: str, parameters: dict) -> Non
     client = create_client(
         capability=Capability.TEXT_GENERATION,
         provider=provider,
+        model=model,
     )
     prompt = "Hi"
     max_tokens = 30
 
     # Act
     response = await client.generate(
         prompt=prompt,
-        model=model,
         max_tokens=max_tokens,
         **parameters,
     )

packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py

@@ -37,14 +37,14 @@ async def test_stream(provider: Provider, model: str, parameters: dict) -> None:
     client = create_client(
         capability=Capability.TEXT_GENERATION,
         provider=provider,
+        model=model,
     )
     prompt = "Hi"
     max_tokens = 30
 
     # Act - Create stream
     stream = client.stream(
         prompt=prompt,
-        model=model,
         max_tokens=max_tokens,
         **parameters,
     )