From 91365268634b8f70a37a5b8cd23daeb2deaf73b0 Mon Sep 17 00:00:00 2001 From: cybermaggedon Date: Fri, 10 Jul 2026 15:28:56 +0100 Subject: [PATCH] feat: LLM-native structured output via JSON schema enforcement (#1037) Thread existing JSON schemas from prompt definitions through the text-completion service to LLM backends' native structured output APIs. When a prompt has response-type "json" and a strict-mode compatible schema, the LLM constrains token selection at the logit level to guarantee schema-valid output. Wire-level changes: - Add response_format and schema fields to TextCompletionRequest - Update translator to encode/decode new fields - Pass new fields through LlmService, TextCompletionClient, and PromptManager Runtime schema compatibility checker: - New is_strict_mode_compatible() utility validates schemas against LLM provider constraints (additionalProperties, required fields, no unsupported constraints, no open-ended objects) - Per-prompt eligibility decision: compliant schemas use structured output, non-compliant schemas fall back to free-text + post-hoc validation LLM backend implementations: - OpenAI: response_format with json_schema, variant-aware top-level array rejection (openai variant blocks, llama/vllm variants allow) - New vllm variant for the OpenAI backend - vLLM (dedicated): response_format in raw HTTP body - Ollama: format= parameter - Claude: tool-use trick (forced tool call with schema as input_schema) - Mistral: native json_schema response_format - Llamafile, LM Studio: OpenAI SDK response_format - Azure OpenAI: AzureOpenAI SDK response_format - Azure serverless: response_format in raw HTTP body - TGI: response_format in raw HTTP body - VertexAI Gemini: response_mime_type + response_schema - VertexAI Claude: tool-use trick - Google AI Studio: response_mime_type + response_schema - Bedrock, Cohere: signature-only (no structured output yet) Post-hoc jsonschema.validate() retained as defence-in-depth. Tech spec added: docs/tech-specs/structured-output.md Update tests --- Makefile | 3 +- docs/tech-specs/structured-output.md | 263 +++++++++++++++++ .../test_prompt_streaming_integration.py | 8 +- .../test_base/test_schema_compatibility.py | 268 ++++++++++++++++++ trustgraph-base/trustgraph/base/__init__.py | 1 + .../trustgraph/base/llm_service.py | 13 +- .../trustgraph/base/schema_compatibility.py | 90 ++++++ .../trustgraph/base/text_completion_client.py | 12 +- .../messaging/translators/text_completion.py | 15 +- .../trustgraph/schema/services/llm.py | 4 +- .../model/text_completion/bedrock/llm.py | 10 +- .../model/text_completion/azure/llm.py | 36 ++- .../model/text_completion/azure_openai/llm.py | 23 +- .../model/text_completion/claude/llm.py | 44 ++- .../model/text_completion/cohere/llm.py | 10 +- .../model/text_completion/llamafile/llm.py | 40 ++- .../model/text_completion/lmstudio/llm.py | 40 ++- .../model/text_completion/mistral/llm.py | 42 ++- .../model/text_completion/ollama/llm.py | 25 +- .../model/text_completion/openai/llm.py | 48 +++- .../model/text_completion/openai/variants.py | 14 + .../model/text_completion/tgi/llm.py | 22 +- .../model/text_completion/vllm/llm.py | 32 ++- .../trustgraph/prompt/template/service.py | 16 +- .../trustgraph/template/prompt_manager.py | 16 +- .../text_completion/googleaistudio/llm.py | 18 +- .../model/text_completion/vertexai/llm.py | 47 ++- 27 files changed, 1089 insertions(+), 71 deletions(-) create mode 100644 docs/tech-specs/structured-output.md create mode 100644 tests/unit/test_base/test_schema_compatibility.py create mode 100644 trustgraph-base/trustgraph/base/schema_compatibility.py diff --git a/Makefile b/Makefile index 0f0f37b2..6c5f5ce5 100644 --- a/Makefile +++ b/Makefile @@ -57,7 +57,8 @@ container-bedrock container-vertexai \ container-hf container-ocr \ container-unstructured container-mcp -some-containers: container-base container-flow container-unstructured +some-containers: container-base container-flow +# container-unstructured push: ${DOCKER} push ${CONTAINER_BASE}/trustgraph-base:${VERSION} diff --git a/docs/tech-specs/structured-output.md b/docs/tech-specs/structured-output.md new file mode 100644 index 00000000..a9ef2a5f --- /dev/null +++ b/docs/tech-specs/structured-output.md @@ -0,0 +1,263 @@ +--- +layout: default +title: "Structured Output: LLM-Native JSON Schema Enforcement" +parent: "Tech Specs" +--- + +# Structured Output: LLM-Native JSON Schema Enforcement + +## Problem / Opportunity + +TrustGraph's knowledge-graph pipeline relies on LLMs to produce +structured JSON output — entity extractions, relationship triples, +topic classifications, and other schema-governed artefacts. Today, +the correct structure is requested via natural-language instructions +embedded in the prompt template: the prompt describes the expected +JSON shape, and the system parses the LLM's free-text response, +hoping it conforms. + +This approach has several weaknesses: + +1. **Fragile parsing.** LLM responses may include markdown fencing, + preamble text, trailing commentary, or minor schema violations + (missing fields, wrong types, extra keys). Every consumer must + tolerate or work around these deviations, adding defensive code + and retry logic. + +2. **Wasted tokens and latency.** A significant portion of each + prompt is spent describing the output format in prose. When the + model deviates, retries consume additional tokens and add + end-to-end latency. + +3. **Silent data-quality issues.** Malformed responses that pass + lenient parsing can inject bad data into the knowledge graph — + wrong types, truncated lists, invented field names — without + raising errors. + +4. **Untapped LLM capability.** Most modern LLMs (OpenAI, Google + Gemini, Anthropic Claude, Ollama-hosted models via llama.cpp) + support *structured output* or *guided decoding*: the caller + supplies a JSON schema and the model constrains token selection at + the logit level to guarantee schema-valid output. TrustGraph + already defines the required JSON schemas inside its prompt + definitions but does not pass them through to the LLM backend. + +### Opportunity + +By threading the existing JSON schemas from prompt definitions +through the text-completion service to each LLM backend's native +structured-output API, TrustGraph can: + +- **Guarantee valid output** on every call — no parsing heuristics, + no retries for format errors. +- **Reduce prompt size** by removing prose format instructions that + the schema makes redundant. +- **Improve data quality** in the knowledge graph by eliminating an + entire class of silent ingestion errors. +- **Simplify service code** by removing defensive JSON extraction and + validation logic from every consumer. + +## Scope + +Prompt definitions declare a `response-type` of `"text"`, `"json"`, +or `"jsonl"`. Structured output applies only to prompts that produce +machine-readable output (`"json"` and `"jsonl"`). + +JSONL presents a compatibility challenge: LLM structured-output APIs +enforce a single top-level JSON schema, but JSONL prompts ask the +model to emit one JSON object per line — a format that is not itself +valid JSON. Converting JSONL prompts to request a JSON array would +conflict with the prompt text and sacrifice truncation resilience +(partial JSONL is recoverable line-by-line; a truncated array is +broken JSON). + +This spec takes a three-phase approach: + +- **Phase 1** — plumb schemas through to LLM backends with automatic + compatibility detection; non-compliant schemas fall back to the + current free-text path. +- **Phase 2** — fix up non-compliant schemas so more prompts benefit. +- **Phase 3** — address JSONL prompts. + +--- + +## Phase 1 — Structured Output with Automatic Fallback + +### Design + +Phase 1 threads the JSON schema from the prompt definition through +the text-completion service to the LLM backend's native +structured-output API. Only prompts with `response-type: "json"` are +candidates. + +Not all existing schemas are compatible with LLM structured-output +APIs. Rather than require schema changes up front, Phase 1 includes +a **runtime compatibility check**: if a schema passes, structured +output is used; if not, the prompt falls back to the current +free-text path with post-hoc validation. This makes the feature +safe to deploy immediately. + +### Strict-Mode Schema Requirements + +LLM providers impose constraints beyond standard JSON Schema +validation. A schema is considered strict-mode compatible when: + +- Every `object` has `additionalProperties: false`. +- Every property defined in `properties` appears in `required`. + Optional fields use a nullable type (e.g. `"type": ["string", "null"]`) + instead of omitting the key from `required`. +- No `minimum`, `maximum`, `minLength`, `maxLength`, or `pattern` + constraints (unsupported by most providers' constrained decoders). +- No open-ended objects (`{"type": "object"}` without `properties`). +- A schema is present and non-null. + +### Runtime Compatibility Check + +`PromptManager` (or a shared utility) inspects each schema at load +time against the strict-mode rules above. The result is a boolean +flag per prompt: `structured_output_eligible`. + +- **Eligible** — `response_format` and `schema` are set on the + `TextCompletionRequest`; the LLM enforces the schema at generation + time. +- **Not eligible** — request is sent without schema fields; the + current free-text parsing and `jsonschema.validate()` path is used. + +This is a per-prompt decision, not a global switch. + +### Text-Completion Request Changes + +`TextCompletionRequest` gains two optional fields: + +``` +TextCompletionRequest: + system: str + prompt: str + streaming: bool + response_format: str | None # "json" or None (default) + schema: dict | None # JSON Schema object or None +``` + +When `response_format` is `"json"` and `schema` is provided, the LLM +backend MUST use its native structured-output mechanism. When either +field is absent or null, behaviour is unchanged. + +### LLM Backend Mapping + +Each backend maps `response_format` + `schema` to its provider's +native API: + +| Backend | API mechanism | +|------------|-------------------------------------------------------| +| OpenAI | `response_format={"type": "json_schema", "json_schema": {"name": "...", "schema": ...}}` | +| Claude | `tool_use` with a single tool whose `input_schema` is the target schema, or the `response_format` parameter when available | +| Gemini | `response_mime_type="application/json"` + `response_schema=...` | +| Ollama | `format="json"` + schema in the `format` field (llama.cpp guided decoding) | +| Llamafile | `response_format={"type": "json_object"}` + schema | + +Backends that do not support schema-level enforcement (e.g. older +Ollama versions) fall back to `response_format=json` without a schema +and rely on post-hoc validation. + +### Current Prompt Compatibility + +Of the nine `response-type: "json"` prompts, two are strict-mode +compatible today: + +| Prompt | Status | Issue | +|--------------------------|-----------|------------------------------------| +| `schema-selection` | Ready | — | +| `supervisor-decompose` | Ready | — | +| `plan-create` | Fixable | Optional fields not in `required` | +| `graphql-generation` | Blocked | Open-ended `variables` object; `minimum`/`maximum` on `confidence` | +| `plan-step-execute` | Blocked | Open-ended `arguments` object | +| `diagnose-structured-data` | No schema | — | +| `diagnose-xml` | No schema | — | +| `diagnose-json` | No schema | — | +| `diagnose-csv` | No schema | — | + +### What Does Not Change + +- Prompt templates and their text content. +- The `"text"` and `"jsonl"` response-type paths. +- The `TextCompletionResponse` schema. +- Post-hoc validation (retained as a defence-in-depth measure). + +--- + +## Phase 2 — Schema Remediation + +Phase 2 expands structured-output coverage by fixing schemas that +failed the Phase 1 compatibility check. + +### Fixable Schemas + +**`plan-create`** — `tool_hint` and `depends_on` are optional +(present in `properties` but absent from `required`). Fix: add both +to `required` and change their types to nullable: + +```json +"tool_hint": {"type": ["string", "null"]}, +"depends_on": { + "type": ["array", "null"], + "items": {"type": "integer"} +} +``` + +### Schemas Requiring Design Decisions + +**`graphql-generation`** — Two issues: + +- `variables` is an open-ended object (`"additionalProperties": true`) + with no fixed properties. Constrained decoding cannot handle + arbitrary keys. Options: remove `variables` from the schema and + accept it as free-form text within a wrapper, or restructure as a + JSON-encoded string field. +- `confidence` uses `"minimum": 0.0, "maximum": 1.0`. Fix: remove + the numeric bounds; accept any number and clamp in application code + if needed. + +**`plan-step-execute`** — `arguments` is an open-ended object with no +fixed properties. Same constraint as `graphql-generation.variables`. + +### Missing Schemas + +The four `diagnose-*` prompts have `response-type: "json"` but no +schema. Adding schemas for these prompts would bring them into +structured-output scope. This requires defining the expected +response shape for each diagnostic prompt. + +--- + +## Phase 3 (Future) — Structured Output for JSONL Prompts + +JSONL prompts ask the LLM to emit multiple JSON objects, one per +line. Each object is validated individually against the prompt's +schema. The current approach is tolerant of truncation and +malformed lines — useful properties for large extractions. + +### Options + +**Option A — Array wrapper.** Change the prompt text to request a +JSON array instead of JSONL. Wrap the schema as +`{"type": "array", "items": }` and use structured +output. Trade-off: loses line-by-line truncation resilience; requires +updating every JSONL prompt template. + +**Option B — Structured output per chunk.** Split the input so each +text-completion call produces a single JSON object, then aggregate. +Trade-off: more LLM calls; higher latency and cost; may not suit +prompts that extract variable-length lists from a single chunk. + +**Option C — Hybrid.** Use structured output with the array-wrapped +schema but retain the post-hoc JSONL parser as a fallback for +backends that do not support structured output or when the response +is truncated. Trade-off: two code paths to maintain. + +**Option D — Status quo.** Leave JSONL prompts on the free-text path +with post-hoc validation. Structured output for `"json"` prompts +already covers the most schema-sensitive cases; JSONL extraction is +inherently more tolerant of partial results. + +Phase 3 design will be selected after earlier phases are deployed and +real-world structured-output behaviour is observed across backends. diff --git a/tests/integration/test_prompt_streaming_integration.py b/tests/integration/test_prompt_streaming_integration.py index 84a3cdec..686b7769 100644 --- a/tests/integration/test_prompt_streaming_integration.py +++ b/tests/integration/test_prompt_streaming_integration.py @@ -34,7 +34,7 @@ class TestPromptStreaming: " of", " artificial", " intelligence", "." ] - async def streaming_text_completion_stream(system, prompt, handler, timeout=600): + async def streaming_text_completion_stream(system, prompt, handler, timeout=600, response_format=None, schema=None): """Simulate streaming text completion via text_completion_stream""" for i, chunk_text in enumerate(chunks): response = TextCompletionResponse( @@ -58,7 +58,7 @@ class TestPromptStreaming: model="test-model", ) - async def non_streaming_text_completion(system, prompt, timeout=600): + async def non_streaming_text_completion(system, prompt, timeout=600, response_format=None, schema=None): """Simulate non-streaming text completion""" full_text = "Machine learning is a field of artificial intelligence." return TextCompletionResult( @@ -230,7 +230,7 @@ class TestPromptStreaming: # Mock text completion client that raises an error text_completion_client = AsyncMock() - async def failing_stream(system, prompt, handler, timeout=600): + async def failing_stream(system, prompt, handler, timeout=600, response_format=None, schema=None): raise RuntimeError("Text completion error") text_completion_client.text_completion_stream = AsyncMock( @@ -316,7 +316,7 @@ class TestPromptStreaming: # Mock text completion that sends empty chunks text_completion_client = AsyncMock() - async def empty_streaming(system, prompt, handler, timeout=600): + async def empty_streaming(system, prompt, handler, timeout=600, response_format=None, schema=None): # Send empty chunk followed by final marker await handler(TextCompletionResponse( response="", diff --git a/tests/unit/test_base/test_schema_compatibility.py b/tests/unit/test_base/test_schema_compatibility.py new file mode 100644 index 00000000..215a71b3 --- /dev/null +++ b/tests/unit/test_base/test_schema_compatibility.py @@ -0,0 +1,268 @@ +""" +Unit tests for schema_compatibility.is_strict_mode_compatible +""" + +import pytest +from trustgraph.base.schema_compatibility import is_strict_mode_compatible + + +class TestIsStrictModeCompatible: + + def test_none_schema(self): + assert is_strict_mode_compatible(None) is False + + def test_empty_dict(self): + assert is_strict_mode_compatible({}) is True + + def test_simple_string(self): + assert is_strict_mode_compatible({"type": "string"}) is True + + def test_compliant_object(self): + schema = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"}, + }, + "required": ["name", "age"], + "additionalProperties": False, + } + assert is_strict_mode_compatible(schema) is True + + def test_missing_additional_properties(self): + schema = { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + } + assert is_strict_mode_compatible(schema) is False + + def test_additional_properties_true(self): + schema = { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"], + "additionalProperties": True, + } + assert is_strict_mode_compatible(schema) is False + + def test_property_not_in_required(self): + schema = { + "type": "object", + "properties": { + "name": {"type": "string"}, + "nickname": {"type": "string"}, + }, + "required": ["name"], + "additionalProperties": False, + } + assert is_strict_mode_compatible(schema) is False + + def test_open_ended_object_no_properties(self): + schema = { + "type": "object", + } + assert is_strict_mode_compatible(schema) is False + + def test_implicit_object_with_properties_key(self): + schema = { + "properties": { + "x": {"type": "number"}, + }, + "required": ["x"], + "additionalProperties": False, + } + assert is_strict_mode_compatible(schema) is True + + def test_nested_object_compliant(self): + schema = { + "type": "object", + "properties": { + "address": { + "type": "object", + "properties": { + "street": {"type": "string"}, + }, + "required": ["street"], + "additionalProperties": False, + }, + }, + "required": ["address"], + "additionalProperties": False, + } + assert is_strict_mode_compatible(schema) is True + + def test_nested_object_non_compliant(self): + schema = { + "type": "object", + "properties": { + "metadata": { + "type": "object", + }, + }, + "required": ["metadata"], + "additionalProperties": False, + } + assert is_strict_mode_compatible(schema) is False + + def test_array_with_compliant_items(self): + schema = { + "type": "array", + "items": { + "type": "object", + "properties": { + "id": {"type": "integer"}, + }, + "required": ["id"], + "additionalProperties": False, + }, + } + assert is_strict_mode_compatible(schema) is True + + def test_array_with_non_compliant_items(self): + schema = { + "type": "array", + "items": { + "type": "object", + "properties": {"id": {"type": "integer"}}, + }, + } + assert is_strict_mode_compatible(schema) is False + + def test_array_with_simple_items(self): + schema = { + "type": "array", + "items": {"type": "string"}, + } + assert is_strict_mode_compatible(schema) is True + + def test_oneof_all_compliant(self): + schema = { + "oneOf": [ + {"type": "string"}, + {"type": "integer"}, + ] + } + assert is_strict_mode_compatible(schema) is True + + def test_oneof_with_non_compliant_branch(self): + schema = { + "oneOf": [ + {"type": "string"}, + {"type": "object"}, + ] + } + assert is_strict_mode_compatible(schema) is False + + def test_anyof(self): + schema = { + "anyOf": [ + {"type": "string"}, + {"type": "number"}, + ] + } + assert is_strict_mode_compatible(schema) is True + + def test_allof(self): + schema = { + "allOf": [ + { + "type": "object", + "properties": {"a": {"type": "string"}}, + "required": ["a"], + "additionalProperties": False, + }, + ] + } + assert is_strict_mode_compatible(schema) is True + + def test_unsupported_minimum(self): + schema = {"type": "integer", "minimum": 0} + assert is_strict_mode_compatible(schema) is False + + def test_unsupported_maximum(self): + schema = {"type": "integer", "maximum": 100} + assert is_strict_mode_compatible(schema) is False + + def test_unsupported_pattern(self): + schema = {"type": "string", "pattern": "^[a-z]+$"} + assert is_strict_mode_compatible(schema) is False + + def test_unsupported_min_length(self): + schema = {"type": "string", "minLength": 1} + assert is_strict_mode_compatible(schema) is False + + def test_unsupported_max_length(self): + schema = {"type": "string", "maxLength": 255} + assert is_strict_mode_compatible(schema) is False + + def test_unsupported_min_items(self): + schema = {"type": "array", "items": {"type": "string"}, "minItems": 1} + assert is_strict_mode_compatible(schema) is False + + def test_unsupported_max_items(self): + schema = {"type": "array", "items": {"type": "string"}, "maxItems": 10} + assert is_strict_mode_compatible(schema) is False + + def test_unsupported_exclusive_minimum(self): + schema = {"type": "number", "exclusiveMinimum": 0} + assert is_strict_mode_compatible(schema) is False + + def test_unsupported_min_max_properties(self): + schema = { + "type": "object", + "properties": {"a": {"type": "string"}}, + "required": ["a"], + "additionalProperties": False, + "minProperties": 1, + } + assert is_strict_mode_compatible(schema) is False + + def test_unsupported_constraint_inside_nested_property(self): + schema = { + "type": "object", + "properties": { + "score": {"type": "integer", "minimum": 0, "maximum": 100}, + }, + "required": ["score"], + "additionalProperties": False, + } + assert is_strict_mode_compatible(schema) is False + + def test_nullable_property(self): + schema = { + "type": "object", + "properties": { + "name": {"type": ["string", "null"]}, + }, + "required": ["name"], + "additionalProperties": False, + } + assert is_strict_mode_compatible(schema) is True + + def test_realistic_compliant_schema(self): + schema = { + "type": "object", + "properties": { + "action": {"type": "string"}, + "services": { + "type": "array", + "items": {"type": "string"}, + }, + }, + "required": ["action", "services"], + "additionalProperties": False, + } + assert is_strict_mode_compatible(schema) is True + + def test_realistic_non_compliant_optional_field(self): + schema = { + "type": "object", + "properties": { + "action": {"type": "string"}, + "reason": {"type": "string"}, + }, + "required": ["action"], + "additionalProperties": False, + } + assert is_strict_mode_compatible(schema) is False diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index e5dd28de..941901ed 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -49,4 +49,5 @@ from . keyword_index_client import KeywordIndexClientSpec, KeywordIndexClient from . row_embeddings_query_client import RowEmbeddingsQueryClientSpec from . collection_config_handler import CollectionConfigHandler from . audit_publisher import AuditPublisher +from . schema_compatibility import is_strict_mode_compatible diff --git a/trustgraph-base/trustgraph/base/llm_service.py b/trustgraph-base/trustgraph/base/llm_service.py index 6af11670..cad1875d 100644 --- a/trustgraph-base/trustgraph/base/llm_service.py +++ b/trustgraph-base/trustgraph/base/llm_service.py @@ -126,6 +126,8 @@ class LlmService(FlowProcessor): # Check if streaming is requested and supported streaming = getattr(request, 'streaming', False) + response_format = getattr(request, 'response_format', None) + schema = getattr(request, 'schema', None) if streaming and self.supports_streaming(): @@ -136,7 +138,8 @@ class LlmService(FlowProcessor): ).time(): async for chunk in self.generate_content_stream( - request.system, request.prompt, model, temperature + request.system, request.prompt, model, temperature, + response_format=response_format, schema=schema, ): await flow("response").send( TextCompletionResponse( @@ -159,7 +162,8 @@ class LlmService(FlowProcessor): ).time(): response = await self.generate_content( - request.system, request.prompt, model, temperature + request.system, request.prompt, model, temperature, + response_format=response_format, schema=schema, ) await flow("response").send( @@ -215,7 +219,10 @@ class LlmService(FlowProcessor): """ return False - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """ Override in subclass to implement streaming. Should yield LlmChunk objects. diff --git a/trustgraph-base/trustgraph/base/schema_compatibility.py b/trustgraph-base/trustgraph/base/schema_compatibility.py new file mode 100644 index 00000000..22376f27 --- /dev/null +++ b/trustgraph-base/trustgraph/base/schema_compatibility.py @@ -0,0 +1,90 @@ + +import logging + +logger = logging.getLogger(__name__) + + +def is_strict_mode_compatible(schema): + """ + Check whether a JSON schema is compatible with LLM structured-output + strict mode. Returns True if the schema can be passed directly to + providers like OpenAI, vLLM, etc. + """ + + if schema is None: + return False + + try: + _check_node(schema) + return True + except _IncompatibleSchema as e: + logger.debug("Schema not strict-mode compatible: %s", e) + return False + + +class _IncompatibleSchema(Exception): + pass + + +def _check_node(node): + + if not isinstance(node, dict): + return + + node_type = node.get("type") + + if node_type == "object" or ( + node_type is None and "properties" in node + ): + _check_object(node) + + if node_type == "array": + items = node.get("items") + if items: + _check_node(items) + + for keyword in ("oneOf", "anyOf", "allOf"): + for child in node.get(keyword, []): + _check_node(child) + + _check_unsupported_constraints(node) + + +def _check_object(node): + + props = node.get("properties") + if props is None: + raise _IncompatibleSchema( + "object without properties (open-ended)" + ) + + if node.get("additionalProperties") is not False: + raise _IncompatibleSchema( + "object missing additionalProperties: false" + ) + + required = set(node.get("required", [])) + for key in props: + if key not in required: + raise _IncompatibleSchema( + f"property '{key}' not in required" + ) + + for value in props.values(): + _check_node(value) + + +UNSUPPORTED_KEYWORDS = { + "minimum", "maximum", "exclusiveMinimum", "exclusiveMaximum", + "minLength", "maxLength", "pattern", + "minItems", "maxItems", + "minProperties", "maxProperties", +} + + +def _check_unsupported_constraints(node): + found = UNSUPPORTED_KEYWORDS & node.keys() + if found: + raise _IncompatibleSchema( + f"unsupported constraints: {', '.join(sorted(found))}" + ) diff --git a/trustgraph-base/trustgraph/base/text_completion_client.py b/trustgraph-base/trustgraph/base/text_completion_client.py index 876d71df..87dc1874 100644 --- a/trustgraph-base/trustgraph/base/text_completion_client.py +++ b/trustgraph-base/trustgraph/base/text_completion_client.py @@ -14,11 +14,15 @@ class TextCompletionResult: class TextCompletionClient(RequestResponse): - async def text_completion(self, system, prompt, timeout=600): + async def text_completion( + self, system, prompt, timeout=600, + response_format=None, schema=None, + ): resp = await self.request( TextCompletionRequest( - system = system, prompt = prompt, streaming = False + system=system, prompt=prompt, streaming=False, + response_format=response_format, schema=schema, ), timeout=timeout ) @@ -35,6 +39,7 @@ class TextCompletionClient(RequestResponse): async def text_completion_stream( self, system, prompt, handler, timeout=600, + response_format=None, schema=None, ): """ Streaming text completion. `handler` is an async callable invoked @@ -54,7 +59,8 @@ class TextCompletionClient(RequestResponse): final = await self.request( TextCompletionRequest( - system = system, prompt = prompt, streaming = True + system=system, prompt=prompt, streaming=True, + response_format=response_format, schema=schema, ), recipient=on_chunk, timeout=timeout, diff --git a/trustgraph-base/trustgraph/messaging/translators/text_completion.py b/trustgraph-base/trustgraph/messaging/translators/text_completion.py index 62cc4afb..584c144b 100644 --- a/trustgraph-base/trustgraph/messaging/translators/text_completion.py +++ b/trustgraph-base/trustgraph/messaging/translators/text_completion.py @@ -10,14 +10,21 @@ class TextCompletionRequestTranslator(MessageTranslator): return TextCompletionRequest( system=data["system"], prompt=data["prompt"], - streaming=data.get("streaming", False) + streaming=data.get("streaming", False), + response_format=data.get("response_format"), + schema=data.get("schema"), ) - + def encode(self, obj: TextCompletionRequest) -> Dict[str, Any]: - return { + result = { "system": obj.system, - "prompt": obj.prompt + "prompt": obj.prompt, } + if obj.response_format is not None: + result["response_format"] = obj.response_format + if obj.schema is not None: + result["schema"] = obj.schema + return result class TextCompletionResponseTranslator(MessageTranslator): diff --git a/trustgraph-base/trustgraph/schema/services/llm.py b/trustgraph-base/trustgraph/schema/services/llm.py index 89c0cd54..3a319e3b 100644 --- a/trustgraph-base/trustgraph/schema/services/llm.py +++ b/trustgraph-base/trustgraph/schema/services/llm.py @@ -11,7 +11,9 @@ from ..core.primitives import Error class TextCompletionRequest: system: str = "" prompt: str = "" - streaming: bool = False # Default false for backward compatibility + streaming: bool = False + response_format: str | None = None + schema: dict | None = None @dataclass class TextCompletionResponse: diff --git a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py index 4e07b271..dee00849 100755 --- a/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py +++ b/trustgraph-bedrock/trustgraph/model/text_completion/bedrock/llm.py @@ -247,7 +247,10 @@ class Processor(LlmService): return self.model_variants[cache_key] - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -311,7 +314,10 @@ class Processor(LlmService): """Bedrock supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from Bedrock""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py index 915b6ef1..00056539 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure/llm.py @@ -55,7 +55,10 @@ class Processor(LlmService): self.max_output = max_output self.default_model = model - def build_prompt(self, system, content, temperature=None, stream=False, model=None): + def build_prompt( + self, system, content, temperature=None, stream=False, + model=None, response_format=None, schema=None, + ): # Use provided temperature or fall back to default effective_temperature = temperature if temperature is not None else self.temperature model_name = model or self.default_model @@ -79,6 +82,17 @@ class Processor(LlmService): data["stream"] = True data["stream_options"] = {"include_usage": True} + if response_format == "json" and schema is not None: + data["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + "strict": True, + }, + } + logger.debug("Structured output enabled") + body = json.dumps(data) return body @@ -109,7 +123,10 @@ class Processor(LlmService): return result - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -125,7 +142,9 @@ class Processor(LlmService): system, prompt, effective_temperature, - model=model_name + model=model_name, + response_format=response_format, + schema=schema, ) response = self.call_llm(prompt) @@ -169,7 +188,10 @@ class Processor(LlmService): """Azure serverless endpoints support streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from Azure serverless endpoint""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature @@ -178,7 +200,11 @@ class Processor(LlmService): logger.debug(f"Using temperature: {effective_temperature}") try: - body = self.build_prompt(system, prompt, effective_temperature, stream=True, model=model_name) + body = self.build_prompt( + system, prompt, effective_temperature, stream=True, + model=model_name, response_format=response_format, + schema=schema, + ) url = self.endpoint api_key = self.token diff --git a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py index def44bd4..f77fb281 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/azure_openai/llm.py @@ -62,7 +62,10 @@ class Processor(LlmService): azure_endpoint = endpoint, ) - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -76,6 +79,18 @@ class Processor(LlmService): try: + kwargs = {} + if response_format == "json" and schema is not None: + kwargs["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + "strict": True, + }, + } + logger.debug("Structured output enabled") + resp = self.openai.chat.completions.create( model=model_name, messages=[ @@ -92,6 +107,7 @@ class Processor(LlmService): temperature=effective_temperature, max_completion_tokens=self.max_output, top_p=1, + **kwargs, ) inputtokens = resp.usage.prompt_tokens @@ -129,7 +145,10 @@ class Processor(LlmService): """Azure OpenAI supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """ Stream content generation from Azure OpenAI. Yields LlmChunk objects with is_final=True on the last chunk. diff --git a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py index 2e7573d0..d4fcddcd 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/claude/llm.py @@ -48,7 +48,10 @@ class Processor(LlmService): logger.info("Claude LLM service initialized") - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -60,11 +63,27 @@ class Processor(LlmService): try: + kwargs = {} + use_tool_extract = False + + if response_format == "json" and schema is not None: + kwargs["tools"] = [{ + "name": "structured_response", + "description": "Return the structured response", + "input_schema": schema, + }] + kwargs["tool_choice"] = { + "type": "tool", + "name": "structured_response", + } + use_tool_extract = True + logger.debug("Structured output enabled (tool-use)") + response = message = self.claude.messages.create( model=model_name, max_tokens=self.max_output, temperature=effective_temperature, - system = system, + system=system, messages=[ { "role": "user", @@ -75,10 +94,22 @@ class Processor(LlmService): } ] } - ] + ], + **kwargs, ) - resp = response.content[0].text + if use_tool_extract: + import json + tool_block = next( + (b for b in response.content if b.type == "tool_use"), + None, + ) + if tool_block: + resp = json.dumps(tool_block.input) + else: + resp = response.content[0].text + else: + resp = response.content[0].text inputtokens = response.usage.input_tokens outputtokens = response.usage.output_tokens logger.debug(f"LLM response: {resp}") @@ -110,7 +141,10 @@ class Processor(LlmService): """Claude/Anthropic supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from Claude""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature diff --git a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py index 4190cb98..10c72d85 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/cohere/llm.py @@ -46,7 +46,10 @@ class Processor(LlmService): logger.info("Cohere LLM service initialized") - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -104,7 +107,10 @@ class Processor(LlmService): """Cohere supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from Cohere""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature diff --git a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py index 276727b5..1061d602 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/llamafile/llm.py @@ -50,7 +50,10 @@ class Processor(LlmService): logger.info("Llamafile LLM service initialized") - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -64,6 +67,18 @@ class Processor(LlmService): try: + if response_format == "json" and schema is not None: + fmt = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + }, + } + logger.debug("Structured output enabled") + else: + fmt = {"type": "text"} + resp = self.openai.chat.completions.create( model=model_name, messages=[ @@ -74,9 +89,7 @@ class Processor(LlmService): top_p=1, frequency_penalty=0, presence_penalty=0, - response_format={ - "type": "text" - } + response_format=fmt, ) inputtokens = resp.usage.prompt_tokens @@ -106,7 +119,10 @@ class Processor(LlmService): """LlamaFile supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from LlamaFile""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature @@ -117,6 +133,18 @@ class Processor(LlmService): prompt = system + "\n\n" + prompt try: + if response_format == "json" and schema is not None: + fmt = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + }, + } + logger.debug("Structured output enabled (streaming)") + else: + fmt = {"type": "text"} + response = self.openai.chat.completions.create( model=model_name, messages=[{"role": "user", "content": prompt}], @@ -125,7 +153,7 @@ class Processor(LlmService): top_p=1, frequency_penalty=0, presence_penalty=0, - response_format={"type": "text"}, + response_format=fmt, stream=True, stream_options={"include_usage": True} ) diff --git a/trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py b/trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py index b057f58d..3930e3e0 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/lmstudio/llm.py @@ -50,7 +50,10 @@ class Processor(LlmService): logger.info("LMStudio LLM service initialized") - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -66,6 +69,18 @@ class Processor(LlmService): logger.debug(f"Prompt: {prompt}") + if response_format == "json" and schema is not None: + fmt = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + }, + } + logger.debug("Structured output enabled") + else: + fmt = {"type": "text"} + resp = self.openai.chat.completions.create( model=model_name, messages=[ @@ -76,9 +91,7 @@ class Processor(LlmService): top_p=1, frequency_penalty=0, presence_penalty=0, - response_format={ - "type": "text" - } + response_format=fmt, ) logger.debug(f"Full response: {resp}") @@ -110,7 +123,10 @@ class Processor(LlmService): """LM Studio supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from LM Studio""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature @@ -121,6 +137,18 @@ class Processor(LlmService): prompt = system + "\n\n" + prompt try: + if response_format == "json" and schema is not None: + fmt = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + }, + } + logger.debug("Structured output enabled (streaming)") + else: + fmt = {"type": "text"} + response = self.openai.chat.completions.create( model=model_name, messages=[{"role": "user", "content": prompt}], @@ -129,7 +157,7 @@ class Processor(LlmService): top_p=1, frequency_penalty=0, presence_penalty=0, - response_format={"type": "text"}, + response_format=fmt, stream=True, stream_options={"include_usage": True} ) diff --git a/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py b/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py index e53f6f6e..75335f62 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/mistral/llm.py @@ -48,7 +48,10 @@ class Processor(LlmService): logger.info("Mistral LLM service initialized") - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -62,6 +65,19 @@ class Processor(LlmService): try: + if response_format == "json" and schema is not None: + fmt = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + "strict": True, + }, + } + logger.debug("Structured output enabled") + else: + fmt = {"type": "text"} + resp = self.mistral.chat.complete( model=model_name, messages=[ @@ -80,9 +96,7 @@ class Processor(LlmService): top_p=1, frequency_penalty=0, presence_penalty=0, - response_format={ - "type": "text" - } + response_format=fmt, ) inputtokens = resp.usage.prompt_tokens @@ -120,7 +134,10 @@ class Processor(LlmService): """Mistral supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from Mistral""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature @@ -131,6 +148,19 @@ class Processor(LlmService): prompt = system + "\n\n" + prompt try: + if response_format == "json" and schema is not None: + fmt = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + "strict": True, + }, + } + logger.debug("Structured output enabled (streaming)") + else: + fmt = {"type": "text"} + stream = self.mistral.chat.stream( model=model_name, messages=[ @@ -149,7 +179,7 @@ class Processor(LlmService): top_p=1, frequency_penalty=0, presence_penalty=0, - response_format={"type": "text"} + response_format=fmt, ) total_input_tokens = 0 diff --git a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py index 2e537fde..68148959 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/ollama/llm.py @@ -62,7 +62,10 @@ class Processor(LlmService): else: logger.warning(f"Failed to check Ollama model '{model_name}': {e}") - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -79,7 +82,12 @@ class Processor(LlmService): try: - response = await self.llm.generate(model_name, prompt, options={'temperature': effective_temperature}) + kwargs = {} + if response_format == "json" and schema is not None: + kwargs["format"] = schema + logger.debug("Structured output enabled") + + response = await self.llm.generate(model_name, prompt, options={'temperature': effective_temperature}, **kwargs) response_text = response['response'] logger.debug("Sending response...") @@ -108,7 +116,10 @@ class Processor(LlmService): """Ollama supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from Ollama""" model_name = model or self.default_model @@ -123,11 +134,17 @@ class Processor(LlmService): prompt = system + "\n\n" + prompt try: + kwargs = {} + if response_format == "json" and schema is not None: + kwargs["format"] = schema + logger.debug("Structured output enabled (streaming)") + stream = await self.llm.generate( model_name, prompt, options={'temperature': effective_temperature}, - stream=True + stream=True, + **kwargs, ) total_input_tokens = 0 diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py index 57958bc0..e50129fc 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/llm.py @@ -82,7 +82,10 @@ class Processor(LlmService): return self.variant.extract_content(message) return message.content - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature @@ -96,6 +99,25 @@ class Processor(LlmService): api_kwargs = self._build_kwargs(model_name, effective_temperature) + if response_format == "json" and schema is not None: + is_top_level_array = schema.get("type") == "array" + if is_top_level_array and not self.variant.supports_top_level_array(): + logger.debug( + "Variant %s does not support top-level array " + "schemas, falling back to free-text", + self.variant.name, + ) + else: + api_kwargs["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + "strict": True, + }, + } + logger.debug("Structured output enabled") + messages = [ { "role": "user", @@ -160,7 +182,10 @@ class Processor(LlmService): """OpenAI supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """ Stream content generation from OpenAI. Yields LlmChunk objects with is_final=True on the last chunk. @@ -176,6 +201,25 @@ class Processor(LlmService): try: api_kwargs = self._build_kwargs(model_name, effective_temperature) + if response_format == "json" and schema is not None: + is_top_level_array = schema.get("type") == "array" + if is_top_level_array and not self.variant.supports_top_level_array(): + logger.debug( + "Variant %s does not support top-level array " + "schemas, falling back to free-text (streaming)", + self.variant.name, + ) + else: + api_kwargs["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + "strict": True, + }, + } + logger.debug("Structured output enabled (streaming)") + messages = [ { "role": "user", diff --git a/trustgraph-flow/trustgraph/model/text_completion/openai/variants.py b/trustgraph-flow/trustgraph/model/text_completion/openai/variants.py index 87de725d..f05903d5 100644 --- a/trustgraph-flow/trustgraph/model/text_completion/openai/variants.py +++ b/trustgraph-flow/trustgraph/model/text_completion/openai/variants.py @@ -62,6 +62,10 @@ class Variant: """Extract thinking content from a streaming delta.""" return getattr(delta, "reasoning_content", None) + def supports_top_level_array(self): + """Whether this provider accepts a top-level array JSON schema.""" + return True + def create_completion(self, client, model, messages, **kwargs): """Call the completions API. Override for non-standard SDKs.""" return client.chat.completions.create( @@ -84,6 +88,9 @@ class OpenAIVariant(Variant): token_param = "max_completion_tokens" temperature_with_thinking = False + def supports_top_level_array(self): + return False + def thinking_kwargs(self, effort): return {"reasoning_effort": effort} @@ -195,6 +202,12 @@ class LlamaVariant(Variant): return re.sub(r".*?", "", content, flags=re.DOTALL).strip() +class VllmVariant(LlamaVariant): + """vLLM via OpenAI-compatible API. Supports full structured output.""" + + name = "vllm" + + VARIANTS = { "openai": OpenAIVariant, "deepseek": DeepSeekVariant, @@ -203,6 +216,7 @@ VARIANTS = { "dashscope": DashScopeVariant, "glm": GlmVariant, "llama": LlamaVariant, + "vllm": VllmVariant, } DEFAULT_VARIANT = "openai" diff --git a/trustgraph-flow/trustgraph/model/text_completion/tgi/llm.py b/trustgraph-flow/trustgraph/model/text_completion/tgi/llm.py index 5caeb9be..15d6f24f 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/tgi/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/tgi/llm.py @@ -51,7 +51,10 @@ class Processor(LlmService): logger.info(f"Using TGI service at {base_url}") logger.info("TGI LLM service initialized") - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -79,7 +82,17 @@ class Processor(LlmService): ], "max_tokens": self.max_output, "temperature": effective_temperature, - } + } + + if response_format == "json" and schema is not None: + request["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + }, + } + logger.debug("Structured output enabled") try: @@ -125,7 +138,10 @@ class Processor(LlmService): """TGI supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from TGI""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature diff --git a/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py b/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py index 7570fa40..ca904e6b 100755 --- a/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py +++ b/trustgraph-flow/trustgraph/model/text_completion/vllm/llm.py @@ -52,7 +52,10 @@ class Processor(LlmService): logger.info(f"Using vLLM service at {base_url}") logger.info("vLLM LLM service initialized") - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -71,7 +74,17 @@ class Processor(LlmService): "prompt": system + "\n\n" + prompt, "max_tokens": self.max_output, "temperature": effective_temperature, - } + } + + if response_format == "json" and schema is not None: + request["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + }, + } + logger.debug("Structured output enabled") try: @@ -127,7 +140,10 @@ class Processor(LlmService): """vLLM supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from vLLM""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature @@ -148,6 +164,16 @@ class Processor(LlmService): "stream_options": {"include_usage": True}, } + if response_format == "json" and schema is not None: + request["response_format"] = { + "type": "json_schema", + "json_schema": { + "name": "response", + "schema": schema, + }, + } + logger.debug("Structured output enabled (streaming)") + try: url = f"{self.base_url.rstrip('/')}/completions" diff --git a/trustgraph-flow/trustgraph/prompt/template/service.py b/trustgraph-flow/trustgraph/prompt/template/service.py index 5da329d3..4733f74e 100755 --- a/trustgraph-flow/trustgraph/prompt/template/service.py +++ b/trustgraph-flow/trustgraph/prompt/template/service.py @@ -155,7 +155,10 @@ class Processor(FlowProcessor): # For streaming, we need to intercept LLM responses # and forward them as they arrive - async def llm_streaming(system, prompt): + async def llm_streaming( + system, prompt, + response_format=None, schema=None, + ): logger.debug(f"System prompt: {system}") logger.debug(f"User prompt: {prompt}") @@ -179,6 +182,8 @@ class Processor(FlowProcessor): system=system, prompt=prompt, handler=forward_chunks, timeout=600, + response_format=response_format, + schema=schema, ) # Return empty string since we already sent all chunks @@ -195,14 +200,19 @@ class Processor(FlowProcessor): # Non-streaming path (original behavior) usage = {} - async def llm(system, prompt): + async def llm( + system, prompt, + response_format=None, schema=None, + ): logger.debug(f"System prompt: {system}") logger.debug(f"User prompt: {prompt}") try: result = await flow("text-completion-request").text_completion( - system = system, prompt = prompt, + system=system, prompt=prompt, + response_format=response_format, + schema=schema, ) usage["in_token"] = result.in_token usage["out_token"] = result.out_token diff --git a/trustgraph-flow/trustgraph/template/prompt_manager.py b/trustgraph-flow/trustgraph/template/prompt_manager.py index 976d3695..eed4d214 100644 --- a/trustgraph-flow/trustgraph/template/prompt_manager.py +++ b/trustgraph-flow/trustgraph/template/prompt_manager.py @@ -5,6 +5,8 @@ from jsonschema import validate import re import logging +from trustgraph.base import is_strict_mode_compatible + # Module logger logger = logging.getLogger(__name__) @@ -145,12 +147,24 @@ class PromptManager: terms = self.terms | self.prompts[id].terms | input resp_type = self.prompts[id].response_type + schema = self.prompts[id].schema prompt = { "system": self.system_template.render(terms), - "prompt": self.render(id, input) + "prompt": self.render(id, input), } + use_structured = ( + resp_type == "json" + and schema is not None + and is_strict_mode_compatible(schema) + ) + + if use_structured: + logger.debug("Using structured output for prompt '%s'", id) + prompt["response_format"] = "json" + prompt["schema"] = schema + resp = await llm(**prompt) if resp_type == "text": diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/googleaistudio/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/googleaistudio/llm.py index b01ff410..faebe42a 100644 --- a/trustgraph-vertexai/trustgraph/model/text_completion/googleaistudio/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/googleaistudio/llm.py @@ -109,7 +109,10 @@ class Processor(LlmService): return self.generation_configs[cache_key] - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -123,6 +126,14 @@ class Processor(LlmService): # Set system instruction per request (can't be cached) generation_config.system_instruction = system + if response_format == "json" and schema is not None: + generation_config.response_mime_type = "application/json" + generation_config.response_schema = schema + logger.debug("Structured output enabled (Gemini)") + else: + generation_config.response_mime_type = "text/plain" + generation_config.response_schema = None + try: response = self.client.models.generate_content( @@ -174,7 +185,10 @@ class Processor(LlmService): """Google AI Studio supports streaming""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """Stream content generation from Google AI Studio""" model_name = model or self.default_model effective_temperature = temperature if temperature is not None else self.temperature diff --git a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py index d7a7dd2a..94104fa6 100755 --- a/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py +++ b/trustgraph-vertexai/trustgraph/model/text_completion/vertexai/llm.py @@ -166,7 +166,10 @@ class Processor(LlmService): return self.generation_configs[cache_key] - async def generate_content(self, system, prompt, model=None, temperature=None): + async def generate_content( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): # Use provided model or fall back to default model_name = model or self.default_model @@ -182,6 +185,22 @@ class Processor(LlmService): logger.debug(f"Sending request to Anthropic model '{model_name}'...") client = self._get_anthropic_client() + kwargs = {} + use_tool_extract = False + + if response_format == "json" and schema is not None: + kwargs["tools"] = [{ + "name": "structured_response", + "description": "Return the structured response", + "input_schema": schema, + }] + kwargs["tool_choice"] = { + "type": "tool", + "name": "structured_response", + } + use_tool_extract = True + logger.debug("Structured output enabled (tool-use)") + response = client.messages.create( model=model_name, system=system, @@ -190,10 +209,21 @@ class Processor(LlmService): temperature=effective_temperature, top_p=self.api_params['top_p'], top_k=self.api_params['top_k'], + **kwargs, ) + if use_tool_extract: + import json + tool_block = next( + (b for b in response.content if b.type == "tool_use"), + None, + ) + text = json.dumps(tool_block.input) if tool_block else response.content[0].text + else: + text = response.content[0].text + resp = LlmResult( - text=response.content[0].text, + text=text, in_token=response.usage.input_tokens, out_token=response.usage.output_tokens, model=model_name @@ -206,6 +236,14 @@ class Processor(LlmService): # Set system instruction per request (can't be cached) generation_config.system_instruction = system + if response_format == "json" and schema is not None: + generation_config.response_mime_type = "application/json" + generation_config.response_schema = schema + logger.debug("Structured output enabled (Gemini)") + else: + generation_config.response_mime_type = "text/plain" + generation_config.response_schema = None + response = self.client.models.generate_content( model=model_name, config=generation_config, @@ -248,7 +286,10 @@ class Processor(LlmService): """VertexAI supports streaming for both Gemini and Claude models""" return True - async def generate_content_stream(self, system, prompt, model=None, temperature=None): + async def generate_content_stream( + self, system, prompt, model=None, temperature=None, + response_format=None, schema=None, + ): """ Stream content generation from VertexAI (Gemini or Claude). Yields LlmChunk objects with is_final=True on the last chunk.