From 5a998924a3c5cac82742fc1733f7e3bd82d82753 Mon Sep 17 00:00:00 2001 From: Motty Chen Date: Fri, 10 Apr 2026 17:22:59 -0500 Subject: [PATCH] fix(#309): provider-specific JSON Schema for structured LLM output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add llm_json_schema_adaptation: oneOf→anyOf, strip discriminator for OpenAI strict mode - Apply same union normalization for Vertex Gemini response_schema - Bedrock unchanged (passes through; ignores schema) - Unit tests for EmailInsightResponse / provider routing Made-with: Cursor --- .../ai_engine/llm_json_schema_adaptation.py | 81 +++++++++++++++++++ .../ai_engine/unified_ai_engine.py | 37 ++++++++- .../test_llm_json_schema_adaptation.py | 59 ++++++++++++++ 3 files changed, 173 insertions(+), 4 deletions(-) create mode 100644 coaching/src/application/ai_engine/llm_json_schema_adaptation.py create mode 100644 coaching/tests/unit/application/ai_engine/test_llm_json_schema_adaptation.py diff --git a/coaching/src/application/ai_engine/llm_json_schema_adaptation.py b/coaching/src/application/ai_engine/llm_json_schema_adaptation.py new file mode 100644 index 00000000..80ae1231 --- /dev/null +++ b/coaching/src/application/ai_engine/llm_json_schema_adaptation.py @@ -0,0 +1,81 @@ +"""Provider-specific JSON Schema adaptation for LLM structured output. + +Pydantic ``model_json_schema()`` emits ``oneOf`` for discriminated unions (e.g. +``EmailInsightResponse.blocks``). OpenAI strict structured outputs reject +``oneOf`` in some positions (e.g. under ``items``) with ``invalid_json_schema``. + +This module adapts already-prepared schemas (e.g. after +``additionalProperties: false`` normalization) for each provider's accepted +dialect without changing the application-level Pydantic models: API responses +are still validated with the original models after generation. + +References: +- https://platform.openai.com/docs/guides/structured_outputs +- https://ai.google.dev/gemini-api/docs/structured-output +""" + +from __future__ import annotations + +import copy +from typing import Any + + +def adapt_json_schema_for_openai_structured_output(schema: dict[str, Any]) -> dict[str, Any]: + """Return a deep copy of ``schema`` adjusted for OpenAI Responses API strict JSON schema. + + - Recursively renames ``oneOf`` to ``anyOf`` (union of object shapes). + - Strips ``discriminator`` (OpenAI strict mode does not use Pydantic's discriminator metadata). + + Args: + schema: Schema already passed through strict OpenAI prep (e.g. + ``UnifiedAIEngine._prepare_schema_for_structured_output``). + + Returns: + Adapted schema safe to pass as ``text.format.schema`` with strict mode. + """ + adapted = copy.deepcopy(schema) + _transform_openai_inplace(adapted) + return adapted + + +def adapt_json_schema_for_vertex_response_schema(schema: dict[str, Any]) -> dict[str, Any]: + """Return a deep copy of ``schema`` adjusted for Vertex Gemini ``response_schema``. + + Uses the same ``oneOf`` → ``anyOf`` normalization as OpenAI for union blocks; + Gemini's JSON schema subset also commonly rejects or mishandles ``oneOf`` in + array ``items``. If Google diverges, split logic here. + + Args: + schema: Schema after ``_prepare_schema_for_structured_output`` (or equivalent). + + Returns: + Adapted schema for ``GenerateContentConfig.response_schema``. + """ + adapted = copy.deepcopy(schema) + _transform_openai_inplace(adapted) + return adapted + + +def _transform_openai_inplace(node: Any) -> None: + """Recursively apply OpenAI/Gemini-friendly union transforms in place.""" + if isinstance(node, dict): + if "oneOf" in node: + node["anyOf"] = node.pop("oneOf") + node.pop("discriminator", None) + + defs = node.get("$defs") + if isinstance(defs, dict): + for defn in defs.values(): + _transform_openai_inplace(defn) + + for key, value in node.items(): + if key == "$defs": + continue + if isinstance(value, dict): + _transform_openai_inplace(value) + elif isinstance(value, list): + for item in value: + _transform_openai_inplace(item) + elif isinstance(node, list): + for item in node: + _transform_openai_inplace(item) diff --git a/coaching/src/application/ai_engine/unified_ai_engine.py b/coaching/src/application/ai_engine/unified_ai_engine.py index 889ba67b..5ad21309 100644 --- a/coaching/src/application/ai_engine/unified_ai_engine.py +++ b/coaching/src/application/ai_engine/unified_ai_engine.py @@ -12,6 +12,12 @@ from typing import TYPE_CHECKING, Any import structlog +from pydantic import BaseModel + +from coaching.src.application.ai_engine.llm_json_schema_adaptation import ( + adapt_json_schema_for_openai_structured_output, + adapt_json_schema_for_vertex_response_schema, +) from coaching.src.application.ai_engine.response_serializer import ResponseSerializer from coaching.src.application.llm_usage.llm_invocation_context import LlmInvocationContext from coaching.src.application.llm_usage.llm_usage_recording_service import LlmUsageRecordingService @@ -26,7 +32,6 @@ from coaching.src.infrastructure.llm.provider_factory import LLMProviderFactory from coaching.src.repositories.topic_repository import TopicRepository from coaching.src.services.s3_prompt_storage import S3PromptStorage -from pydantic import BaseModel if TYPE_CHECKING: from coaching.src.services.template_parameter_processor import TemplateParameterProcessor @@ -551,6 +556,11 @@ async def _build_single_shot_context( ) provider, model_name = self.provider_factory.get_provider_for_model(model_code) + response_schema_for_provider = self._adapt_structured_schema_for_llm_provider( + response_schema, + provider.provider_name, + ) + # Step 7: Call LLM with topic configuration messages = [LLMMessage(role="user", content=rendered_user)] @@ -564,7 +574,7 @@ async def _build_single_shot_context( "max_tokens": topic.max_tokens, "messages": [{"role": msg.role, "content": msg.content} for msg in messages], "system_prompt": rendered_system, - "has_response_schema": response_schema is not None, + "has_response_schema": response_schema_for_provider is not None, }, ) @@ -590,7 +600,7 @@ async def _build_single_shot_context( temperature=topic.temperature, max_tokens=topic.max_tokens, system_prompt=rendered_system, - response_schema=response_schema, # Pass schema for structured output + response_schema=response_schema_for_provider, ) except Exception: wall_ms = int((time.perf_counter() - llm_start) * 1000) @@ -707,7 +717,7 @@ async def _build_single_shot_context( rendered_system_prompt=rendered_system, rendered_user_prompt=rendered_user, enriched_parameters=enriched_params, - response_schema=response_schema, + response_schema=response_schema_for_provider, llm_response=llm_response, serialized_response=serialized, ) @@ -1025,6 +1035,25 @@ def _inject_response_format_with_schema( ) return system_prompt, None + @staticmethod + def _adapt_structured_schema_for_llm_provider( + schema: dict[str, Any] | None, + provider_name: str, + ) -> dict[str, Any] | None: + """Map canonical strict schema to the dialect accepted by the active LLM provider. + + Pydantic unions (discriminated or not) often emit ``oneOf``; OpenAI strict + structured outputs reject ``oneOf`` in some nested positions. Vertex + Gemini may accept a similar adjustment. Bedrock ignores ``response_schema``. + """ + if schema is None: + return None + if provider_name == "openai": + return adapt_json_schema_for_openai_structured_output(schema) + if provider_name == "google_vertex": + return adapt_json_schema_for_vertex_response_schema(schema) + return schema + def _prepare_schema_for_structured_output( self, schema: dict[str, Any], diff --git a/coaching/tests/unit/application/ai_engine/test_llm_json_schema_adaptation.py b/coaching/tests/unit/application/ai_engine/test_llm_json_schema_adaptation.py new file mode 100644 index 00000000..686500ed --- /dev/null +++ b/coaching/tests/unit/application/ai_engine/test_llm_json_schema_adaptation.py @@ -0,0 +1,59 @@ +"""Tests for provider-specific JSON Schema adaptation (#309).""" + +from __future__ import annotations + +from coaching.src.application.ai_engine.llm_json_schema_adaptation import ( + adapt_json_schema_for_openai_structured_output, + adapt_json_schema_for_vertex_response_schema, +) +from coaching.src.application.ai_engine.unified_ai_engine import UnifiedAIEngine +from coaching.src.models.responses import EmailInsightResponse + + +def _prepare_like_engine(schema: dict, model_name: str) -> dict: + """Mirror UnifiedAIEngine._prepare_schema_for_structured_output without full engine.""" + engine = UnifiedAIEngine.__new__(UnifiedAIEngine) + return engine._prepare_schema_for_structured_output(schema, model_name) + + +def test_openai_adaptation_replaces_oneof_with_anyof_for_email_insight() -> None: + full = EmailInsightResponse.model_json_schema(by_alias=True) + prepared = _prepare_like_engine(full, "EmailInsightResponse") + assert "oneOf" in str(prepared) or any( + isinstance(v, dict) and "oneOf" in v for v in prepared.get("$defs", {}).values() + ) + + adapted = adapt_json_schema_for_openai_structured_output(prepared) + dumped = str(adapted) + assert "oneOf" not in dumped + assert "anyOf" in dumped + + +def test_openai_adaptation_removes_discriminator() -> None: + full = EmailInsightResponse.model_json_schema(by_alias=True) + prepared = _prepare_like_engine(full, "EmailInsightResponse") + adapted = adapt_json_schema_for_openai_structured_output(prepared) + assert "discriminator" not in str(adapted) + + +def test_vertex_adapter_matches_openai_union_handling() -> None: + full = EmailInsightResponse.model_json_schema(by_alias=True) + prepared = _prepare_like_engine(full, "EmailInsightResponse") + v = adapt_json_schema_for_vertex_response_schema(prepared) + assert "oneOf" not in str(v) + assert "anyOf" in str(v) + + +def test_unified_engine_adapt_routes_by_provider() -> None: + full = EmailInsightResponse.model_json_schema(by_alias=True) + prepared = _prepare_like_engine(full, "EmailInsightResponse") + o = UnifiedAIEngine._adapt_structured_schema_for_llm_provider(prepared, "openai") + g = UnifiedAIEngine._adapt_structured_schema_for_llm_provider(prepared, "google_vertex") + b = UnifiedAIEngine._adapt_structured_schema_for_llm_provider(prepared, "bedrock") + assert o is not None and "anyOf" in str(o) + assert g is not None and "anyOf" in str(g) + assert b == prepared + + +def test_adapt_none_returns_none() -> None: + assert UnifiedAIEngine._adapt_structured_schema_for_llm_provider(None, "openai") is None