Add output format for other LLMs

This commit is contained in:
Nicolas Mowen 2026-03-09 17:58:07 -06:00
parent 65dacf29ff
commit 720d987894
4 changed files with 50 additions and 14 deletions

View File

@ -42,13 +42,18 @@ class OpenAIClient(GenAIClient):
azure_endpoint=azure_endpoint,
)
def _send(self, prompt: str, images: list[bytes]) -> Optional[str]:
def _send(
self,
prompt: str,
images: list[bytes],
response_format: Optional[dict] = None,
) -> Optional[str]:
"""Submit a request to Azure OpenAI."""
encoded_images = [base64.b64encode(image).decode("utf-8") for image in images]
try:
result = self.provider.chat.completions.create(
model=self.genai_config.model,
messages=[
request_params = {
"model": self.genai_config.model,
"messages": [
{
"role": "user",
"content": [{"type": "text", "text": prompt}]
@ -64,9 +69,12 @@ class OpenAIClient(GenAIClient):
],
},
],
timeout=self.timeout,
"timeout": self.timeout,
**self.genai_config.runtime_options,
)
}
if response_format:
request_params["response_format"] = response_format
result = self.provider.chat.completions.create(**request_params)
except Exception as e:
logger.warning("Azure OpenAI returned an error: %s", str(e))
return None

View File

@ -42,7 +42,12 @@ class GeminiClient(GenAIClient):
http_options=types.HttpOptions(**http_options_dict),
)
def _send(self, prompt: str, images: list[bytes]) -> Optional[str]:
def _send(
self,
prompt: str,
images: list[bytes],
response_format: Optional[dict] = None,
) -> Optional[str]:
"""Submit a request to Gemini."""
contents = [
types.Part.from_bytes(data=img, mime_type="image/jpeg") for img in images
@ -52,6 +57,12 @@ class GeminiClient(GenAIClient):
generation_config_dict = {"candidate_count": 1}
generation_config_dict.update(self.genai_config.runtime_options)
if response_format and response_format.get("type") == "json_schema":
generation_config_dict["response_mime_type"] = "application/json"
schema = response_format.get("json_schema", {}).get("schema")
if schema:
generation_config_dict["response_schema"] = schema
response = self.provider.models.generate_content(
model=self.genai_config.model,
contents=contents,

View File

@ -53,7 +53,12 @@ class OllamaClient(GenAIClient):
logger.warning("Error initializing Ollama: %s", str(e))
return None
def _send(self, prompt: str, images: list[bytes]) -> Optional[str]:
def _send(
self,
prompt: str,
images: list[bytes],
response_format: Optional[dict] = None,
) -> Optional[str]:
"""Submit a request to Ollama"""
if self.provider is None:
logger.warning(
@ -65,6 +70,10 @@ class OllamaClient(GenAIClient):
**self.provider_options,
**self.genai_config.runtime_options,
}
if response_format and response_format.get("type") == "json_schema":
schema = response_format.get("json_schema", {}).get("schema")
if schema:
ollama_options["format"] = schema
result = self.provider.generate(
self.genai_config.model,
prompt,

View File

@ -36,7 +36,12 @@ class OpenAIClient(GenAIClient):
return OpenAI(api_key=self.genai_config.api_key, **provider_opts)
def _send(self, prompt: str, images: list[bytes]) -> Optional[str]:
def _send(
self,
prompt: str,
images: list[bytes],
response_format: Optional[dict] = None,
) -> Optional[str]:
"""Submit a request to OpenAI."""
encoded_images = [base64.b64encode(image).decode("utf-8") for image in images]
messages_content = []
@ -57,17 +62,20 @@ class OpenAIClient(GenAIClient):
}
)
try:
result = self.provider.chat.completions.create(
model=self.genai_config.model,
messages=[
request_params = {
"model": self.genai_config.model,
"messages": [
{
"role": "user",
"content": messages_content,
},
],
timeout=self.timeout,
"timeout": self.timeout,
**self.genai_config.runtime_options,
)
}
if response_format:
request_params["response_format"] = response_format
result = self.provider.chat.completions.create(**request_params)
if (
result is not None
and hasattr(result, "choices")