Split formatter across multiple stages

Signed-off-by: Samuel Monson <[email protected]>

Author: sjmonson

src/guidellm/backends/openai.py

@@ -12,14 +12,19 @@
 
 import asyncio
 import time
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from typing import Any
 
 import httpx
 
 from guidellm.backends.backend import Backend
 from guidellm.backends.response_handlers import GenerationResponseHandlerFactory
-from guidellm.schemas import GenerationRequest, GenerationResponse, RequestInfo
+from guidellm.schemas import (
+    GenerationRequest,
+    GenerationRequestArguments,
+    GenerationResponse,
+    RequestInfo,
+)
 
 __all__ = ["OpenAIHTTPBackend"]
 
@@ -59,6 +64,10 @@ def __init__(
         follow_redirects: bool = True,
         verify: bool = False,
         validate_backend: bool | str | dict[str, Any] = True,
+        stream: bool = True,
+        extras: dict[str, Any] | GenerationRequestArguments | None = None,
+        max_tokens: int | None = None,
+        max_completion_tokens: int | None = None,
     ):
         """
         Initialize OpenAI HTTP backend with server configuration.
@@ -96,11 +105,28 @@ def __init__(
         self.validate_backend: dict[str, Any] | None = self._resolve_validate_kwargs(
             validate_backend
         )
+        self.stream: bool = stream
+        self.extras = (
+            GenerationRequestArguments(**extras)
+            if extras and isinstance(extras, dict)
+            else extras
+        )
+        self.max_tokens: int | None = max_tokens or max_completion_tokens
 
         # Runtime state
         self._in_process = False
         self._async_client: httpx.AsyncClient | None = None
 
+        # TODO: Find a better way to register formatters
+        self.request_formatters: dict[
+            str, Callable[[GenerationRequest], GenerationRequestArguments]
+        ] = {
+            "text_completions": self.formatter_text_completions,
+            "chat_completions": self.formatter_chat_completions,
+            "audio_transcriptions": self.formatter_audio_transcriptions,
+            "audio_translations": self.formatter_audio_transcriptions,
+        }
+
     @property
     def info(self) -> dict[str, Any]:
         """
@@ -227,31 +253,35 @@ async def resolve(  # type: ignore[override]
         if history is not None:
             raise NotImplementedError("Multi-turn requests not yet supported")
 
+        arguments: GenerationRequestArguments = self.request_formatters[
+            request.request_type
+        ](request)
+
         if (request_path := self.api_routes.get(request.request_type)) is None:
             raise ValueError(f"Unsupported request type '{request.request_type}'")
 
         request_url = f"{self.target}/{request_path}"
         request_files = (
             {
                 key: tuple(value) if isinstance(value, list) else value
-                for key, value in request.arguments.files.items()
+                for key, value in arguments.files.items()
             }
-            if request.arguments.files
+            if arguments.files
             else None
         )
-        request_json = request.arguments.body if not request_files else None
-        request_data = request.arguments.body if request_files else None
+        request_json = arguments.body if not request_files else None
+        request_data = arguments.body if request_files else None
         response_handler = GenerationResponseHandlerFactory.create(
             request.request_type, handler_overrides=self.response_handlers
         )
 
-        if not request.arguments.stream:
+        if not arguments.stream:
             request_info.timings.request_start = time.time()
             response = await self._async_client.request(
-                request.arguments.method or "POST",
+                arguments.method or "POST",
                 request_url,
-                params=request.arguments.params,
-                headers=request.arguments.headers,
+                params=arguments.params,
+                headers=arguments.headers,
                 json=request_json,
                 data=request_data,
                 files=request_files,
@@ -266,10 +296,10 @@ async def resolve(  # type: ignore[override]
             request_info.timings.request_start = time.time()
 
             async with self._async_client.stream(
-                request.arguments.method or "POST",
+                arguments.method or "POST",
                 request_url,
-                params=request.arguments.params,
-                headers=request.arguments.headers,
+                params=arguments.params,
+                headers=arguments.headers,
                 json=request_json,
                 data=request_data,
                 files=request_files,
@@ -332,3 +362,177 @@ def _resolve_validate_kwargs(
             validate_kwargs["method"] = "GET"
 
         return validate_kwargs
+
+    def formatter_text_completions(
+        self, data: GenerationRequest
+    ) -> GenerationRequestArguments:
+        arguments: GenerationRequestArguments = GenerationRequestArguments()
+        arguments.body = {}  # The type checker works better setting this field here
+
+        # Add model
+        if self.model is not None:
+            arguments.body["model"] = self.model
+
+        # Configure streaming
+        if self.stream:
+            arguments.stream = True
+            arguments.body["stream"] = True
+            arguments.body["stream_options"] = {"include_usage": True}
+
+        # Handle output tokens
+        if data.output_metrics.text_tokens:
+            arguments.body["max_tokens"] = data.output_metrics.text_tokens
+            arguments.body["stop"] = None
+            arguments.body["ignore_eos"] = True
+        elif self.max_tokens is not None:
+            arguments.body["max_tokens"] = self.max_tokens
+
+        # Apply extra arguments
+        if self.extras:
+            arguments.model_combine(self.extras)
+
+        # Build prompt
+        prefix = "".join(pre for pre in data.columns.get("prefix_column", []) if pre)
+        text = "".join(txt for txt in data.columns.get("text_column", []) if txt)
+        if prefix or text:
+            prompt = prefix + text
+            arguments.body["prompt"] = prompt
+
+        return arguments
+
+    def formatter_chat_completions(  # noqa: C901, PLR0912, PLR0915
+        self, data: GenerationRequest
+    ) -> GenerationRequestArguments:
+        arguments = GenerationRequestArguments()
+        arguments.body = {}  # The type checker works best with body assigned here
+
+        # Add model
+        if self.model is not None:
+            arguments.body["model"] = self.model
+
+        # Configure streaming
+        if self.stream:
+            arguments.stream = True
+            arguments.body["stream"] = True
+            arguments.body["stream_options"] = {"include_usage": True}
+
+        # Handle output tokens
+        if data.output_metrics.text_tokens:
+            arguments.body.update(
+                {
+                    "max_completion_tokens": data.output_metrics.text_tokens,
+                    "stop": None,
+                    "ignore_eos": True,
+                }
+            )
+        elif self.max_tokens is not None:
+            arguments.body["max_completion_tokens"] = self.max_tokens
+
+        # Apply extra arguments
+        if self.extras:
+            arguments.model_combine(self.extras)
+
+        # Build messages
+        arguments.body["messages"] = []
+
+        for prefix in data.columns.get("prefix_column", []):
+            if not prefix:
+                continue
+
+            arguments.body["messages"].append({"role": "system", "content": prefix})
+
+        for text in data.columns.get("text_column", []):
+            if not text:
+                continue
+
+            arguments.body["messages"].append(
+                {"role": "user", "content": [{"type": "text", "text": text}]}
+            )
+
+        for image in data.columns.get("image_column", []):
+            if not image:
+                continue
+
+            arguments.body["messages"].append(
+                {
+                    "role": "user",
+                    "content": [{"type": "image_url", "image_url": image.get("image")}],
+                }
+            )
+
+        for video in data.columns.get("video_column", []):
+            if not video:
+                continue
+
+            arguments.body["messages"].append(
+                {
+                    "role": "user",
+                    "content": [{"type": "video_url", "video_url": video.get("video")}],
+                }
+            )
+
+        for audio in data.columns.get("audio_column", []):
+            if not audio:
+                continue
+
+            arguments.body["messages"].append(
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "input_audio",
+                            "input_audio": {
+                                "data": audio.get("audio"),
+                                "format": audio.get("format"),
+                            },
+                        }
+                    ],
+                }
+            )
+
+        return arguments
+
+    def formatter_audio_transcriptions(  # noqa: C901
+        self, data: GenerationRequest
+    ) -> GenerationRequestArguments:
+        arguments = GenerationRequestArguments(files={})
+        arguments.body = {}
+
+        # Add model
+        if self.model is not None:
+            arguments.body["model"] = self.model
+
+        # Configure streaming
+        if self.stream:
+            arguments.stream = True
+            arguments.body["stream"] = True
+            arguments.body["stream_options"] = {"include_usage": True}
+
+        # Apply extra arguments
+        if self.extras:
+            arguments.model_combine(self.extras)
+
+        # Build audio input
+        audio_columns = data.columns.get("audio_column", [])
+        if len(audio_columns) != 1:
+            raise ValueError(
+                f"GenerativeAudioTranscriptionRequestFormatter expects exactly "
+                f"one audio column, but got {len(audio_columns)}."
+            )
+
+        arguments.files = {
+            "file": (
+                audio_columns[0].get("file_name", "audio_input"),
+                audio_columns[0].get("audio"),
+                audio_columns[0].get("mimetype"),
+            )
+        }
+
+        # Build prompt
+        prefix = "".join(pre for pre in data.columns.get("prefix_column", []) if pre)
+        text = "".join(txt for txt in data.columns.get("text_column", []) if txt)
+        if prefix or text:
+            prompt = prefix + text
+            arguments.body["prompt"] = prompt
+
+        return arguments

src/guidellm/data/__init__.py

@@ -4,12 +4,12 @@
     DatasetDeserializer,
     DatasetDeserializerFactory,
 )
+from .finalizers import DatasetFinalizer, FinalizerRegistry
 from .loaders import DataLoader, DatasetsIterator
 from .preprocessors import (
     DataDependentPreprocessor,
     DatasetPreprocessor,
     PreprocessorRegistry,
-    RequestFormatter,
 )
 from .processor import ProcessorFactory
 from .schemas import GenerativeDatasetColumnType
@@ -20,11 +20,12 @@
     "DataNotSupportedError",
     "DatasetDeserializer",
     "DatasetDeserializerFactory",
+    "DatasetFinalizer",
     "DatasetPreprocessor",
     "DatasetsIterator",
+    "FinalizerRegistry",
     "GenerativeDatasetColumnType",
     "GenerativeRequestCollator",
     "PreprocessorRegistry",
     "ProcessorFactory",
-    "RequestFormatter",
 ]