Implement structured output for prompt steps (Phase 1)
Add schema validation and retry logic for AI responses:
- Add output_schema, max_retries, plain_text fields to PromptStep
- Add DEFAULT_OUTPUT_SCHEMA with output/reasoning fields
- Implement append_schema_instructions() for prompt augmentation
- Implement validate_schema() using jsonschema library
- Add retry logic with error feedback on validation failure
- Enhance substitute_variables() for nested field access ({var.field})
- Add jsonschema to dependencies
- Update GUI dialog with plain_text checkbox and max_retries spinner
- Add comprehensive tests for structured output and nested access
Tools using plain_text=True bypass validation (current behavior).
Tools without plain_text get structured JSON output with schema enforcement.
Migration Phase 1 complete. Next: Apply plain_text=True to all 72 tools.
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
parent
f980fe05f1
commit
8581775002
|
|
@ -35,6 +35,7 @@ dependencies = [
|
|||
"PySide6>=6.5",
|
||||
"NodeGraphQt>=0.6.0",
|
||||
"setuptools", # Required for distutils compatibility (Python 3.12+)
|
||||
"jsonschema>=4.0", # JSON schema validation for structured output
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ from PySide6.QtWidgets import (
|
|||
QDialog, QVBoxLayout, QFormLayout, QLineEdit,
|
||||
QComboBox, QPushButton, QHBoxLayout, QLabel,
|
||||
QPlainTextEdit, QSplitter, QGroupBox, QTextEdit, QMessageBox,
|
||||
QCheckBox
|
||||
QCheckBox, QSpinBox
|
||||
)
|
||||
from PySide6.QtCore import Qt, QThread, Signal
|
||||
|
||||
|
|
@ -73,6 +73,31 @@ class PromptStepDialog(QDialog):
|
|||
self.strip_fences_check = QCheckBox("Strip markdown code fences from output")
|
||||
form.addRow("", self.strip_fences_check)
|
||||
|
||||
# Structured output options
|
||||
form.addRow(QLabel("")) # Spacer
|
||||
structured_label = QLabel("<b>Structured Output</b>")
|
||||
form.addRow(structured_label)
|
||||
|
||||
# Plain text checkbox (bypass structured output)
|
||||
self.plain_text_check = QCheckBox("Plain text mode (bypass JSON validation)")
|
||||
self.plain_text_check.setToolTip(
|
||||
"When checked, output is returned as-is without JSON parsing.\n"
|
||||
"When unchecked, output must be valid JSON matching a schema."
|
||||
)
|
||||
self.plain_text_check.stateChanged.connect(self._on_plain_text_changed)
|
||||
form.addRow("", self.plain_text_check)
|
||||
|
||||
# Max retries (only visible when structured output is enabled)
|
||||
self.retries_spin = QSpinBox()
|
||||
self.retries_spin.setRange(0, 5)
|
||||
self.retries_spin.setValue(1)
|
||||
self.retries_spin.setToolTip(
|
||||
"Number of retries if JSON validation fails.\n"
|
||||
"The AI will receive error feedback and try again."
|
||||
)
|
||||
self.retries_label = QLabel("Max retries:")
|
||||
form.addRow(self.retries_label, self.retries_spin)
|
||||
|
||||
layout.addLayout(form)
|
||||
|
||||
# Prompt text
|
||||
|
|
@ -103,6 +128,12 @@ class PromptStepDialog(QDialog):
|
|||
|
||||
layout.addLayout(buttons)
|
||||
|
||||
def _on_plain_text_changed(self, state):
|
||||
"""Show/hide structured output options based on plain_text checkbox."""
|
||||
is_plain = bool(state)
|
||||
self.retries_label.setVisible(not is_plain)
|
||||
self.retries_spin.setVisible(not is_plain)
|
||||
|
||||
def _load_step(self, step: PromptStep):
|
||||
"""Load step data into form."""
|
||||
# Load name
|
||||
|
|
@ -125,6 +156,11 @@ class PromptStepDialog(QDialog):
|
|||
self.prompt_input.setPlainText(step.prompt)
|
||||
self.strip_fences_check.setChecked(step.strip_fences)
|
||||
|
||||
# Structured output fields
|
||||
self.plain_text_check.setChecked(step.plain_text)
|
||||
self.retries_spin.setValue(step.max_retries)
|
||||
self._on_plain_text_changed(step.plain_text)
|
||||
|
||||
def _validate_and_accept(self):
|
||||
"""Validate and accept."""
|
||||
prompt = self.prompt_input.toPlainText().strip()
|
||||
|
|
@ -153,7 +189,9 @@ class PromptStepDialog(QDialog):
|
|||
output_var=self.output_input.text().strip(),
|
||||
profile=profile,
|
||||
name=name,
|
||||
strip_fences=self.strip_fences_check.isChecked()
|
||||
strip_fences=self.strip_fences_check.isChecked(),
|
||||
plain_text=self.plain_text_check.isChecked(),
|
||||
max_retries=self.retries_spin.value()
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
"""Tool execution engine."""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
|
@ -16,6 +17,80 @@ from .profiles import load_profile
|
|||
# Maximum recursion depth for nested tool calls
|
||||
MAX_TOOL_DEPTH = 10
|
||||
|
||||
# Default schema for structured output
|
||||
# Separates AI reasoning from actual output
|
||||
DEFAULT_OUTPUT_SCHEMA = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"output": {
|
||||
"type": "string",
|
||||
"description": "The actual response content"
|
||||
},
|
||||
"reasoning": {
|
||||
"type": "string",
|
||||
"description": "Internal thoughts, analysis, or step-by-step reasoning"
|
||||
}
|
||||
},
|
||||
"required": ["output"]
|
||||
}
|
||||
|
||||
|
||||
class SchemaValidationError(Exception):
|
||||
"""Raised when response doesn't match expected schema."""
|
||||
pass
|
||||
|
||||
|
||||
def validate_schema(data: dict, schema: dict) -> None:
|
||||
"""
|
||||
Validate data against a JSON schema.
|
||||
|
||||
Args:
|
||||
data: Parsed JSON data
|
||||
schema: JSON schema to validate against
|
||||
|
||||
Raises:
|
||||
SchemaValidationError: If validation fails
|
||||
"""
|
||||
try:
|
||||
from jsonschema import validate, ValidationError
|
||||
validate(instance=data, schema=schema)
|
||||
except ValidationError as e:
|
||||
# Extract path info for clearer error
|
||||
path = ".".join(str(p) for p in e.absolute_path) if e.absolute_path else "root"
|
||||
raise SchemaValidationError(f"Field '{path}': {e.message}")
|
||||
except ImportError:
|
||||
# jsonschema not installed - skip validation with warning
|
||||
print("Warning: jsonschema not installed, skipping schema validation", file=sys.stderr)
|
||||
|
||||
|
||||
def append_schema_instructions(prompt: str, schema: dict) -> str:
|
||||
"""
|
||||
Append schema instructions to a prompt.
|
||||
|
||||
Args:
|
||||
prompt: Original prompt text
|
||||
schema: JSON schema the response must match
|
||||
|
||||
Returns:
|
||||
Augmented prompt with schema instructions
|
||||
"""
|
||||
schema_json = json.dumps(schema, indent=2)
|
||||
|
||||
return f"""{prompt}
|
||||
|
||||
---
|
||||
RESPONSE FORMAT:
|
||||
You must respond with ONLY valid JSON matching this exact schema:
|
||||
{schema_json}
|
||||
|
||||
CRITICAL RULES:
|
||||
- Output ONLY the JSON object, no other text before or after
|
||||
- Put any thinking, reasoning, or analysis in the "reasoning" field
|
||||
- Put your actual answer/response in the "output" field
|
||||
- Do not include markdown code fences
|
||||
- Do not include any preamble like "Here is my response:"
|
||||
"""
|
||||
|
||||
|
||||
def check_system_dependencies(tool: Tool) -> list:
|
||||
"""
|
||||
|
|
@ -114,14 +189,52 @@ def check_dependencies(tool: Tool, checked: set = None) -> list[str]:
|
|||
return missing
|
||||
|
||||
|
||||
def _get_nested_value(obj, path_parts: list):
|
||||
"""
|
||||
Navigate nested dict/JSON to get a value.
|
||||
|
||||
Args:
|
||||
obj: Starting object (dict or JSON string)
|
||||
path_parts: List of keys to navigate
|
||||
|
||||
Returns:
|
||||
Value at path, or None if not found
|
||||
"""
|
||||
value = obj
|
||||
|
||||
for part in path_parts:
|
||||
if isinstance(value, dict):
|
||||
value = value.get(part)
|
||||
elif isinstance(value, str):
|
||||
# Try parsing as JSON
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
if isinstance(parsed, dict):
|
||||
value = parsed.get(part)
|
||||
else:
|
||||
return None
|
||||
except json.JSONDecodeError:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
|
||||
if value is None:
|
||||
return None
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def substitute_variables(template: str, variables: dict, warn_non_scalar: bool = False) -> str:
|
||||
"""
|
||||
Substitute {variable} placeholders in a template.
|
||||
|
||||
Supports escaping: use {{ for literal { and }} for literal }
|
||||
Supports:
|
||||
- Simple: {varname}
|
||||
- Nested: {varname.field} or {varname.field.subfield}
|
||||
- Settings: {settings.key}
|
||||
- Escaping: {{ for literal { and }} for literal }
|
||||
|
||||
Also supports {settings.key} syntax for accessing top-level scalar
|
||||
values from the settings dict.
|
||||
For nested access, if the variable is a JSON string, it will be parsed.
|
||||
|
||||
Args:
|
||||
template: String with {var} placeholders
|
||||
|
|
@ -137,6 +250,8 @@ def substitute_variables(template: str, variables: dict, warn_non_scalar: bool =
|
|||
'Use {braces}'
|
||||
>>> substitute_variables("Backend: {settings.backend}", {"settings": {"backend": "piper"}})
|
||||
'Backend: piper'
|
||||
>>> substitute_variables("Answer: {result.output}", {"result": '{"output": "Hello"}'})
|
||||
'Answer: Hello'
|
||||
"""
|
||||
import re
|
||||
|
||||
|
|
@ -147,34 +262,60 @@ def substitute_variables(template: str, variables: dict, warn_non_scalar: bool =
|
|||
# First, replace escaped braces with placeholders
|
||||
result = template.replace("{{", ESCAPE_OPEN).replace("}}", ESCAPE_CLOSE)
|
||||
|
||||
# Handle settings.key syntax (top-level scalar values only)
|
||||
settings = variables.get("settings", {})
|
||||
if settings and isinstance(settings, dict):
|
||||
def replace_settings(match):
|
||||
key = match.group(1)
|
||||
# Pattern for {var} or {var.field.subfield}
|
||||
pattern = r'\{([a-zA-Z_][a-zA-Z0-9_]*(?:\.[a-zA-Z_][a-zA-Z0-9_]*)*)\}'
|
||||
|
||||
def replace_var(match):
|
||||
path = match.group(1)
|
||||
parts = path.split('.')
|
||||
|
||||
# Get base variable
|
||||
base_name = parts[0]
|
||||
|
||||
# Check if variable exists in dict (distinguish None value from missing key)
|
||||
if base_name not in variables:
|
||||
return match.group(0) # Leave unchanged for missing variables
|
||||
|
||||
base_value = variables[base_name]
|
||||
|
||||
# If no nested parts, simple substitution
|
||||
if len(parts) == 1:
|
||||
if base_name == "settings":
|
||||
# Settings dict - leave unchanged (handled specially)
|
||||
return match.group(0)
|
||||
# Replace with empty string for None, str() for other values
|
||||
return "" if base_value is None else str(base_value)
|
||||
|
||||
# Handle nested access
|
||||
if base_name == "settings":
|
||||
# Settings dict - only allow scalar values in templates
|
||||
settings = base_value if isinstance(base_value, dict) else {}
|
||||
key = parts[1] # Only support one level: settings.key
|
||||
if key in settings:
|
||||
value = settings[key]
|
||||
if isinstance(value, (str, int, float, bool)):
|
||||
return str(value) if value is not None else ""
|
||||
else:
|
||||
# Non-scalar value - leave placeholder, warn if requested
|
||||
if warn_non_scalar:
|
||||
print(
|
||||
f"Warning: {{settings.{key}}} is not a scalar value, "
|
||||
f"access it in code steps instead",
|
||||
file=sys.stderr
|
||||
)
|
||||
return match.group(0) # Return unchanged
|
||||
return match.group(0) # Key not found, leave unchanged
|
||||
return match.group(0)
|
||||
return match.group(0)
|
||||
|
||||
result = re.sub(r'\{settings\.([^}]+)\}', replace_settings, result)
|
||||
# Navigate nested path for other variables
|
||||
value = _get_nested_value(base_value, parts[1:])
|
||||
if value is None:
|
||||
return match.group(0) # Leave unchanged if not found
|
||||
|
||||
# Now do regular variable substitution (skip 'settings' since we handled it)
|
||||
for name, value in variables.items():
|
||||
if name == "settings":
|
||||
continue # Already handled above
|
||||
# Use 'is not None' to preserve falsey values like 0 and False
|
||||
result = result.replace(f"{{{name}}}", "" if value is None else str(value))
|
||||
# Convert to string for template
|
||||
if isinstance(value, (dict, list)):
|
||||
return json.dumps(value)
|
||||
return str(value)
|
||||
|
||||
result = re.sub(pattern, replace_var, result)
|
||||
|
||||
# Finally, restore escaped braces as single braces
|
||||
result = result.replace(ESCAPE_OPEN, "{").replace(ESCAPE_CLOSE, "}")
|
||||
|
|
@ -198,7 +339,12 @@ def execute_prompt_step(
|
|||
|
||||
Returns:
|
||||
Tuple of (output_value, success)
|
||||
|
||||
If plain_text=True, returns raw provider output.
|
||||
Otherwise, enforces structured JSON output with schema validation.
|
||||
"""
|
||||
import re
|
||||
|
||||
# Build prompt with variable substitution
|
||||
prompt = substitute_variables(step.prompt, variables, warn_non_scalar=verbose)
|
||||
|
||||
|
|
@ -209,25 +355,101 @@ def execute_prompt_step(
|
|||
# Prepend system prompt to user prompt
|
||||
prompt = f"{profile.system_prompt}\n\n---\n\n{prompt}"
|
||||
|
||||
# Call provider (support variable substitution in provider name, e.g. {settings.provider})
|
||||
# Determine provider
|
||||
provider = provider_override or substitute_variables(step.provider, variables, warn_non_scalar=verbose)
|
||||
|
||||
if provider.lower() == "mock":
|
||||
result = mock_provider(prompt)
|
||||
else:
|
||||
result = call_provider(provider, prompt)
|
||||
# Plain text mode - bypass structured output
|
||||
if step.plain_text:
|
||||
if provider.lower() == "mock":
|
||||
result = mock_provider(prompt)
|
||||
else:
|
||||
result = call_provider(provider, prompt)
|
||||
|
||||
if not result.success:
|
||||
print(f"Error in prompt step: {result.error}", file=sys.stderr)
|
||||
return "", False
|
||||
if not result.success:
|
||||
print(f"Error in prompt step: {result.error}", file=sys.stderr)
|
||||
return "", False
|
||||
|
||||
text = result.text
|
||||
if step.strip_fences:
|
||||
import re
|
||||
text = re.sub(r'^```\w*\n', '', text.strip())
|
||||
text = re.sub(r'\n```\s*$', '', text)
|
||||
text = result.text
|
||||
if step.strip_fences:
|
||||
text = re.sub(r'^```\w*\n', '', text.strip())
|
||||
text = re.sub(r'\n```\s*$', '', text)
|
||||
|
||||
return text, True
|
||||
return text, True
|
||||
|
||||
# Structured output mode - enforce JSON with schema validation
|
||||
schema = step.output_schema or DEFAULT_OUTPUT_SCHEMA
|
||||
|
||||
# Augment prompt with schema instructions
|
||||
augmented_prompt = append_schema_instructions(prompt, schema)
|
||||
|
||||
max_attempts = step.max_retries + 1
|
||||
last_response = None
|
||||
last_error = None
|
||||
|
||||
for attempt in range(max_attempts):
|
||||
current_prompt = augmented_prompt
|
||||
|
||||
if attempt > 0 and last_response:
|
||||
# Retry with error feedback
|
||||
current_prompt = f"""{augmented_prompt}
|
||||
|
||||
Your previous response was:
|
||||
{last_response[:1000]}{"..." if len(last_response) > 1000 else ""}
|
||||
|
||||
This failed validation: {last_error}
|
||||
Please try again with valid JSON matching the schema exactly."""
|
||||
|
||||
if verbose:
|
||||
if attempt > 0:
|
||||
print(f"[verbose] Retry {attempt}/{step.max_retries} after validation failure", file=sys.stderr)
|
||||
|
||||
# Call provider
|
||||
if provider.lower() == "mock":
|
||||
result = mock_provider(current_prompt)
|
||||
else:
|
||||
result = call_provider(provider, current_prompt)
|
||||
|
||||
if not result.success:
|
||||
print(f"Error in prompt step: {result.error}", file=sys.stderr)
|
||||
return "", False
|
||||
|
||||
last_response = result.text
|
||||
|
||||
# Strip markdown code fences if present (common model behavior)
|
||||
text = result.text.strip()
|
||||
if text.startswith("```"):
|
||||
text = re.sub(r'^```\w*\n', '', text)
|
||||
text = re.sub(r'\n```\s*$', '', text)
|
||||
|
||||
# Try to parse as JSON
|
||||
try:
|
||||
parsed = json.loads(text)
|
||||
except json.JSONDecodeError as e:
|
||||
last_error = f"Invalid JSON: {e}"
|
||||
if attempt == max_attempts - 1:
|
||||
print(f"Prompt step failed after {step.max_retries} retry(s): {last_error}", file=sys.stderr)
|
||||
if verbose:
|
||||
print(f"[verbose] Last response: {last_response[:500]}...", file=sys.stderr)
|
||||
return "", False
|
||||
continue
|
||||
|
||||
# Validate against schema
|
||||
try:
|
||||
validate_schema(parsed, schema)
|
||||
except SchemaValidationError as e:
|
||||
last_error = str(e)
|
||||
if attempt == max_attempts - 1:
|
||||
print(f"Prompt step failed after {step.max_retries} retry(s): {last_error}", file=sys.stderr)
|
||||
if verbose:
|
||||
print(f"[verbose] Parsed response: {json.dumps(parsed)[:500]}...", file=sys.stderr)
|
||||
return "", False
|
||||
continue
|
||||
|
||||
# Success - return normalized JSON string
|
||||
return json.dumps(parsed), True
|
||||
|
||||
# Should not reach here, but just in case
|
||||
return "", False
|
||||
|
||||
|
||||
def execute_code_step(
|
||||
|
|
|
|||
|
|
@ -101,6 +101,10 @@ class PromptStep:
|
|||
profile: Optional[str] = None # Optional AI persona profile name
|
||||
name: Optional[str] = None # Optional display name for the step
|
||||
strip_fences: bool = False # Strip markdown code fences from output
|
||||
# Structured output fields
|
||||
output_schema: Optional[dict] = None # JSON schema for validation (None = use default)
|
||||
max_retries: int = 1 # Retry count on validation failure
|
||||
plain_text: bool = False # Bypass structured output enforcement
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
d = {
|
||||
|
|
@ -117,6 +121,12 @@ class PromptStep:
|
|||
d["name"] = self.name
|
||||
if self.strip_fences:
|
||||
d["strip_fences"] = self.strip_fences
|
||||
if self.output_schema:
|
||||
d["output_schema"] = self.output_schema
|
||||
if self.max_retries != 1:
|
||||
d["max_retries"] = self.max_retries
|
||||
if self.plain_text:
|
||||
d["plain_text"] = self.plain_text
|
||||
return d
|
||||
|
||||
@classmethod
|
||||
|
|
@ -128,7 +138,10 @@ class PromptStep:
|
|||
prompt_file=data.get("prompt_file"),
|
||||
profile=data.get("profile"),
|
||||
name=data.get("name"),
|
||||
strip_fences=data.get("strip_fences", False)
|
||||
strip_fences=data.get("strip_fences", False),
|
||||
output_schema=data.get("output_schema"),
|
||||
max_retries=data.get("max_retries", 1),
|
||||
plain_text=data.get("plain_text", False)
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -184,7 +184,7 @@ class TestRunCommand:
|
|||
tool = Tool(
|
||||
name="summarize",
|
||||
steps=[
|
||||
PromptStep(prompt="Summarize: {input}", provider="mock", output_var="summary")
|
||||
PromptStep(prompt="Summarize: {input}", provider="mock", output_var="summary", plain_text=True)
|
||||
],
|
||||
output="{summary}"
|
||||
)
|
||||
|
|
@ -223,7 +223,7 @@ class TestTestCommand:
|
|||
tool = Tool(
|
||||
name="test-me",
|
||||
steps=[
|
||||
PromptStep(prompt="Test: {input}", provider="claude", output_var="result")
|
||||
PromptStep(prompt="Test: {input}", provider="claude", output_var="result", plain_text=True)
|
||||
],
|
||||
output="{result}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -108,6 +108,82 @@ Line 3: {var1} again"""
|
|||
assert result == "Enabled: True"
|
||||
|
||||
|
||||
class TestNestedFieldAccess:
|
||||
"""Tests for nested field access in variable substitution."""
|
||||
|
||||
def test_nested_dict_access(self):
|
||||
"""Access nested dict field via dot notation."""
|
||||
result = substitute_variables(
|
||||
"Answer: {result.output}",
|
||||
{"result": {"output": "Hello", "reasoning": "User said hi"}}
|
||||
)
|
||||
assert result == "Answer: Hello"
|
||||
|
||||
def test_nested_json_string_access(self):
|
||||
"""Access field in JSON string via dot notation."""
|
||||
result = substitute_variables(
|
||||
"Answer: {result.output}",
|
||||
{"result": '{"output": "Hello", "reasoning": "User said hi"}'}
|
||||
)
|
||||
assert result == "Answer: Hello"
|
||||
|
||||
def test_nested_deep_access(self):
|
||||
"""Access deeply nested fields."""
|
||||
result = substitute_variables(
|
||||
"Value: {data.nested.deep}",
|
||||
{"data": {"nested": {"deep": "found"}}}
|
||||
)
|
||||
assert result == "Value: found"
|
||||
|
||||
def test_nested_missing_field_unchanged(self):
|
||||
"""Missing nested field leaves placeholder unchanged."""
|
||||
result = substitute_variables(
|
||||
"Value: {result.missing}",
|
||||
{"result": {"output": "Hello"}}
|
||||
)
|
||||
assert result == "Value: {result.missing}"
|
||||
|
||||
def test_nested_on_non_dict_unchanged(self):
|
||||
"""Nested access on non-dict value leaves placeholder unchanged."""
|
||||
result = substitute_variables(
|
||||
"Value: {result.field}",
|
||||
{"result": "just a string"}
|
||||
)
|
||||
assert result == "Value: {result.field}"
|
||||
|
||||
def test_nested_invalid_json_unchanged(self):
|
||||
"""Invalid JSON string leaves placeholder unchanged."""
|
||||
result = substitute_variables(
|
||||
"Value: {result.field}",
|
||||
{"result": "not valid json {"}
|
||||
)
|
||||
assert result == "Value: {result.field}"
|
||||
|
||||
def test_nested_array_in_json(self):
|
||||
"""Nested array values should be JSON-serialized."""
|
||||
result = substitute_variables(
|
||||
"Items: {result.items}",
|
||||
{"result": '{"items": [1, 2, 3]}'}
|
||||
)
|
||||
assert result == "Items: [1, 2, 3]"
|
||||
|
||||
def test_nested_object_in_json(self):
|
||||
"""Nested object values should be JSON-serialized."""
|
||||
result = substitute_variables(
|
||||
"Data: {result.data}",
|
||||
{"result": '{"data": {"key": "value"}}'}
|
||||
)
|
||||
assert result == 'Data: {"key": "value"}'
|
||||
|
||||
def test_settings_access_unchanged(self):
|
||||
"""Settings access should use existing special handling."""
|
||||
result = substitute_variables(
|
||||
"Backend: {settings.backend}",
|
||||
{"settings": {"backend": "piper"}}
|
||||
)
|
||||
assert result == "Backend: piper"
|
||||
|
||||
|
||||
class TestExecutePromptStep:
|
||||
"""Tests for prompt step execution."""
|
||||
|
||||
|
|
@ -118,10 +194,12 @@ class TestExecutePromptStep:
|
|||
success=True
|
||||
)
|
||||
|
||||
# Use plain_text=True to test legacy behavior (no schema enforcement)
|
||||
step = PromptStep(
|
||||
prompt="Summarize: {input}",
|
||||
provider="claude",
|
||||
output_var="summary"
|
||||
output_var="summary",
|
||||
plain_text=True
|
||||
)
|
||||
variables = {"input": "Some text to summarize"}
|
||||
|
||||
|
|
@ -142,7 +220,7 @@ class TestExecutePromptStep:
|
|||
error="Provider error"
|
||||
)
|
||||
|
||||
step = PromptStep(prompt="Test", provider="claude", output_var="out")
|
||||
step = PromptStep(prompt="Test", provider="claude", output_var="out", plain_text=True)
|
||||
output, success = execute_prompt_step(step, {"input": ""})
|
||||
|
||||
assert success is False
|
||||
|
|
@ -152,7 +230,7 @@ class TestExecutePromptStep:
|
|||
def test_mock_provider_used(self, mock_mock):
|
||||
mock_mock.return_value = ProviderResult(text="mock response", success=True)
|
||||
|
||||
step = PromptStep(prompt="Test", provider="mock", output_var="out")
|
||||
step = PromptStep(prompt="Test", provider="mock", output_var="out", plain_text=True)
|
||||
output, success = execute_prompt_step(step, {"input": ""})
|
||||
|
||||
assert success is True
|
||||
|
|
@ -162,13 +240,180 @@ class TestExecutePromptStep:
|
|||
def test_provider_override(self, mock_call):
|
||||
mock_call.return_value = ProviderResult(text="response", success=True)
|
||||
|
||||
step = PromptStep(prompt="Test", provider="claude", output_var="out")
|
||||
step = PromptStep(prompt="Test", provider="claude", output_var="out", plain_text=True)
|
||||
execute_prompt_step(step, {"input": ""}, provider_override="gpt4")
|
||||
|
||||
# Should use override, not step's provider
|
||||
assert mock_call.call_args[0][0] == "gpt4"
|
||||
|
||||
|
||||
class TestStructuredOutput:
|
||||
"""Tests for structured output enforcement."""
|
||||
|
||||
@patch('cmdforge.runner.call_provider')
|
||||
def test_structured_output_valid_json(self, mock_call):
|
||||
"""Valid JSON matching default schema should succeed."""
|
||||
mock_call.return_value = ProviderResult(
|
||||
text='{"output": "Hello, world!", "reasoning": "Greeting requested"}',
|
||||
success=True
|
||||
)
|
||||
|
||||
step = PromptStep(
|
||||
prompt="Say hello",
|
||||
provider="claude",
|
||||
output_var="result"
|
||||
# plain_text defaults to False - structured output enabled
|
||||
)
|
||||
|
||||
output, success = execute_prompt_step(step, {"input": ""})
|
||||
|
||||
assert success is True
|
||||
# Output should be normalized JSON
|
||||
import json
|
||||
parsed = json.loads(output)
|
||||
assert parsed["output"] == "Hello, world!"
|
||||
assert parsed["reasoning"] == "Greeting requested"
|
||||
|
||||
@patch('cmdforge.runner.call_provider')
|
||||
def test_structured_output_strips_markdown_fences(self, mock_call):
|
||||
"""Markdown code fences should be stripped before parsing."""
|
||||
mock_call.return_value = ProviderResult(
|
||||
text='```json\n{"output": "Hello"}\n```',
|
||||
success=True
|
||||
)
|
||||
|
||||
step = PromptStep(
|
||||
prompt="Say hello",
|
||||
provider="claude",
|
||||
output_var="result"
|
||||
)
|
||||
|
||||
output, success = execute_prompt_step(step, {"input": ""})
|
||||
|
||||
assert success is True
|
||||
import json
|
||||
parsed = json.loads(output)
|
||||
assert parsed["output"] == "Hello"
|
||||
|
||||
@patch('cmdforge.runner.call_provider')
|
||||
def test_structured_output_retry_on_invalid_json(self, mock_call):
|
||||
"""Invalid JSON should trigger retry with error feedback."""
|
||||
# First call returns invalid JSON, second returns valid
|
||||
mock_call.side_effect = [
|
||||
ProviderResult(text="Not JSON at all", success=True),
|
||||
ProviderResult(text='{"output": "Fixed!"}', success=True),
|
||||
]
|
||||
|
||||
step = PromptStep(
|
||||
prompt="Test",
|
||||
provider="claude",
|
||||
output_var="result",
|
||||
max_retries=1
|
||||
)
|
||||
|
||||
output, success = execute_prompt_step(step, {"input": ""})
|
||||
|
||||
assert success is True
|
||||
assert mock_call.call_count == 2
|
||||
# Second call should include error feedback
|
||||
second_call_prompt = mock_call.call_args_list[1][0][1]
|
||||
assert "failed validation" in second_call_prompt.lower()
|
||||
|
||||
@patch('cmdforge.runner.call_provider')
|
||||
def test_structured_output_fails_after_retries(self, mock_call):
|
||||
"""Should fail after exhausting retries."""
|
||||
mock_call.return_value = ProviderResult(
|
||||
text="Still not JSON",
|
||||
success=True
|
||||
)
|
||||
|
||||
step = PromptStep(
|
||||
prompt="Test",
|
||||
provider="claude",
|
||||
output_var="result",
|
||||
max_retries=1 # 1 retry = 2 total attempts
|
||||
)
|
||||
|
||||
output, success = execute_prompt_step(step, {"input": ""})
|
||||
|
||||
assert success is False
|
||||
assert mock_call.call_count == 2
|
||||
|
||||
@patch('cmdforge.runner.call_provider')
|
||||
def test_structured_output_custom_schema(self, mock_call):
|
||||
"""Custom schema should be validated."""
|
||||
mock_call.return_value = ProviderResult(
|
||||
text='{"score": 0.8, "reasoning": "Good match"}',
|
||||
success=True
|
||||
)
|
||||
|
||||
custom_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"score": {"type": "number", "minimum": 0, "maximum": 1},
|
||||
"reasoning": {"type": "string"}
|
||||
},
|
||||
"required": ["score"]
|
||||
}
|
||||
|
||||
step = PromptStep(
|
||||
prompt="Score this",
|
||||
provider="claude",
|
||||
output_var="result",
|
||||
output_schema=custom_schema
|
||||
)
|
||||
|
||||
output, success = execute_prompt_step(step, {"input": ""})
|
||||
|
||||
assert success is True
|
||||
import json
|
||||
parsed = json.loads(output)
|
||||
assert parsed["score"] == 0.8
|
||||
|
||||
@patch('cmdforge.runner.call_provider')
|
||||
def test_structured_output_schema_violation_triggers_retry(self, mock_call):
|
||||
"""Schema violation should trigger retry."""
|
||||
# First call missing required 'output' field
|
||||
mock_call.side_effect = [
|
||||
ProviderResult(text='{"reasoning": "thinking..."}', success=True),
|
||||
ProviderResult(text='{"output": "Fixed!", "reasoning": "done"}', success=True),
|
||||
]
|
||||
|
||||
step = PromptStep(
|
||||
prompt="Test",
|
||||
provider="claude",
|
||||
output_var="result",
|
||||
max_retries=1
|
||||
)
|
||||
|
||||
output, success = execute_prompt_step(step, {"input": ""})
|
||||
|
||||
assert success is True
|
||||
assert mock_call.call_count == 2
|
||||
|
||||
@patch('cmdforge.runner.call_provider')
|
||||
def test_plain_text_bypasses_validation(self, mock_call):
|
||||
"""plain_text=True should skip all JSON validation."""
|
||||
mock_call.return_value = ProviderResult(
|
||||
text="Just plain text, no JSON here!",
|
||||
success=True
|
||||
)
|
||||
|
||||
step = PromptStep(
|
||||
prompt="Write prose",
|
||||
provider="claude",
|
||||
output_var="result",
|
||||
plain_text=True
|
||||
)
|
||||
|
||||
output, success = execute_prompt_step(step, {"input": ""})
|
||||
|
||||
assert success is True
|
||||
assert output == "Just plain text, no JSON here!"
|
||||
# Should only call provider once (no retries for plain text)
|
||||
assert mock_call.call_count == 1
|
||||
|
||||
|
||||
class TestExecuteCodeStep:
|
||||
"""Tests for code step execution."""
|
||||
|
||||
|
|
@ -387,7 +632,8 @@ class TestRunTool:
|
|||
PromptStep(
|
||||
prompt="Summarize: {input}",
|
||||
provider="claude",
|
||||
output_var="summary"
|
||||
output_var="summary",
|
||||
plain_text=True # Use plain text mode for legacy behavior
|
||||
)
|
||||
],
|
||||
output="{summary}"
|
||||
|
|
@ -423,7 +669,7 @@ class TestRunTool:
|
|||
name="multi-step",
|
||||
steps=[
|
||||
CodeStep(code="preprocessed = input.strip().upper()", output_var="preprocessed"),
|
||||
PromptStep(prompt="Process: {preprocessed}", provider="claude", output_var="response"),
|
||||
PromptStep(prompt="Process: {preprocessed}", provider="claude", output_var="response", plain_text=True),
|
||||
CodeStep(code="final = f'Result: {response}'", output_var="final")
|
||||
],
|
||||
output="{final}"
|
||||
|
|
@ -439,7 +685,7 @@ class TestRunTool:
|
|||
tool = Tool(
|
||||
name="test",
|
||||
steps=[
|
||||
PromptStep(prompt="Test", provider="claude", output_var="response")
|
||||
PromptStep(prompt="Test", provider="claude", output_var="response", plain_text=True)
|
||||
],
|
||||
output="{response}"
|
||||
)
|
||||
|
|
@ -458,7 +704,7 @@ class TestRunTool:
|
|||
tool = Tool(
|
||||
name="test",
|
||||
steps=[
|
||||
PromptStep(prompt="Test", provider="claude", output_var="response")
|
||||
PromptStep(prompt="Test", provider="claude", output_var="response", plain_text=True)
|
||||
],
|
||||
output="{response}"
|
||||
)
|
||||
|
|
@ -475,7 +721,7 @@ class TestRunTool:
|
|||
tool = Tool(
|
||||
name="test",
|
||||
steps=[
|
||||
PromptStep(prompt="Test", provider="claude", output_var="response")
|
||||
PromptStep(prompt="Test", provider="claude", output_var="response", plain_text=True)
|
||||
],
|
||||
output="{response}"
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in New Issue