mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-13 23:32:11 +02:00
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=<schema> 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
This commit is contained in:
cybermaggedon
GitHub
parent
f106ae2103
commit
9136526863
27 changed files with 1089 additions and 71 deletions
3
Makefile
3
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}
|
||||
|
|
|
|||
263
docs/tech-specs/structured-output.md
Normal file
263
docs/tech-specs/structured-output.md
Normal file
|
|
@ -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": <existing-schema>}` 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.
|
||||
|
|
@ -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="",
|
||||
|
|
|
|||
268
tests/unit/test_base/test_schema_compatibility.py
Normal file
268
tests/unit/test_base/test_schema_compatibility.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
90
trustgraph-base/trustgraph/base/schema_compatibility.py
Normal file
90
trustgraph-base/trustgraph/base/schema_compatibility.py
Normal file
|
|
@ -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))}"
|
||||
)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"<think>.*?</think>", "", 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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue