Restore AI-assisted code generation in Code Step dialog

Feature was present in TUI but missing after GUI conversion. Now includes:
- Split layout: Code editor (left) | AI Assist panel (right)
- Provider selector dropdown for AI calls
- Prompt editor with smart default template showing available variables
- "Generate Code" button that calls AI in background thread
- Response feedback area showing success/error status
- Automatic markdown fence stripping from AI responses
- Available variables computed from tool arguments + previous step outputs

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
rob 2026-01-14 14:20:58 -04:00
parent bee74061b1
commit 2438aec831
2 changed files with 217 additions and 33 deletions

View File

@ -3,12 +3,12 @@
from PySide6.QtWidgets import (
QDialog, QVBoxLayout, QFormLayout, QLineEdit,
QComboBox, QPushButton, QHBoxLayout, QLabel,
QPlainTextEdit
QPlainTextEdit, QSplitter, QGroupBox, QTextEdit
)
from PySide6.QtCore import Qt
from PySide6.QtCore import Qt, QThread, Signal
from ...tool import PromptStep, CodeStep
from ...providers import load_providers
from ...providers import load_providers, call_provider
from ...profiles import list_profiles
@ -136,63 +136,142 @@ class PromptStepDialog(QDialog):
)
class CodeStepDialog(QDialog):
"""Dialog for editing code steps."""
class AIGenerateWorker(QThread):
"""Background worker for AI code generation."""
finished = Signal(str)
error = Signal(str)
def __init__(self, parent, step: CodeStep = None):
def __init__(self, provider: str, prompt: str):
super().__init__()
self.provider = provider
self.prompt = prompt
def run(self):
try:
result = call_provider(self.provider, self.prompt)
if result.success:
self.finished.emit(result.text)
else:
self.error.emit(result.error or "Unknown error")
except Exception as e:
self.error.emit(str(e))
class CodeStepDialog(QDialog):
"""Dialog for editing code steps with AI assist."""
def __init__(self, parent, step: CodeStep = None, available_vars: list = None):
super().__init__(parent)
self.setWindowTitle("Edit Code Step" if step else "Add Code Step")
self.setMinimumSize(600, 500)
self.setMinimumSize(900, 600)
self._step = step
self._available_vars = available_vars or ["input"]
self._worker = None
self._setup_ui()
if step:
self._load_step(step)
def _setup_ui(self):
"""Set up the UI."""
"""Set up the UI with code editor and AI assist panel."""
layout = QVBoxLayout(self)
layout.setSpacing(16)
layout.setSpacing(12)
# Form
# Top: Output variable
form = QFormLayout()
form.setSpacing(12)
# Output variable
form.setSpacing(8)
self.output_input = QLineEdit()
self.output_input.setPlaceholderText("result")
self.output_input.setText("result")
form.addRow("Output variable:", self.output_input)
layout.addLayout(form)
# Code editor
code_label = QLabel("Code:")
layout.addWidget(code_label)
# Available variables display
vars_text = ", ".join(self._available_vars)
vars_label = QLabel(f"Available variables: {vars_text}")
vars_label.setStyleSheet("color: #718096; font-size: 11px;")
layout.addWidget(vars_label)
# Main splitter: Code editor | AI Assist
splitter = QSplitter(Qt.Horizontal)
# Left: Code editor
code_group = QGroupBox("Code")
code_layout = QVBoxLayout(code_group)
self.code_input = QPlainTextEdit()
self.code_input.setPlaceholderText(
"# Python code here\n"
"# Access input with: input_text\n"
"# Access args with their variable names\n"
"# Access previous step outputs with their variable names\n\n"
"result = input_text.upper()"
"# Access input with: input\n"
"# Access args with their variable names\n\n"
"result = input.upper()"
)
font = self.code_input.font()
font.setFamily("Courier New, Consolas, monospace")
font.setFamily("Consolas, Monaco, monospace")
self.code_input.setFont(font)
layout.addWidget(self.code_input, 1)
code_layout.addWidget(self.code_input)
# Help text
help_text = QLabel(
"The code will be executed with variables in scope. "
"The last expression or the output variable will be captured."
splitter.addWidget(code_group)
# Right: AI Assist panel
ai_group = QGroupBox("AI Assisted Code Generation")
ai_layout = QVBoxLayout(ai_group)
ai_layout.setSpacing(8)
# Provider selector
provider_layout = QHBoxLayout()
provider_layout.addWidget(QLabel("Provider:"))
self.ai_provider_combo = QComboBox()
providers = load_providers()
for provider in sorted(providers, key=lambda p: p.name):
self.ai_provider_combo.addItem(provider.name)
# Add common defaults if not present
for default in ["claude", "gpt", "mock"]:
if self.ai_provider_combo.findText(default) < 0:
self.ai_provider_combo.addItem(default)
self.ai_provider_combo.setMinimumWidth(150)
provider_layout.addWidget(self.ai_provider_combo)
provider_layout.addStretch()
ai_layout.addLayout(provider_layout)
# Prompt editor
prompt_label = QLabel("Describe what you want the code to do:")
ai_layout.addWidget(prompt_label)
self.ai_prompt_input = QPlainTextEdit()
self.ai_prompt_input.setPlaceholderText(
"Example: Parse the input as JSON and extract the 'name' field"
)
help_text.setStyleSheet("color: #718096; font-size: 11px;")
help_text.setWordWrap(True)
layout.addWidget(help_text)
# Set default prompt template
self._set_default_prompt()
ai_layout.addWidget(self.ai_prompt_input, 2)
# Buttons
# Generate button
btn_layout = QHBoxLayout()
self.btn_generate = QPushButton("Generate Code")
self.btn_generate.clicked.connect(self._generate_code)
self.btn_generate.setMinimumHeight(32)
btn_layout.addStretch()
btn_layout.addWidget(self.btn_generate)
btn_layout.addStretch()
ai_layout.addLayout(btn_layout)
# Output/feedback area
feedback_label = QLabel("AI Response:")
ai_layout.addWidget(feedback_label)
self.ai_output = QTextEdit()
self.ai_output.setReadOnly(True)
self.ai_output.setPlaceholderText("AI response will appear here...")
self.ai_output.setMaximumHeight(100)
ai_layout.addWidget(self.ai_output, 1)
splitter.addWidget(ai_group)
splitter.setSizes([450, 450])
layout.addWidget(splitter, 1)
# Bottom: OK/Cancel buttons
buttons = QHBoxLayout()
buttons.addStretch()
@ -207,6 +286,81 @@ class CodeStepDialog(QDialog):
layout.addLayout(buttons)
def _set_default_prompt(self):
"""Set the default AI prompt template."""
vars_formatted = ', '.join(f'"{{{v}}}"' for v in self._available_vars)
default_prompt = f"""Write inline Python code (NOT a function definition) according to my instruction.
The code runs directly with variable substitution. Assign any available variables to a local variable first using triple quotes for multi-line content safety.
Example:
my_var = \"\"\"{{input}}\"\"\"
result = my_var.upper()
INSTRUCTION: [Describe what you want the code to do]
CURRENT CODE:
{{code}}
AVAILABLE VARIABLES: {vars_formatted}
IMPORTANT: Return ONLY executable inline Python code. No function definitions, no markdown fencing, no explanations - just the code."""
self.ai_prompt_input.setPlainText(default_prompt)
def _generate_code(self):
"""Generate code using AI."""
prompt_template = self.ai_prompt_input.toPlainText().strip()
if not prompt_template:
self.ai_output.setHtml("<span style='color: #e53e3e;'>Please enter a prompt</span>")
return
# Replace {code} placeholder with current code
current_code = self.code_input.toPlainText().strip() or "# No code yet"
prompt = prompt_template.replace("{code}", current_code)
provider = self.ai_provider_combo.currentText()
# Disable button and show loading state
self.btn_generate.setEnabled(False)
self.btn_generate.setText("Generating...")
self.ai_output.setHtml(f"<span style='color: #718096;'>Calling {provider}...</span>")
# Start worker thread
self._worker = AIGenerateWorker(provider, prompt)
self._worker.finished.connect(self._on_generate_finished)
self._worker.error.connect(self._on_generate_error)
self._worker.start()
def _on_generate_finished(self, result: str):
"""Handle successful AI generation."""
self.btn_generate.setEnabled(True)
self.btn_generate.setText("Generate Code")
# Clean up the result - strip markdown code fences if present
code = result.strip()
if code.startswith("```python"):
code = code[9:]
elif code.startswith("```"):
code = code[3:]
if code.endswith("```"):
code = code[:-3]
code = code.strip()
# Update code editor
self.code_input.setPlainText(code)
# Show success message
self.ai_output.setHtml(
f"<span style='color: #38a169;'>Code generated successfully!</span><br>"
f"<span style='color: #718096;'>Response length: {len(result)} chars</span>"
)
def _on_generate_error(self, error: str):
"""Handle AI generation error."""
self.btn_generate.setEnabled(True)
self.btn_generate.setText("Generate Code")
self.ai_output.setHtml(f"<span style='color: #e53e3e;'>Error: {error}</span>")
def _load_step(self, step: CodeStep):
"""Load step data into form."""
self.output_input.setText(step.output_var)

View File

@ -261,6 +261,34 @@ class ToolBuilderPage(QWidget):
del self._tool.arguments[idx]
self._refresh_arguments()
def _get_available_vars(self, up_to_step: int = -1) -> list:
"""Get available variables for a step.
Args:
up_to_step: Include variables from steps before this index.
-1 means include all existing steps (for new steps).
"""
vars_list = ["input"]
# Add argument variables
if self._tool and self._tool.arguments:
for arg in self._tool.arguments:
if arg.variable and arg.variable not in vars_list:
vars_list.append(arg.variable)
# Add output variables from previous steps
if self._tool and self._tool.steps:
steps_to_check = self._tool.steps if up_to_step < 0 else self._tool.steps[:up_to_step]
for step in steps_to_check:
if hasattr(step, 'output_var') and step.output_var:
# Handle comma-separated output vars
for var in step.output_var.split(','):
var = var.strip()
if var and var not in vars_list:
vars_list.append(var)
return vars_list
def _add_prompt_step(self):
"""Add a prompt step."""
from ..dialogs.step_dialog import PromptStepDialog
@ -275,7 +303,8 @@ class ToolBuilderPage(QWidget):
def _add_code_step(self):
"""Add a code step."""
from ..dialogs.step_dialog import CodeStepDialog
dialog = CodeStepDialog(self)
available_vars = self._get_available_vars()
dialog = CodeStepDialog(self, available_vars=available_vars)
if dialog.exec():
step = dialog.get_step()
if not self._tool:
@ -297,7 +326,8 @@ class ToolBuilderPage(QWidget):
dialog = PromptStepDialog(self, step)
elif isinstance(step, CodeStep):
from ..dialogs.step_dialog import CodeStepDialog
dialog = CodeStepDialog(self, step)
available_vars = self._get_available_vars(up_to_step=idx)
dialog = CodeStepDialog(self, step, available_vars=available_vars)
else:
return