feat(ai): enhance chat functionality and orchestration capabilities
- Introduce context optimization to manage chat history size and prevent unbounded growth. - Implement complexity assessment for design prompts to determine orchestration needs. - Add orchestrator functionality for parallel design generation, improving performance and responsiveness. - Update chat message parsing to include explicit status handling for orchestrator steps. - Refactor design generator to support orchestration and streamline context building.
This commit is contained in:
Kayshen-X
parent
80dcda0bf1
commit
cbd211df4b
|
|
@ -19,6 +19,7 @@ import {
|
|||
extractAndApplyDesign,
|
||||
extractAndApplyDesignModification
|
||||
} from '@/services/ai/design-generator'
|
||||
import { trimChatHistory } from '@/services/ai/context-optimizer'
|
||||
import type { ChatMessage as ChatMessageType } from '@/services/ai/ai-types'
|
||||
import type { AIProviderType } from '@/types/agent-settings'
|
||||
import ClaudeLogo from '@/components/icons/claude-logo'
|
||||
|
|
@ -227,12 +228,14 @@ function useChatHandlers() {
|
|||
} else {
|
||||
// --- CHAT MODE ---
|
||||
chatHistory.push({ role: 'user', content: fullUserMessage })
|
||||
// Trim history to prevent unbounded context growth
|
||||
const trimmedHistory = trimChatHistory(chatHistory)
|
||||
// Resolve which provider the currently selected model belongs to
|
||||
const currentProvider = useAIStore.getState().modelGroups.find((g) =>
|
||||
g.models.some((m) => m.value === model),
|
||||
)?.provider
|
||||
let chatThinking = false
|
||||
for await (const chunk of streamChat(CHAT_SYSTEM_PROMPT, chatHistory, model, undefined, currentProvider)) {
|
||||
for await (const chunk of streamChat(CHAT_SYSTEM_PROMPT, trimmedHistory, model, undefined, currentProvider)) {
|
||||
if (chunk.type === 'thinking') {
|
||||
// Show a brief thinking indicator so the UI isn't stuck on empty
|
||||
if (!chatThinking && !accumulated) {
|
||||
|
|
|
|||
|
|
@ -49,16 +49,22 @@ function stripToolCallXml(text: string): string {
|
|||
export interface ParsedStep {
|
||||
title: string
|
||||
content: string
|
||||
/** Explicit status from orchestrator steps (undefined for normal steps) */
|
||||
status?: 'pending' | 'streaming' | 'done' | 'error'
|
||||
}
|
||||
|
||||
export function parseStepBlocks(text: string, isStreaming?: boolean): ParsedStep[] {
|
||||
const stepRegex = /<step(?:[^>]*title="([^"]+)")?[^>]*>([\s\S]*?)<\/step>/gi
|
||||
const stepRegex = /<step([^>]*)>([\s\S]*?)<\/step>/gi
|
||||
const parsed: ParsedStep[] = []
|
||||
let match: RegExpExecArray | null
|
||||
|
||||
while ((match = stepRegex.exec(text)) !== null) {
|
||||
const attrs = match[1]
|
||||
const titleMatch = attrs.match(/title="([^"]+)"/)
|
||||
const statusMatch = attrs.match(/status="([^"]+)"/)
|
||||
parsed.push({
|
||||
title: (match[1] ?? 'Processing').trim() || 'Processing',
|
||||
title: (titleMatch?.[1] ?? 'Processing').trim() || 'Processing',
|
||||
status: (statusMatch?.[1] as ParsedStep['status']) ?? undefined,
|
||||
content: (match[2] ?? '').trim(),
|
||||
})
|
||||
}
|
||||
|
|
@ -68,9 +74,11 @@ export function parseStepBlocks(text: string, isStreaming?: boolean): ParsedStep
|
|||
if (isStreaming && lastOpen > lastClose) {
|
||||
const partial = text.slice(lastOpen)
|
||||
const titleMatch = partial.match(/title="([^"]+)"/i)
|
||||
const statusMatch = partial.match(/status="([^"]+)"/i)
|
||||
const contentStart = partial.indexOf('>')
|
||||
parsed.push({
|
||||
title: (titleMatch?.[1] ?? 'Design').trim() || 'Design',
|
||||
status: (statusMatch?.[1] as ParsedStep['status']) ?? undefined,
|
||||
content:
|
||||
contentStart >= 0
|
||||
? partial
|
||||
|
|
@ -149,7 +157,17 @@ export function buildPipelineProgress(
|
|||
return steps.map((s) => ({ label: s.title, done: true, active: false }))
|
||||
}
|
||||
|
||||
// Map each step to done/active/pending based on completed JSON blocks.
|
||||
// If steps have explicit status (orchestrator mode), use that directly
|
||||
const hasExplicitStatus = steps.some((s) => s.status !== undefined)
|
||||
if (hasExplicitStatus) {
|
||||
return steps.map((s) => ({
|
||||
label: s.title,
|
||||
done: s.status === 'done',
|
||||
active: s.status === 'streaming',
|
||||
}))
|
||||
}
|
||||
|
||||
// Fallback: Map each step to done/active/pending based on completed JSON blocks.
|
||||
// Step[i] is done when jsonBlockCount > i.
|
||||
// The step at jsonBlockCount is active (currently being generated).
|
||||
return steps.map((s, index) => {
|
||||
|
|
|
|||
|
|
@ -7,6 +7,12 @@ const DEFAULT_STREAM_NO_TEXT_TIMEOUT_MS = 75_000
|
|||
interface StreamChatOptions {
|
||||
hardTimeoutMs?: number
|
||||
noTextTimeoutMs?: number
|
||||
/**
|
||||
* Whether thinking events should reset the no-text timeout.
|
||||
* Default: true (backward compatible). Set to false for fast calls
|
||||
* where thinking should NOT prevent the no-text timeout from firing.
|
||||
*/
|
||||
thinkingResetsTimeout?: boolean
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -22,6 +28,7 @@ export async function* streamChat(
|
|||
): AsyncGenerator<AIStreamChunk> {
|
||||
const hardTimeoutMs = Math.max(10_000, options?.hardTimeoutMs ?? DEFAULT_STREAM_HARD_TIMEOUT_MS)
|
||||
const noTextTimeoutMs = Math.max(10_000, options?.noTextTimeoutMs ?? DEFAULT_STREAM_NO_TEXT_TIMEOUT_MS)
|
||||
const thinkingResetsTimeout = options?.thinkingResetsTimeout ?? true
|
||||
|
||||
const controller = new AbortController()
|
||||
let abortReason: 'hard_timeout' | 'no_text_timeout' | null = null
|
||||
|
|
@ -113,8 +120,11 @@ export async function* streamChat(
|
|||
continue
|
||||
}
|
||||
|
||||
// Any non-empty content (text or thinking) counts as activity
|
||||
if ((chunk.type === 'text' || chunk.type === 'thinking') && chunk.content.trim().length > 0) {
|
||||
// Any non-empty text counts as activity; thinking only resets
|
||||
// the timeout when thinkingResetsTimeout is true (default).
|
||||
if (chunk.type === 'text' && chunk.content.trim().length > 0) {
|
||||
resetActivityTimeout()
|
||||
} else if (chunk.type === 'thinking' && chunk.content.trim().length > 0 && thinkingResetsTimeout) {
|
||||
resetActivityTimeout()
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -27,3 +27,55 @@ export interface AIStreamChunk {
|
|||
type: 'text' | 'thinking' | 'done' | 'error' | 'ping'
|
||||
content: string
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Orchestrator types — used for parallel sub-agent design generation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/** A sub-task produced by the orchestrator (lightweight — only spatial info) */
|
||||
export interface SubTask {
|
||||
/** Unique ID for this sub-task, e.g. "sidebar", "header" */
|
||||
id: string
|
||||
/** Human-readable label for progress UI */
|
||||
label: string
|
||||
/** Spatial region allocated on the canvas */
|
||||
region: { width: number; height: number }
|
||||
/** ID prefix for all nodes generated by this sub-agent (assigned at runtime) */
|
||||
idPrefix: string
|
||||
/** Parent frame ID to insert into (assigned at runtime) */
|
||||
parentFrameId: string | null
|
||||
}
|
||||
|
||||
/** Plan produced by the orchestrator (lightweight — only structure) */
|
||||
export interface OrchestratorPlan {
|
||||
rootFrame: {
|
||||
id: string
|
||||
name: string
|
||||
width: number
|
||||
height: number
|
||||
layout?: 'none' | 'vertical' | 'horizontal'
|
||||
gap?: number
|
||||
fill?: Array<{ type: string; color: string }>
|
||||
}
|
||||
subtasks: SubTask[]
|
||||
}
|
||||
|
||||
/** Progress state for the orchestrated generation */
|
||||
export interface OrchestrationProgress {
|
||||
phase: 'planning' | 'generating' | 'merging' | 'done' | 'error'
|
||||
subtasks: Array<{
|
||||
id: string
|
||||
label: string
|
||||
status: 'pending' | 'streaming' | 'done' | 'error'
|
||||
nodeCount: number
|
||||
}>
|
||||
totalNodes: number
|
||||
}
|
||||
|
||||
/** Result from a single sub-agent */
|
||||
export interface SubAgentResult {
|
||||
subtaskId: string
|
||||
nodes: import('@/types/pen').PenNode[]
|
||||
rawResponse: string
|
||||
error?: string
|
||||
}
|
||||
|
|
|
|||
90
src/services/ai/complexity-classifier.ts
Normal file
90
src/services/ai/complexity-classifier.ts
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
/**
|
||||
* Heuristic complexity classifier for design prompts.
|
||||
* Determines whether a prompt should be routed through
|
||||
* the orchestrator (parallel sub-agents) or handled directly.
|
||||
*
|
||||
* No API calls — runs in <1ms. False positives are cheap
|
||||
* (one extra planning call), false negatives are safe
|
||||
* (fall back to existing single-call path).
|
||||
*/
|
||||
|
||||
// Structural keywords indicating distinct spatial sections
|
||||
const SECTION_KEYWORDS = [
|
||||
'sidebar', 'header', 'footer', 'nav', 'navigation',
|
||||
'hero', 'stats', 'statistics', 'chart', 'table',
|
||||
'form', 'modal', 'dialog', 'card section', 'cards section',
|
||||
'profile', 'settings', 'feed', 'content area',
|
||||
'grid', 'gallery', 'carousel', 'tabs',
|
||||
'search bar', 'filter', 'toolbar', 'breadcrumb',
|
||||
'notification', 'calendar', 'timeline',
|
||||
// Chinese equivalents
|
||||
'侧边栏', '头部', '底部', '导航', '英雄区',
|
||||
'统计', '图表', '表格', '表单', '弹窗',
|
||||
'卡片区', '个人资料', '设置', '时间线',
|
||||
]
|
||||
|
||||
// Full-page keywords that almost always produce complex output
|
||||
const COMPLEX_PAGE_KEYWORDS = [
|
||||
'dashboard', 'landing page', 'homepage', 'e-commerce',
|
||||
'admin panel', 'social media', 'email client',
|
||||
'settings page', 'analytics', 'crm', 'erp',
|
||||
'portfolio', 'blog', 'marketplace', 'checkout',
|
||||
// Chinese
|
||||
'仪表盘', '着陆页', '首页', '电商', '管理后台',
|
||||
'社交', '邮件', '分析页',
|
||||
]
|
||||
|
||||
// Simple single-component keywords (suppress orchestration)
|
||||
const SIMPLE_KEYWORDS = [
|
||||
'button', 'input', 'avatar', 'badge', 'tooltip',
|
||||
'toggle', 'switch', 'checkbox', 'radio',
|
||||
'tag', 'chip', 'divider', 'spinner', 'icon',
|
||||
'按钮', '输入框', '头像', '标签', '开关',
|
||||
]
|
||||
|
||||
/** Minimum sections required to trigger orchestrator */
|
||||
const COMPLEXITY_THRESHOLD = 3
|
||||
|
||||
export interface ComplexityAssessment {
|
||||
isComplex: boolean
|
||||
estimatedSections: number
|
||||
reason: string
|
||||
}
|
||||
|
||||
export function assessComplexity(prompt: string): ComplexityAssessment {
|
||||
const lower = prompt.toLowerCase()
|
||||
|
||||
// Count distinct section keywords mentioned
|
||||
const mentionedSections = SECTION_KEYWORDS.filter((kw) => lower.includes(kw))
|
||||
const isFullPage = COMPLEX_PAGE_KEYWORDS.some((kw) => lower.includes(kw))
|
||||
const isSimpleComponent =
|
||||
SIMPLE_KEYWORDS.some((kw) => lower.includes(kw)) && mentionedSections.length <= 1
|
||||
|
||||
// Count conjunctions as section indicators
|
||||
const conjunctions = (lower.match(/\band\b/g) || []).length
|
||||
+ (lower.match(/[,，、]/g) || []).length
|
||||
|
||||
const estimatedSections = isFullPage
|
||||
? Math.max(mentionedSections.length, 4)
|
||||
: mentionedSections.length + Math.min(Math.floor(conjunctions / 2), 2)
|
||||
|
||||
if (isSimpleComponent) {
|
||||
return {
|
||||
isComplex: false,
|
||||
estimatedSections: 1,
|
||||
reason: 'Simple single component request',
|
||||
}
|
||||
}
|
||||
|
||||
const isComplex =
|
||||
estimatedSections >= COMPLEXITY_THRESHOLD ||
|
||||
(isFullPage && estimatedSections >= 2)
|
||||
|
||||
return {
|
||||
isComplex,
|
||||
estimatedSections,
|
||||
reason: isComplex
|
||||
? `Detected ${mentionedSections.length} sections (${mentionedSections.join(', ')})`
|
||||
: `Below threshold: ${estimatedSections} sections`,
|
||||
}
|
||||
}
|
||||
56
src/services/ai/context-optimizer.ts
Normal file
56
src/services/ai/context-optimizer.ts
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
/**
|
||||
* Context optimization utilities for AI chat.
|
||||
* Prevents unbounded growth of chat history and context size.
|
||||
*/
|
||||
|
||||
const DEFAULT_MAX_MESSAGES = 10
|
||||
const DEFAULT_MAX_CHARS = 32_000
|
||||
|
||||
/**
|
||||
* Sliding window for chat history.
|
||||
* Keeps the most recent messages while respecting character limits.
|
||||
* Always preserves the first user message for context continuity.
|
||||
*/
|
||||
export function trimChatHistory<T extends { role: string; content: string }>(
|
||||
messages: T[],
|
||||
maxMessages: number = DEFAULT_MAX_MESSAGES,
|
||||
maxChars: number = DEFAULT_MAX_CHARS,
|
||||
): T[] {
|
||||
if (messages.length <= maxMessages) {
|
||||
const totalChars = messages.reduce((sum, m) => sum + m.content.length, 0)
|
||||
if (totalChars <= maxChars) return messages
|
||||
}
|
||||
|
||||
// Always keep the first user message for context continuity
|
||||
const firstUser = messages.find((m) => m.role === 'user')
|
||||
const recentMessages = messages.slice(-maxMessages)
|
||||
|
||||
const window: T[] = []
|
||||
let charCount = 0
|
||||
|
||||
// Add first user message if it's not already in the recent window
|
||||
if (firstUser && !recentMessages.includes(firstUser)) {
|
||||
window.push(firstUser)
|
||||
charCount += firstUser.content.length
|
||||
}
|
||||
|
||||
// Add recent messages, respecting char limit
|
||||
for (const msg of recentMessages) {
|
||||
const msgChars = msg.content.length
|
||||
if (charCount + msgChars > maxChars) {
|
||||
// Truncate this message to fit
|
||||
const remaining = maxChars - charCount
|
||||
if (remaining > 200) {
|
||||
window.push({
|
||||
...msg,
|
||||
content: msg.content.slice(0, remaining) + '\n[...truncated...]',
|
||||
} as T)
|
||||
}
|
||||
break
|
||||
}
|
||||
window.push(msg)
|
||||
charCount += msgChars
|
||||
}
|
||||
|
||||
return window
|
||||
}
|
||||
|
|
@ -11,10 +11,13 @@ import {
|
|||
startNewAnimationBatch,
|
||||
resetAnimationState,
|
||||
} from './design-animation'
|
||||
import { assessComplexity } from './complexity-classifier'
|
||||
import { executeOrchestration } from './orchestrator'
|
||||
|
||||
const DESIGN_STREAM_TIMEOUTS = {
|
||||
hardTimeoutMs: 180_000,
|
||||
noTextTimeoutMs: 60_000,
|
||||
thinkingResetsTimeout: false,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
|
@ -24,13 +27,13 @@ const DESIGN_STREAM_TIMEOUTS = {
|
|||
|
||||
const generationRemappedIds = new Map<string, string>()
|
||||
|
||||
function resetGenerationRemapping(): void {
|
||||
export function resetGenerationRemapping(): void {
|
||||
generationRemappedIds.clear()
|
||||
}
|
||||
|
||||
// Helper to find all complete JSON blocks in text
|
||||
|
||||
function extractJsonFromResponse(text: string): PenNode[] | null {
|
||||
export function extractJsonFromResponse(text: string): PenNode[] | null {
|
||||
const parsedBlocks = extractAllJsonBlocks(text)
|
||||
.map((block) => tryParseNodes(block))
|
||||
.filter(Boolean) as PenNode[][]
|
||||
|
|
@ -97,7 +100,7 @@ function buildContextMessage(request: AIDesignRequest): string {
|
|||
}
|
||||
|
||||
/** Build a concise summary of document variables for AI context. */
|
||||
function buildVariableContext(
|
||||
export function buildVariableContext(
|
||||
variables?: Record<string, VariableDefinition>,
|
||||
themes?: Record<string, string[]>,
|
||||
): string | null {
|
||||
|
|
@ -159,7 +162,7 @@ interface StreamingNodeResult {
|
|||
* Uses brace-counting to detect complete objects before the block closes.
|
||||
* Each object is expected to have a `_parent` field for tree insertion.
|
||||
*/
|
||||
function extractStreamingNodes(
|
||||
export function extractStreamingNodes(
|
||||
text: string,
|
||||
processedOffset: number,
|
||||
): { results: StreamingNodeResult[]; newOffset: number } {
|
||||
|
|
@ -260,7 +263,7 @@ function parseJsonlToTree(text: string): PenNode[] | null {
|
|||
* Insert a single streaming node into the canvas with animation.
|
||||
* Handles root frame replacement and parent ID remapping.
|
||||
*/
|
||||
function insertStreamingNode(
|
||||
export function insertStreamingNode(
|
||||
node: PenNode,
|
||||
parentId: string | null,
|
||||
): void {
|
||||
|
|
@ -341,6 +344,17 @@ export async function generateDesign(
|
|||
animated?: boolean
|
||||
}
|
||||
): Promise<{ nodes: PenNode[]; rawResponse: string }> {
|
||||
// Route complex prompts through orchestrator for parallel generation
|
||||
const { isComplex } = assessComplexity(request.prompt)
|
||||
if (isComplex) {
|
||||
try {
|
||||
return await executeOrchestration(request, callbacks)
|
||||
} catch (err) {
|
||||
// Orchestrator failed — silently fall back to single-call generation
|
||||
console.warn('Orchestrator failed, falling back to direct generation:', err)
|
||||
}
|
||||
}
|
||||
|
||||
const userMessage = buildContextMessage(request)
|
||||
let fullResponse = ''
|
||||
let streamingOffset = 0 // Tracks how far we've parsed in the streaming text
|
||||
|
|
|
|||
20
src/services/ai/orchestrator-prompts.ts
Normal file
20
src/services/ai/orchestrator-prompts.ts
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
/**
|
||||
* Orchestrator prompt — ultra-lightweight, only splits into sections.
|
||||
* No design details, no prompt rewriting. Just structure.
|
||||
*/
|
||||
|
||||
export const ORCHESTRATOR_PROMPT = `Split a UI request into 2-6 spatial sections. Output ONLY JSON, start with {.
|
||||
|
||||
FORMAT:
|
||||
{"rootFrame":{"id":"page","name":"Page","width":1200,"height":800,"layout":"vertical","fill":[{"type":"solid","color":"#16171B"}]},"subtasks":[{"id":"nav","label":"Navigation Bar","region":{"width":1200,"height":56}},{"id":"hero","label":"Hero Section","region":{"width":1200,"height":400}}]}
|
||||
|
||||
RULES:
|
||||
- Each subtask = one section with id, label, region (width+height).
|
||||
- Regions tile to fill rootFrame. vertical = top-to-bottom.
|
||||
- Mobile: 375x812. Desktop: 1200x800.
|
||||
- NO explanation. NO markdown. JUST the JSON object.`
|
||||
|
||||
export const ORCHESTRATOR_TIMEOUTS = {
|
||||
hardTimeoutMs: 30_000,
|
||||
noTextTimeoutMs: 20_000,
|
||||
}
|
||||
476
src/services/ai/orchestrator.ts
Normal file
476
src/services/ai/orchestrator.ts
Normal file
|
|
@ -0,0 +1,476 @@
|
|||
/**
|
||||
* Orchestrator for parallel design generation.
|
||||
*
|
||||
* Flow:
|
||||
* 1. Fast "architect" API call decomposes the prompt into spatial sub-tasks
|
||||
* 2. Root frame is created on canvas
|
||||
* 3. Multiple sub-agents execute in parallel, each streaming JSONL
|
||||
* 4. Nodes are inserted to canvas in real-time with animation
|
||||
*
|
||||
* Falls back to single-call generation on any orchestrator failure.
|
||||
*/
|
||||
|
||||
import type { PenNode, FrameNode } from '@/types/pen'
|
||||
import type { VariableDefinition } from '@/types/variables'
|
||||
import type {
|
||||
AIDesignRequest,
|
||||
OrchestratorPlan,
|
||||
OrchestrationProgress,
|
||||
SubTask,
|
||||
SubAgentResult,
|
||||
} from './ai-types'
|
||||
import { streamChat } from './ai-service'
|
||||
import { DESIGN_GENERATOR_PROMPT } from './ai-prompts'
|
||||
import { ORCHESTRATOR_PROMPT, ORCHESTRATOR_TIMEOUTS } from './orchestrator-prompts'
|
||||
import {
|
||||
extractStreamingNodes,
|
||||
extractJsonFromResponse,
|
||||
insertStreamingNode,
|
||||
buildVariableContext,
|
||||
resetGenerationRemapping,
|
||||
} from './design-generator'
|
||||
import { useDocumentStore } from '@/stores/document-store'
|
||||
import { useHistoryStore } from '@/stores/history-store'
|
||||
import {
|
||||
resetAnimationState,
|
||||
startNewAnimationBatch,
|
||||
} from './design-animation'
|
||||
|
||||
const SUB_AGENT_TIMEOUTS = {
|
||||
hardTimeoutMs: 120_000,
|
||||
noTextTimeoutMs: 45_000,
|
||||
thinkingResetsTimeout: true,
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Public API
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export async function executeOrchestration(
|
||||
request: AIDesignRequest,
|
||||
callbacks?: {
|
||||
onApplyPartial?: (count: number) => void
|
||||
onTextUpdate?: (text: string) => void
|
||||
animated?: boolean
|
||||
},
|
||||
): Promise<{ nodes: PenNode[]; rawResponse: string }> {
|
||||
const animated = callbacks?.animated ?? false
|
||||
|
||||
// -- Phase 1: Planning --
|
||||
callbacks?.onTextUpdate?.(
|
||||
'<step title="Planning layout" status="streaming">Analyzing design structure...</step>',
|
||||
)
|
||||
|
||||
const plan = await callOrchestrator(request.prompt, (thinking) => {
|
||||
// Forward thinking progress to UI so user sees activity
|
||||
const truncated = thinking.length > 200
|
||||
? thinking.slice(-200) + '...'
|
||||
: thinking
|
||||
callbacks?.onTextUpdate?.(
|
||||
`<step title="Planning layout" status="streaming">${truncated}</step>`,
|
||||
)
|
||||
})
|
||||
|
||||
// Assign ID prefixes
|
||||
for (const st of plan.subtasks) {
|
||||
st.idPrefix = st.id
|
||||
st.parentFrameId = plan.rootFrame.id
|
||||
}
|
||||
|
||||
// Show planning done + all subtask steps as pending
|
||||
emitProgress(plan, {
|
||||
phase: 'generating',
|
||||
subtasks: plan.subtasks.map((st) => ({
|
||||
id: st.id, label: st.label, status: 'pending' as const, nodeCount: 0,
|
||||
})),
|
||||
totalNodes: 0,
|
||||
}, callbacks)
|
||||
|
||||
// -- Phase 2: Setup canvas --
|
||||
resetGenerationRemapping()
|
||||
|
||||
if (animated) {
|
||||
resetAnimationState()
|
||||
useHistoryStore.getState().startBatch(useDocumentStore.getState().document)
|
||||
}
|
||||
|
||||
// Insert root frame
|
||||
const rootNode: FrameNode = {
|
||||
id: plan.rootFrame.id,
|
||||
type: 'frame',
|
||||
name: plan.rootFrame.name,
|
||||
x: 0,
|
||||
y: 0,
|
||||
width: plan.rootFrame.width,
|
||||
height: plan.rootFrame.height,
|
||||
layout: plan.rootFrame.layout ?? 'vertical',
|
||||
gap: plan.rootFrame.gap ?? 0,
|
||||
fill: (plan.rootFrame.fill as FrameNode['fill']) ?? [
|
||||
{ type: 'solid', color: '#FFFFFF' },
|
||||
],
|
||||
children: [],
|
||||
}
|
||||
|
||||
insertStreamingNode(rootNode, null)
|
||||
|
||||
// -- Phase 3: Parallel sub-agent execution --
|
||||
const progress: OrchestrationProgress = {
|
||||
phase: 'generating',
|
||||
subtasks: plan.subtasks.map((st) => ({
|
||||
id: st.id,
|
||||
label: st.label,
|
||||
status: 'pending' as const,
|
||||
nodeCount: 0,
|
||||
})),
|
||||
totalNodes: 0,
|
||||
}
|
||||
|
||||
let results: SubAgentResult[]
|
||||
try {
|
||||
results = await executeSubAgentsSequentially(
|
||||
plan,
|
||||
request,
|
||||
progress,
|
||||
callbacks,
|
||||
)
|
||||
} finally {
|
||||
if (animated) {
|
||||
useHistoryStore.getState().endBatch(useDocumentStore.getState().document)
|
||||
}
|
||||
}
|
||||
|
||||
// -- Phase 4: Collect results --
|
||||
// Mark all completed subtasks as done in final progress
|
||||
for (const entry of progress.subtasks) {
|
||||
if (entry.status !== 'error') {
|
||||
entry.status = 'done'
|
||||
}
|
||||
}
|
||||
progress.phase = 'done'
|
||||
emitProgress(plan, progress, callbacks)
|
||||
|
||||
const allNodes: PenNode[] = [rootNode]
|
||||
for (const r of results) {
|
||||
allNodes.push(...r.nodes)
|
||||
}
|
||||
|
||||
// Build final rawResponse that includes step tags so the chat message
|
||||
// shows the complete pipeline progress after streaming ends
|
||||
const finalStepTags = buildFinalStepTags(plan, progress)
|
||||
|
||||
return { nodes: allNodes, rawResponse: finalStepTags }
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Orchestrator call — fast decomposition
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/** Max prompt length for the orchestrator — it only needs structure, not full detail */
|
||||
const MAX_ORCHESTRATOR_PROMPT_CHARS = 2000
|
||||
|
||||
async function callOrchestrator(
|
||||
prompt: string,
|
||||
onThinking?: (thinking: string) => void,
|
||||
): Promise<OrchestratorPlan> {
|
||||
// Truncate long prompts — the orchestrator only needs high-level structure
|
||||
const truncatedPrompt = prompt.length > MAX_ORCHESTRATOR_PROMPT_CHARS
|
||||
? prompt.slice(0, MAX_ORCHESTRATOR_PROMPT_CHARS) + '\n\n[... prompt truncated for planning — full details will be sent to each section agent]'
|
||||
: prompt
|
||||
|
||||
let rawResponse = ''
|
||||
let thinkingContent = ''
|
||||
|
||||
for await (const chunk of streamChat(
|
||||
ORCHESTRATOR_PROMPT,
|
||||
[{ role: 'user', content: truncatedPrompt }],
|
||||
undefined,
|
||||
{
|
||||
...ORCHESTRATOR_TIMEOUTS,
|
||||
// Don't let thinking indefinitely extend the no-text timeout
|
||||
thinkingResetsTimeout: false,
|
||||
},
|
||||
)) {
|
||||
if (chunk.type === 'text') {
|
||||
rawResponse += chunk.content
|
||||
} else if (chunk.type === 'thinking') {
|
||||
thinkingContent += chunk.content
|
||||
onThinking?.(thinkingContent)
|
||||
} else if (chunk.type === 'error') {
|
||||
throw new Error(`Orchestrator failed: ${chunk.content}`)
|
||||
}
|
||||
}
|
||||
|
||||
const plan = parseOrchestratorResponse(rawResponse)
|
||||
if (!plan) {
|
||||
throw new Error('Failed to parse orchestrator plan')
|
||||
}
|
||||
|
||||
return plan
|
||||
}
|
||||
|
||||
function parseOrchestratorResponse(raw: string): OrchestratorPlan | null {
|
||||
const trimmed = raw.trim()
|
||||
|
||||
// Try direct parse
|
||||
const plan = tryParsePlan(trimmed)
|
||||
if (plan) return plan
|
||||
|
||||
// Try extracting from code fences
|
||||
const fenceMatch = trimmed.match(/```(?:json)?\s*\n?([\s\S]*?)\n?```/)
|
||||
if (fenceMatch) {
|
||||
const fenced = tryParsePlan(fenceMatch[1].trim())
|
||||
if (fenced) return fenced
|
||||
}
|
||||
|
||||
// Try extracting first { ... } block
|
||||
const firstBrace = trimmed.indexOf('{')
|
||||
const lastBrace = trimmed.lastIndexOf('}')
|
||||
if (firstBrace >= 0 && lastBrace > firstBrace) {
|
||||
const braced = tryParsePlan(trimmed.slice(firstBrace, lastBrace + 1))
|
||||
if (braced) return braced
|
||||
}
|
||||
|
||||
return null
|
||||
}
|
||||
|
||||
function tryParsePlan(text: string): OrchestratorPlan | null {
|
||||
try {
|
||||
const obj = JSON.parse(text) as Record<string, unknown>
|
||||
if (!obj.rootFrame || typeof obj.rootFrame !== 'object') return null
|
||||
if (!Array.isArray(obj.subtasks) || obj.subtasks.length === 0) return null
|
||||
|
||||
const rf = obj.rootFrame as Record<string, unknown>
|
||||
if (!rf.id || !rf.width || !rf.height) return null
|
||||
|
||||
for (const st of obj.subtasks as Record<string, unknown>[]) {
|
||||
if (!st.id || !st.region) return null
|
||||
}
|
||||
|
||||
return obj as unknown as OrchestratorPlan
|
||||
} catch {
|
||||
return null
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Parallel sub-agent execution
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
async function executeSubAgentsSequentially(
|
||||
plan: OrchestratorPlan,
|
||||
request: AIDesignRequest,
|
||||
progress: OrchestrationProgress,
|
||||
callbacks?: {
|
||||
onApplyPartial?: (count: number) => void
|
||||
onTextUpdate?: (text: string) => void
|
||||
animated?: boolean
|
||||
},
|
||||
): Promise<SubAgentResult[]> {
|
||||
const results: SubAgentResult[] = []
|
||||
for (let i = 0; i < plan.subtasks.length; i++) {
|
||||
const result = await executeSubAgent(plan.subtasks[i], plan, request, progress, i, callbacks)
|
||||
results.push(result)
|
||||
}
|
||||
return results
|
||||
}
|
||||
|
||||
async function executeSubAgent(
|
||||
subtask: SubTask,
|
||||
plan: OrchestratorPlan,
|
||||
request: AIDesignRequest,
|
||||
progress: OrchestrationProgress,
|
||||
index: number,
|
||||
callbacks?: {
|
||||
onApplyPartial?: (count: number) => void
|
||||
onTextUpdate?: (text: string) => void
|
||||
animated?: boolean
|
||||
},
|
||||
): Promise<SubAgentResult> {
|
||||
const animated = callbacks?.animated ?? false
|
||||
const progressEntry = progress.subtasks[index]
|
||||
progressEntry.status = 'streaming'
|
||||
emitProgress(plan, progress, callbacks)
|
||||
|
||||
const userPrompt = buildSubAgentUserPrompt(
|
||||
subtask,
|
||||
request.prompt,
|
||||
request.context?.variables,
|
||||
request.context?.themes,
|
||||
)
|
||||
|
||||
let rawResponse = ''
|
||||
const nodes: PenNode[] = []
|
||||
let streamOffset = 0
|
||||
|
||||
try {
|
||||
for await (const chunk of streamChat(
|
||||
DESIGN_GENERATOR_PROMPT,
|
||||
[{ role: 'user', content: userPrompt }],
|
||||
undefined,
|
||||
SUB_AGENT_TIMEOUTS,
|
||||
)) {
|
||||
if (chunk.type === 'text') {
|
||||
rawResponse += chunk.content
|
||||
|
||||
if (animated) {
|
||||
const { results, newOffset } = extractStreamingNodes(
|
||||
rawResponse,
|
||||
streamOffset,
|
||||
)
|
||||
if (results.length > 0) {
|
||||
streamOffset = newOffset
|
||||
startNewAnimationBatch()
|
||||
|
||||
for (const { node, parentId } of results) {
|
||||
// Enforce ID prefix
|
||||
ensureIdPrefix(node, subtask.idPrefix)
|
||||
if (parentId !== null) {
|
||||
// Prefix the parent reference too
|
||||
const prefixedParent = ensurePrefixStr(
|
||||
parentId,
|
||||
subtask.idPrefix,
|
||||
)
|
||||
insertStreamingNode(node, prefixedParent)
|
||||
} else {
|
||||
// Sub-agent root → insert under the orchestrator root frame
|
||||
insertStreamingNode(node, plan.rootFrame.id)
|
||||
}
|
||||
nodes.push(node)
|
||||
progressEntry.nodeCount++
|
||||
progress.totalNodes++
|
||||
}
|
||||
callbacks?.onApplyPartial?.(progress.totalNodes)
|
||||
emitProgress(plan, progress, callbacks)
|
||||
}
|
||||
}
|
||||
} else if (chunk.type === 'thinking') {
|
||||
// Forward thinking progress so UI doesn't look stuck
|
||||
emitProgress(plan, progress, callbacks)
|
||||
} else if (chunk.type === 'error') {
|
||||
progressEntry.status = 'error'
|
||||
emitProgress(plan, progress, callbacks)
|
||||
return { subtaskId: subtask.id, nodes, rawResponse, error: chunk.content }
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback: if streaming extraction found nothing, try batch extraction
|
||||
if (nodes.length === 0 && rawResponse.trim().length > 0) {
|
||||
const fallbackNodes = extractJsonFromResponse(rawResponse)
|
||||
if (fallbackNodes && fallbackNodes.length > 0) {
|
||||
startNewAnimationBatch()
|
||||
for (const node of fallbackNodes) {
|
||||
ensureIdPrefix(node, subtask.idPrefix)
|
||||
insertStreamingNode(node, plan.rootFrame.id)
|
||||
nodes.push(node)
|
||||
progressEntry.nodeCount++
|
||||
progress.totalNodes++
|
||||
}
|
||||
callbacks?.onApplyPartial?.(progress.totalNodes)
|
||||
}
|
||||
}
|
||||
|
||||
progressEntry.status = 'done'
|
||||
emitProgress(plan, progress, callbacks)
|
||||
return { subtaskId: subtask.id, nodes, rawResponse }
|
||||
} catch (err) {
|
||||
const msg = err instanceof Error ? err.message : 'Unknown error'
|
||||
progressEntry.status = 'error'
|
||||
emitProgress(plan, progress, callbacks)
|
||||
return { subtaskId: subtask.id, nodes, rawResponse, error: msg }
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Sub-agent prompt builder
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function buildSubAgentUserPrompt(
|
||||
subtask: SubTask,
|
||||
originalPrompt: string,
|
||||
variables?: Record<string, VariableDefinition>,
|
||||
themes?: Record<string, string[]>,
|
||||
): string {
|
||||
const { region } = subtask
|
||||
|
||||
// Extract a brief context line from the original prompt (first 200 chars)
|
||||
const briefContext = originalPrompt.length > 200
|
||||
? originalPrompt.slice(0, 200) + '...'
|
||||
: originalPrompt
|
||||
|
||||
let prompt = `Design: "${subtask.label}"
|
||||
Context: ${briefContext}
|
||||
Canvas: ${region.width}x${region.height}px
|
||||
Root frame: id="${subtask.idPrefix}-root", width=${region.width}, height=${region.height}
|
||||
All node IDs MUST start with "${subtask.idPrefix}-".
|
||||
Generate ONLY this section, not the full page.`
|
||||
|
||||
const varContext = buildVariableContext(variables, themes)
|
||||
if (varContext) {
|
||||
prompt += '\n\n' + varContext
|
||||
}
|
||||
|
||||
return prompt
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// ID namespace isolation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function ensureIdPrefix(node: PenNode, prefix: string): void {
|
||||
if (!node.id.startsWith(`${prefix}-`)) {
|
||||
node.id = `${prefix}-${node.id}`
|
||||
}
|
||||
}
|
||||
|
||||
function ensurePrefixStr(id: string, prefix: string): string {
|
||||
if (id.startsWith(`${prefix}-`)) return id
|
||||
return `${prefix}-${id}`
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Progress emission — updates UI via <step> tags
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function emitProgress(
|
||||
plan: OrchestratorPlan,
|
||||
progress: OrchestrationProgress,
|
||||
callbacks?: {
|
||||
onTextUpdate?: (text: string) => void
|
||||
},
|
||||
): void {
|
||||
if (!callbacks?.onTextUpdate) return
|
||||
|
||||
// Always show "Planning layout" as done first
|
||||
const planningStep = '<step title="Planning layout" status="done">Analyzing design structure...</step>'
|
||||
|
||||
const subtaskSteps = plan.subtasks
|
||||
.map((st, i) => {
|
||||
const entry = progress.subtasks[i]
|
||||
const status = entry.status === 'streaming' ? 'streaming'
|
||||
: entry.status === 'done' ? 'done'
|
||||
: entry.status === 'error' ? 'error'
|
||||
: 'pending'
|
||||
const nodeInfo = entry.nodeCount > 0 ? ` (${entry.nodeCount} elements)` : ''
|
||||
return `<step title="${st.label}${nodeInfo}" status="${status}"></step>`
|
||||
})
|
||||
.join('\n')
|
||||
|
||||
callbacks.onTextUpdate(`${planningStep}\n${subtaskSteps}`)
|
||||
}
|
||||
|
||||
/** Build step tags for the final rawResponse (shown in message after streaming ends) */
|
||||
function buildFinalStepTags(
|
||||
plan: OrchestratorPlan,
|
||||
progress: OrchestrationProgress,
|
||||
): string {
|
||||
const planningStep = '<step title="Planning layout" status="done">Analyzing design structure...</step>'
|
||||
const subtaskSteps = plan.subtasks
|
||||
.map((st, i) => {
|
||||
const entry = progress.subtasks[i]
|
||||
const status = entry.status
|
||||
const nodeInfo = entry.nodeCount > 0 ? ` (${entry.nodeCount} elements)` : ''
|
||||
return `<step title="${st.label}${nodeInfo}" status="${status}"></step>`
|
||||
})
|
||||
.join('\n')
|
||||
return `${planningStep}\n${subtaskSteps}`
|
||||
}
|
||||
Loading…
Reference in a new issue