Enforce typed responses for pipeline stagesxiang

1 files changed, 16 insertions, 1 deletions

diff --git a/src/gap_pipeline/clients.py b/src/gap_pipeline/clients.py
index 3b3e40f..bc35742 100644
--- a/src/gap_pipeline/clients.py
+++ b/src/gap_pipeline/clients.py
@@ -8,6 +8,8 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Protocol
 
+from pydantic import BaseModel
+
 
 @dataclass(frozen=True)
 class ModelResponse:
@@ -27,6 +29,7 @@ class JsonLLM(Protocol):
         system_prompt: str,
         user_prompt: str,
         request_id: str,
+        response_model: type[BaseModel] | None = None,
     ) -> ModelResponse: ...
 
 
@@ -61,14 +64,24 @@ class OpenAIJsonClient:
         system_prompt: str,
         user_prompt: str,
         request_id: str,
+        response_model: type[BaseModel] | None = None,
     ) -> ModelResponse:
+        if response_model is None:
+            response_format: dict[str, Any] = {"type": "json_object"}
+        else:
+            from openai.lib._parsing._completions import (
+                type_to_response_format_param,
+            )
+
+            response_format = type_to_response_format_param(response_model)
+
         kwargs: dict[str, Any] = {
             "model": self.model,
             "messages": [
                 {"role": "system", "content": system_prompt},
                 {"role": "user", "content": user_prompt},
             ],
-            "response_format": {"type": "json_object"},
+            "response_format": response_format,
             "max_completion_tokens": self.max_completion_tokens,
         }
         if self.seed is not None:
@@ -111,6 +124,7 @@ class ReplayClient:
         system_prompt: str,
         user_prompt: str,
         request_id: str,
+        response_model: type[BaseModel] | None = None,
     ) -> ModelResponse:
         if request_id not in self._responses:
             raise KeyError(f"no replay response for {request_id}")
@@ -145,6 +159,7 @@ class ScriptedClient:
         system_prompt: str,
         user_prompt: str,
         request_id: str,
+        response_model: type[BaseModel] | None = None,
     ) -> ModelResponse:
         self.calls.append(request_id)
         prefixes = sorted(