diff --git a/packages/capabilities/text-generation/src/celeste_text_generation/client.py b/packages/capabilities/text-generation/src/celeste_text_generation/client.py
index 4b08856c..03cc2db6 100644
--- a/packages/capabilities/text-generation/src/celeste_text_generation/client.py
+++ b/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)"
diff --git a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/__init__.py b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/__init__.py
index f264f567..432e7b47 100644
--- a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/__init__.py
+++ b/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."""
diff --git a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py
index d3f57e15..2a06f34d 100644
--- a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py
+++ b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_generate.py
@@ -31,6 +31,7 @@ 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
@@ -38,7 +39,6 @@ async def test_generate(provider: Provider, model: str, parameters: dict) -> Non
     # Act
     response = await client.generate(
         prompt=prompt,
-        model=model,
         max_tokens=max_tokens,
         **parameters,
     )
diff --git a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py
index b67b1ad6..f0a6e760 100644
--- a/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py
+++ b/packages/capabilities/text-generation/tests/integration_tests/test_text_generation/test_stream.py
@@ -37,6 +37,7 @@ 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
@@ -44,7 +45,6 @@ async def test_stream(provider: Provider, model: str, parameters: dict) -> None:
     # Act - Create stream
     stream = client.stream(
         prompt=prompt,
-        model=model,
         max_tokens=max_tokens,
         **parameters,
     )