chihlasm/resolutionflow
1
0
Fork 0
Code Issues 23 Pull Requests 1 Actions Packages Projects Releases Wiki Activity

feat(ai): robust response extraction + structured-output foundation

Browse Source 
Harden the Anthropic provider and lay the groundwork for schema-constrained
JSON, optimizing the existing claude-sonnet-4-6 / claude-haiku-4-5 usage
(no model changes).

ai_provider.py:
- _extract_text_from_response replaces fragile response.content[0].text:
  skips non-text leading blocks (e.g. thinking), returns the first text
  block, logs an anthropic.stop_reason warning on max_tokens/refusal
  (truncation now observable), and raises ValueError on a no-text response.
- generate_json gains an optional `schema` param. Anthropic wires it to
  output_config.format (structured outputs); schema=None preserves the exact
  prior call for every existing caller. Gemini accepts-and-ignores it.

kb_conversion_service.py:
- TROUBLESHOOTING_SCHEMA / PROCEDURAL_SCHEMA + _schema_for_target_type(),
  modelled as a strict superset of every field the prompts emit.
- convert_document passes the schema only when the new
  AI_KB_CONVERT_STRUCTURED_OUTPUT setting is True (default False). The
  _try_repair_json fallback stays as belt-and-suspenders.

Tests: 14 provider + 7 schema, TDD (red-green). Live constrained-decoding
smoke-test still required before enabling the flag in production.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Michael Chihlas
2026-05-28 14:54:30 -04:00
parent b1ee46656e
commit 067574ad6a
5 changed files with 503 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

72
backend/app/core/ai_provider.py
View File

@@ -147,6 +147,40 @@ def build_anthropic_chat_messages(
return messages
def _extract_text_from_response(response: Any, model: str) -> str:
"""Return the first text block's text from an Anthropic message response.
Robustness over the naive ``response.content[0].text``:
- Skips non-text leading blocks (e.g. ``thinking``) and returns the first
block whose ``type == "text"``. Indexing ``content[0]`` blindly throws or
returns garbage the moment a non-text block leads the response.
- Surfaces truncation/refusal: when ``stop_reason`` is ``max_tokens`` or
``refusal``, emits a structured warning so silent output corruption
(truncated JSON, empty refusals) is observable rather than handed
downstream to be guessed at.
- Raises ``ValueError`` when no text block is present (e.g. a bare refusal)
instead of returning a non-text block's attributes.
"""
stop_reason = getattr(response, "stop_reason", None)
if stop_reason in ("max_tokens", "refusal"):
logger.warning(
"anthropic.stop_reason",
extra={
"event": "anthropic.stop_reason",
"model": model,
"stop_reason": stop_reason,
},
)
for block in response.content:
if getattr(block, "type", None) == "text":
return block.text
raise ValueError(
f"Anthropic response contained no text block (stop_reason={stop_reason!r})"
)
def _log_anthropic_cache_usage(usage: Any, model: str) -> None:
"""Emit a structured log line capturing cache_read / cache_creation tokens."""
cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
@@ -176,6 +210,7 @@ class AIProvider(ABC):
system_prompt: str | list[SystemBlock],
messages: list[dict[str, Any]],
max_tokens: int = 4096,
schema: dict[str, Any] | None = None,
) -> tuple[str, int, int]:
"""Generate a JSON response from the AI model.
@@ -185,6 +220,15 @@ class AIProvider(ABC):
Anthropic prompt caching per module-docstring policy.
messages: List of message dicts with "role" and "content" keys.
max_tokens: Maximum output tokens.
schema: Optional JSON Schema constraining the response shape.
When provided, the Anthropic backend uses structured outputs
(`output_config.format`) to guarantee valid, parseable JSON —
no markdown fences, no truncated-brace repair. Must satisfy the
structured-output schema limits (every object needs
`additionalProperties: false`; no recursion; numeric/string
constraints are stripped). `None` preserves the legacy
prompt-only behavior. The Gemini backend currently ignores this
argument (it already requests `application/json`).
Returns:
Tuple of (response_text, input_tokens, output_tokens).
@@ -231,7 +275,11 @@ class GeminiProvider(AIProvider):
system_prompt: str | list[SystemBlock],
messages: list[dict[str, Any]],
max_tokens: int = 4096,
schema: dict[str, Any] | None = None,
) -> tuple[str, int, int]:
# `schema` is accepted for interface parity but ignored: Gemini already
# constrains output via response_mime_type="application/json" below.
# Mapping JSON Schema -> Gemini response_schema is deferred.
from google import genai
from google.genai import types as genai_types
@@ -362,18 +410,28 @@ class AnthropicProvider(AIProvider):
system_prompt: str | list[SystemBlock],
messages: list[dict[str, Any]],
max_tokens: int = 4096,
schema: dict[str, Any] | None = None,
) -> tuple[str, int, int]:
client = _get_anthropic_client(self._api_key, self._timeout)
normalized_system = _normalize_system_for_anthropic(system_prompt)
response = await client.messages.create(
model=self._model,
max_tokens=max_tokens,
system=normalized_system,
messages=messages,
)
create_kwargs: dict[str, Any] = {
"model": self._model,
"max_tokens": max_tokens,
"system": normalized_system,
"messages": messages,
}
if schema is not None:
# Structured outputs: constrain the response to valid JSON matching
# the schema (Sonnet 4.6 / Haiku 4.5). Removes the need for
# markdown-fence stripping and truncated-JSON repair downstream.
create_kwargs["output_config"] = {
"format": {"type": "json_schema", "schema": schema}
}
text = response.content[0].text
response = await client.messages.create(**create_kwargs)
text = _extract_text_from_response(response, self._model)
input_tokens = response.usage.input_tokens
output_tokens = response.usage.output_tokens

6
backend/app/core/config.py
View File

@@ -155,6 +155,12 @@ class Settings(BaseSettings):
AI_CONVERSATION_TTL_HOURS: int = 24
AI_MAX_CALLS_PER_FLOW: int = 10
AI_REQUEST_TIMEOUT_SECONDS: int = 120
# When True, KB conversion constrains the Anthropic response with a JSON
# schema (structured outputs) instead of relying on prompt-only JSON +
# downstream fence-stripping / brace-repair. Default OFF: enable in staging
# and smoke-test constrained decoding against the live model before turning
# it on in production. Only affects the Anthropic backend.
AI_KB_CONVERT_STRUCTURED_OUTPUT: bool = False
# AI Provider selection
AI_PROVIDER: str = "anthropic" # "gemini" or "anthropic"
GOOGLE_AI_API_KEY: Optional[str] = None

120
backend/app/core/kb_conversion_service.py
View File

@@ -202,6 +202,115 @@ the engineer attached, NOT from this schema):
9. Return ONLY valid JSON — no markdown fences, no explanation text."""
# ── Structured-output schemas ──
#
# These constrain the model's JSON via Anthropic structured outputs
# (output_config.format) so the response is guaranteed valid and parseable —
# no markdown fences, no truncated-brace repair. They must be a SUPERSET of
# every field the corresponding system prompt instructs the model to emit:
# additionalProperties is False everywhere, so any field the prompt asks for
# but the schema omits would be impossible to produce.
#
# `type`/`field_type` are intentionally left as plain strings (no enum): the
# downstream parser already normalizes/tolerates the type values, and an enum
# risks constraining the model away from a value the prompt would yield.
_TROUBLESHOOTING_OPTION_SCHEMA: dict[str, Any] = {
"type": "object",
"properties": {
"label": {"type": "string"},
"next_node_id": {"type": "string"},
},
"required": ["label", "next_node_id"],
"additionalProperties": False,
}
_TROUBLESHOOTING_NODE_SCHEMA: dict[str, Any] = {
"type": "object",
"properties": {
"id": {"type": "string"},
"type": {"type": "string"},
"question": {"type": "string"},
"options": {"type": "array", "items": _TROUBLESHOOTING_OPTION_SCHEMA},
"next_node_id": {"type": "string"},
"confidence": {"type": "number"},
"source_excerpt": {"type": "string"},
},
# Only the universal fields are required. `question`/`options`/`next_node_id`
# vary by node type and stay optional so a resolution node need not carry
# options and an action node need not carry a question.
"required": ["id", "type", "confidence", "source_excerpt"],
"additionalProperties": False,
}
TROUBLESHOOTING_SCHEMA: dict[str, Any] = {
"type": "object",
"properties": {
"title": {"type": "string"},
"description": {"type": "string"},
"nodes": {"type": "array", "items": _TROUBLESHOOTING_NODE_SCHEMA},
},
"required": ["title", "description", "nodes"],
"additionalProperties": False,
}
_PROCEDURAL_STEP_SCHEMA: dict[str, Any] = {
"type": "object",
"properties": {
"id": {"type": "string"},
"type": {"type": "string"},
"content": {"type": "string"},
"confidence": {"type": "number"},
"source_excerpt": {"type": "string"},
},
"required": ["id", "type", "content", "confidence", "source_excerpt"],
"additionalProperties": False,
}
_PROCEDURAL_INTAKE_SCHEMA: dict[str, Any] = {
"type": "object",
"properties": {
"variable_name": {"type": "string"},
"label": {"type": "string"},
"field_type": {"type": "string"},
"required": {"type": "boolean"},
"display_order": {"type": "integer"},
},
"required": [
"variable_name",
"label",
"field_type",
"required",
"display_order",
],
"additionalProperties": False,
}
PROCEDURAL_SCHEMA: dict[str, Any] = {
"type": "object",
"properties": {
"title": {"type": "string"},
"description": {"type": "string"},
"steps": {"type": "array", "items": _PROCEDURAL_STEP_SCHEMA},
"intake_form": {"type": "array", "items": _PROCEDURAL_INTAKE_SCHEMA},
},
"required": ["title", "description", "steps", "intake_form"],
"additionalProperties": False,
}
def _schema_for_target_type(target_type: str) -> dict[str, Any]:
"""Return the structured-output schema for a KB conversion target type.
Mirrors the prompt selection in ``convert_document``: only
``"troubleshooting"`` uses the decision-tree schema; everything else is
treated as a procedural flow.
"""
if target_type == "troubleshooting":
return TROUBLESHOOTING_SCHEMA
return PROCEDURAL_SCHEMA
def _build_user_message(
source_text: str,
source_metadata: dict[str, Any] | None,
@@ -404,6 +513,16 @@ async def convert_document(
model = settings.get_model_for_action("kb_convert")
provider = get_ai_provider(model=model)
# Structured outputs (flagged): constrain the response to a JSON schema so
# the model can't emit fences or truncated JSON. Falls back to prompt-only
# JSON (schema=None) when disabled; the parse path below stays intact either
# way as a belt-and-suspenders fallback.
schema = (
_schema_for_target_type(kb_import.target_type)
if settings.AI_KB_CONVERT_STRUCTURED_OUTPUT
else None
)
try:
raw_text, input_tokens, output_tokens = await provider.generate_json(
system_prompt=[
@@ -414,6 +533,7 @@ async def convert_document(
],
messages=[{"role": "user", "content": user_message}],
max_tokens=16384,
schema=schema,
)
except Exception as e:
logger.error("AI conversion failed for kb_import=%s: %s", kb_import.id, e)