# AI Chat Builder — Design Document

> **Date:** 2026-02-27
> **Status:** Approved
> **Relationship to existing wizard:** Coexists. Wizard stays as "Quick Build," this becomes "Build with AI" with a separate page and entry point.

---

## Overview

A conversational AI flow builder where the AI acts as a senior MSP engineer, conducting a multi-phase interview to collaboratively build troubleshooting or procedural flows. The user chats naturally; the AI demonstrates domain expertise, suggests diagnostic steps, challenges assumptions, and progressively constructs a TreeStructure that imports into the Tree Editor.

This is distinct from the existing wizard-based AI builder (scaffold → branch detail → assemble). The wizard is fast and automated; the chat builder is collaborative and guided.

---

## Key Decisions

| Decision | Choice | Rationale |
|----------|--------|-----------|
| Relationship to wizard | Keep both, separate entry points | Different use cases: quick generation vs. guided interview |
| SSE streaming | Not in MVP | Full-message responses with loading spinner. Add streaming later. |
| Database model | New `ai_chat_sessions` table | Different workflow/state shape than wizard's `ai_conversations` |
| Backend organization | Follow existing `core/ai_*.py` pattern | Consistent with codebase. No new `services/` directory. |
| Provider failover | Skip for MVP | `get_ai_provider()` selects configured provider. Failover later if needed. |

---

## Backend Data Model

### New Table: `ai_chat_sessions`

| Column | Type | Notes |
|--------|------|-------|
| `id` | UUID PK | `default=uuid4` |
| `user_id` | UUID FK → users | `ondelete=CASCADE` |
| `account_id` | UUID FK → accounts | `ondelete=CASCADE` |
| `status` | VARCHAR | `active`, `completed`, `abandoned` |
| `current_phase` | VARCHAR | `scoping`, `discovery`, `enrichment`, `review`, `generation` |
| `conversation_history` | JSONB | `[{role, content, timestamp}]` |
| `working_tree` | JSONB nullable | Progressive TreeStructure, updated as AI builds |
| `tree_metadata` | JSONB | `{name, description, category_id, tags}` extracted during conversation |
| `flow_type` | VARCHAR | `troubleshooting` or `procedural` |
| `provider_used` | VARCHAR nullable | Which AI provider served this session |
| `message_count` | INTEGER | Per-session exchange cap (prevents runaway token spend) |
| `total_input_tokens` | INTEGER | Running total for cost tracking |
| `total_output_tokens` | INTEGER | Running total |
| `generated_tree_id` | UUID FK → trees nullable | Set when user imports to editor |
| `expires_at` | DateTime(tz) | 24-hour TTL |
| `created_at` | DateTime(tz) | |
| `updated_at` | DateTime(tz) | |

### Quota Tracking

Reuses existing `ai_usage` table with new generation types:
- `chat_message` — each exchange, `counts_toward_quota=False`
- `chat_generate` — final tree generation, `counts_toward_quota=True`

**Required change:** Update `ai_quota_service.py` daily limit query to include chat builder types in the `generation_type.in_(...)` filter.

### Enforcement Layers

1. **Account monthly** — `ai_usage` with `counts_toward_quota=True` (existing, works across both builders)
2. **Account daily** — `ai_usage` by generation type (update filter to include chat types)
3. **Per-session** — `message_count` cap on session table (prevents runaway single-session spend)

---

## Backend API Endpoints

Router: `api/endpoints/ai_chat.py`, prefix `/ai/chat`, tag `ai-chat-builder`

| Method | Endpoint | Purpose |
|--------|----------|---------|
| `POST` | `/ai/chat/sessions` | Start session. Check quota, create session, return AI greeting. |
| `POST` | `/ai/chat/sessions/{id}/messages` | Send message, get AI response. Updates working_tree if AI includes tree update. |
| `GET` | `/ai/chat/sessions/{id}` | Get full session state. For resume after page reload. |
| `POST` | `/ai/chat/sessions/{id}/generate` | Final TreeStructure JSON generation. Validates, one retry on failure. |
| `POST` | `/ai/chat/sessions/{id}/import` | Create Tree record from generated tree. Returns tree ID. |
| `DELETE` | `/ai/chat/sessions/{id}` | Abandon session (soft — sets status to `abandoned`). |

**Rate limiting:** `@limiter.limit("10/minute")` on all mutating endpoints.

**Timeout handling for `/generate`:**
- AI provider timeout → 504 with clear "Generation timed out, please try again" message
- Failed attempts recorded with `counts_toward_quota=False`
- Session stays in `review` phase so user can retry
- Guard: if session already has a completed generation, return cached result (prevents duplicate costs from double-clicks)

**Ownership:** All endpoints verify `session.user_id == current_user.id`.

---

## Backend Service Layer

New file: `core/ai_chat_service.py`

### Functions

**`start_chat_session(flow_type, user_id, account_id, db) → (session, greeting_message)`**
- Creates session record
- Builds system prompt
- Sends initial prime to AI, returns greeting that demonstrates domain knowledge

**`send_message(session, user_message, db) → (ai_response, working_tree_update, phase)`**
- Appends user message to history
- Sends system prompt + full conversation history to AI
- Parses response: extracts conversational text, `[TREE_UPDATE]`, `[PHASE]`, `[METADATA]` markers
- Strips markers before storing the user-visible response
- Updates working_tree and phase if present
- Records usage via `record_ai_usage()`

**`generate_final_tree(session, db) → (tree_structure, metadata)`**
- Sends generation prompt with full conversation context
- Validates output (orphans, circular refs, missing terminals, valid next_node_id refs)
- One retry with correction prompt if validation fails
- Returns validated TreeStructure + metadata

### System Prompt Architecture

Four sections assembled programmatically:

1. **Role & Persona** — Senior MSP engineer, 15+ years, collegial and direct. Domain expertise: Windows Server, AD, Entra ID, M365, networking, virtualization, security, cloud.

2. **Schema Context** — TreeStructure node types with field definitions. Generated from Python dict to stay in sync with actual types.

3. **Interview Protocol** — Five-phase structure with behavioral rules:
   - One focused question at a time
   - Demonstrate domain knowledge immediately
   - Challenge assumptions constructively
   - Capture specific commands with exact syntax
   - Surface edge cases proactively
   - Explain WHY diagnostic order matters

4. **Response Format** — Structured output markers:
   - Conversational text (always present, this is what the user sees)
   - `[TREE_UPDATE]{...json...}[/TREE_UPDATE]` — only when tree structure changes
   - `[PHASE:phase_name]` — when transitioning phases
   - `[METADATA]{...json...}[/METADATA]` — when capturing tree name/description/tags

### Tree Update Rules (Critical for Prompt Quality)

The system prompt must be explicit about when to emit `[TREE_UPDATE]`:

- **Scoping phase:** Never. Still understanding the problem space.
- **Discovery phase:** Only when a concrete node is established — a decision with clear options, or an action with a specific command. Not during exploratory questions.
- **Enrichment phase:** When enriching existing nodes or adding edge case branches.
- **Review phase:** Only if user requests structural changes.
- **General rule:** "If you're asking a question, you're not updating the tree."

> **Known tuning area:** The tree update prompt rules will need iteration. First version will likely be too eager or too conservative. Test with Section 7 scenarios from the implementation guide and adjust.

---

## Frontend Architecture

### Page Layout

Split-panel inside existing AppLayout shell:

```
┌─────────────────────────────────────────────────────┐
│  Toolbar: [phase indicator]     [Import] [Abandon]  │
├──────────────────────────┬──────────────────────────┤
│     Chat Panel (60%)     │   Tree Preview (40%)     │
│                          │                          │
│  AI greeting message     │  Empty state until       │
│  User message            │  discovery phase         │
│  AI response             │                          │
│  ...                     │  TreePreviewPanel        │
│                          │  (reused component)      │
│  [Type a message...]     │                          │
└──────────────────────────┴──────────────────────────┘
```

Responsive: Below 1024px, stack vertically. Preview collapses to expandable panel.

### Route

`/ai/chat` → `AIChatBuilderPage` — under protected AppLayout children in `router.tsx`.

### Zustand Store: `store/aiChatStore.ts`

Plain store (no immer/temporal — conversation is append-only):

```typescript
interface AIChatState {
  sessionId: string | null
  status: 'idle' | 'active' | 'completed' | 'abandoned'
  currentPhase: InterviewPhase
  flowType: 'troubleshooting' | 'procedural' | null
  messages: ChatMessage[]
  isResponding: boolean
  workingTree: TreeStructure | null
  treeMetadata: { name?: string; description?: string; tags?: string[]; category_id?: string } | null
  generatedTree: TreeStructure | null
  isGenerating: boolean
  importedTreeId: string | null
  error: string | null

  startSession: (flowType) => Promise<void>
  sendMessage: (content: string) => Promise<void>
  generateTree: () => Promise<void>
  importToEditor: () => Promise<string>
  abandonSession: () => Promise<void>
  resumeSession: (sessionId: string) => Promise<void>
  reset: () => void
}
```

No localStorage persistence. Session state lives server-side. Page reload → `resumeSession()` rehydrates from API.

### New Components (`components/ai-chat/`)

| Component | Purpose |
|-----------|---------|
| `ChatPanel.tsx` | Scrollable message list + input. Auto-scroll on new messages. |
| `ChatMessage.tsx` | Message bubble. AI uses `MarkdownContent`, user messages plain. |
| `ChatInput.tsx` | Text input. Enter to send, Shift+Enter for newline. Disabled while AI responds. |
| `PhaseIndicator.tsx` | Horizontal breadcrumb: Scoping → Discovery → Enrichment → Review → Generate |
| `ChatToolbar.tsx` | Top bar with phase indicator + action buttons |
| `EmptyPreview.tsx` | Placeholder before tree data exists |

### Reused Components

- `TreePreviewPanel` — right panel tree visualization
- `MarkdownContent` — AI message rendering
- `PageLoader` / `Spinner` — loading states

### Entry Point

"Build with AI" button on `TreeLibraryPage.tsx` alongside existing Create menu. Lucide `Sparkles` icon. Routes to `/ai/chat`.

---

## Files

### New Files

| File | Purpose |
|------|---------|
| `backend/app/models/ai_chat_session.py` | SQLAlchemy model |
| `backend/app/schemas/ai_chat.py` | Pydantic schemas |
| `backend/app/core/ai_chat_service.py` | Conversation loop, system prompt, parsing |
| `backend/app/api/endpoints/ai_chat.py` | API endpoints |
| `backend/alembic/versions/XXX_add_ai_chat_sessions.py` | Migration |
| `frontend/src/types/ai-chat.ts` | TypeScript interfaces |
| `frontend/src/api/aiChat.ts` | API client module |
| `frontend/src/store/aiChatStore.ts` | Zustand store |
| `frontend/src/pages/AIChatBuilderPage.tsx` | Main page |
| `frontend/src/components/ai-chat/ChatPanel.tsx` | Chat message list |
| `frontend/src/components/ai-chat/ChatMessage.tsx` | Single message bubble |
| `frontend/src/components/ai-chat/ChatInput.tsx` | Text input + send |
| `frontend/src/components/ai-chat/PhaseIndicator.tsx` | Phase breadcrumb |
| `frontend/src/components/ai-chat/ChatToolbar.tsx` | Top bar |
| `frontend/src/components/ai-chat/EmptyPreview.tsx` | Preview placeholder |

### Modified Files

| File | Change |
|------|--------|
| `backend/app/api/router.py` | Include `ai_chat.router` |
| `backend/app/models/__init__.py` | Import `AIChatSession` |
| `backend/app/core/ai_quota_service.py` | Add chat builder types to daily limit query |
| `frontend/src/types/index.ts` | Export from `ai-chat.ts` |
| `frontend/src/api/index.ts` | Export `aiChat` module |
| `frontend/src/router.tsx` | Add `/ai/chat` route |
| `frontend/src/pages/TreeLibraryPage.tsx` | Add "Build with AI" button |

---

## Conflicts with Implementation Guide

| Guide Says | Resolution |
|-----------|------------|
| `services/ai/` directory structure | Follow existing `core/ai_*.py` pattern |
| `AIRouter` with automatic failover | Skip for MVP. Single configured provider. |
| SSE streaming | Skip for MVP. Full-message responses. |
| Separate `prompts/` and `validators/` modules | Prompt in service file. Validation reuses existing helpers. |
| `flow_builder_sessions` table name | `ai_chat_sessions` for naming consistency |
| Subscription tier table (Free: 2, Pro: 20, Team: unlimited) | Existing `plan_limits` handles this already |
| Caching for similar problem domains | Skip. Premature optimization. |
| Guide's Phase 4 polish items | Session resume included. Mobile is basic stacking. Full accessibility post-MVP. |

---

## Implementation Phases (from guide, adapted)

**Phase 1: Backend foundation** — Model, migration, service, endpoints, system prompt, validation. Milestone: can POST a message and get a domain-knowledgeable AI response.

**Phase 2: Frontend chat UI** — Page, store, components, routing, entry point. Milestone: full conversation flow with AI, no tree preview yet.

**Phase 3: Tree preview & import** — Wire up TreePreviewPanel with working_tree updates, generate action, import to editor. Milestone: end-to-end flow from conversation to tree in editor.

**Phase 4: Polish** — Error recovery, session resume on reload, prompt tuning with test scenarios, mobile layout.

---

## Test Scenarios (from guide Section 7)

Use after each phase to verify quality:
1. Azure AD Connect Sync Failures (identity/hybrid cloud)
2. Intermittent Site-to-Site VPN Drops (networking)
3. Exchange Online Mail Flow Issues (M365/email)
4. Slow Computer Troubleshooting (Tier 1 desktop support)

---

## Open Items / Known Tuning Areas

- System prompt `[TREE_UPDATE]` emission rules will need iteration with real conversations
- System prompt persona/tone will need tuning to feel like a colleague, not a chatbot
- Final generation prompt may need a more capable model than conversation phase
- Message count cap per session needs a specific number (guide suggests 25 for paid tiers, 10 for free)