diff --git a/docs/plans/2026-02-27-ai-chat-builder-implementation.md b/docs/plans/2026-02-27-ai-chat-builder-implementation.md new file mode 100644 index 00000000..4e98505e --- /dev/null +++ b/docs/plans/2026-02-27-ai-chat-builder-implementation.md @@ -0,0 +1,2626 @@ +# AI Chat Builder Implementation Plan + +> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task. + +**Goal:** Build a conversational AI flow builder where the AI conducts a multi-phase interview as a senior MSP engineer, progressively building a TreeStructure through natural dialogue. + +**Architecture:** New chat-based builder coexists alongside the existing wizard builder. Backend adds a new model, service, and endpoint module in `core/` and `api/endpoints/`. Frontend adds a split-panel page with chat left / tree preview right, a new Zustand store, and new components in `components/ai-chat/`. Reuses existing AI provider abstraction, quota system, and tree validation. + +**Tech Stack:** FastAPI, SQLAlchemy 2.0 (async), Alembic, Pydantic v2 (backend); React 19, Zustand, TypeScript, Tailwind CSS (frontend). AI via existing `AIProvider` abstraction (Gemini/Anthropic). + +**Design doc:** `docs/plans/2026-02-27-ai-chat-builder-design.md` + +--- + +## Phase 1: Backend Foundation + +### Task 1: Database Model — `AIChatSession` + +**Files:** +- Create: `backend/app/models/ai_chat_session.py` +- Modify: `backend/app/models/__init__.py` (lines 29-30, 68-69) + +**Step 1: Create the model file** + +```python +# backend/app/models/ai_chat_session.py +"""AI Chat Builder session tracking. + +Stores conversational flow builder state across the multi-phase interview. +Sessions expire after 24 hours. +""" +import uuid +from datetime import datetime, timezone +from typing import Optional, Any + +from sqlalchemy import String, DateTime, ForeignKey, Integer +from sqlalchemy.orm import Mapped, mapped_column +from sqlalchemy.dialects.postgresql import UUID, JSONB + +from app.core.database import Base + + +class AIChatSession(Base): + __tablename__ = "ai_chat_sessions" + + id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), primary_key=True, default=uuid.uuid4 + ) + user_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), + ForeignKey("users.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + account_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), + ForeignKey("accounts.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) + status: Mapped[str] = mapped_column( + String(20), + nullable=False, + default="active", + comment="active | completed | abandoned", + ) + current_phase: Mapped[str] = mapped_column( + String(20), + nullable=False, + default="scoping", + comment="scoping | discovery | enrichment | review | generation", + ) + flow_type: Mapped[str] = mapped_column( + String(20), + nullable=False, + comment="troubleshooting | procedural", + ) + conversation_history: Mapped[list[dict[str, Any]]] = mapped_column( + JSONB, nullable=False, default=list + ) + working_tree: Mapped[Optional[dict[str, Any]]] = mapped_column( + JSONB, nullable=True + ) + tree_metadata: Mapped[dict[str, Any]] = mapped_column( + JSONB, nullable=False, default=dict + ) + provider_used: Mapped[Optional[str]] = mapped_column( + String(20), nullable=True + ) + message_count: Mapped[int] = mapped_column( + Integer, nullable=False, default=0 + ) + total_input_tokens: Mapped[int] = mapped_column( + Integer, nullable=False, default=0 + ) + total_output_tokens: Mapped[int] = mapped_column( + Integer, nullable=False, default=0 + ) + generated_tree_id: Mapped[Optional[uuid.UUID]] = mapped_column( + UUID(as_uuid=True), + ForeignKey("trees.id", ondelete="SET NULL"), + nullable=True, + ) + expires_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), nullable=False + ) + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), default=lambda: datetime.now(timezone.utc) + ) + updated_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(timezone.utc), + onupdate=lambda: datetime.now(timezone.utc), + ) +``` + +**Step 2: Register the model in `__init__.py`** + +Add import after line 30 (`from .ai_usage import AIUsage`): +```python +from .ai_chat_session import AIChatSession +``` + +Add to `__all__` after line 69 (`"AIUsage"`): +```python + "AIChatSession", +``` + +**Step 3: Create Alembic migration** + +Run: `cd backend && alembic revision -m "add ai_chat_sessions table"` + +Then write the migration manually (do NOT use `--autogenerate` — see CLAUDE.md lesson #17): + +```python +"""add ai_chat_sessions table""" +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects.postgresql import UUID, JSONB + + +def upgrade() -> None: + op.create_table( + "ai_chat_sessions", + sa.Column("id", UUID(as_uuid=True), primary_key=True), + sa.Column("user_id", UUID(as_uuid=True), sa.ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True), + sa.Column("account_id", UUID(as_uuid=True), sa.ForeignKey("accounts.id", ondelete="CASCADE"), nullable=False, index=True), + sa.Column("status", sa.String(20), nullable=False, server_default="active"), + sa.Column("current_phase", sa.String(20), nullable=False, server_default="scoping"), + sa.Column("flow_type", sa.String(20), nullable=False), + sa.Column("conversation_history", JSONB, nullable=False, server_default="[]"), + sa.Column("working_tree", JSONB, nullable=True), + sa.Column("tree_metadata", JSONB, nullable=False, server_default="{}"), + sa.Column("provider_used", sa.String(20), nullable=True), + sa.Column("message_count", sa.Integer, nullable=False, server_default="0"), + sa.Column("total_input_tokens", sa.Integer, nullable=False, server_default="0"), + sa.Column("total_output_tokens", sa.Integer, nullable=False, server_default="0"), + sa.Column("generated_tree_id", UUID(as_uuid=True), sa.ForeignKey("trees.id", ondelete="SET NULL"), nullable=True), + sa.Column("expires_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now()), + sa.Column("updated_at", sa.DateTime(timezone=True), server_default=sa.func.now()), + ) + + +def downgrade() -> None: + op.drop_table("ai_chat_sessions") +``` + +**Step 4: Run migration** + +Run: `cd backend && alembic upgrade head` + +Expected: Migration applied successfully, `ai_chat_sessions` table created. + +**Step 5: Commit** + +```bash +git add backend/app/models/ai_chat_session.py backend/app/models/__init__.py backend/alembic/versions/*ai_chat_sessions* +git commit -m "feat: add ai_chat_sessions database model and migration" +``` + +--- + +### Task 2: Pydantic Schemas + +**Files:** +- Create: `backend/app/schemas/ai_chat.py` + +**Step 1: Create the schemas file** + +```python +# backend/app/schemas/ai_chat.py +"""Pydantic schemas for the AI Chat Builder.""" +from typing import Any, Literal, Optional +from uuid import UUID + +from pydantic import BaseModel, Field + + +# ── Requests ── + + +class AIChatStartRequest(BaseModel): + """Start a new chat builder session.""" + + flow_type: Literal["troubleshooting", "procedural"] = Field( + ..., description="Type of flow to build" + ) + + +class AIChatMessageRequest(BaseModel): + """Send a user message in a chat session.""" + + content: str = Field(..., min_length=1, max_length=5000) + + +class AIChatImportRequest(BaseModel): + """Import generated tree with optional metadata overrides.""" + + name: Optional[str] = Field(None, min_length=1, max_length=255) + description: Optional[str] = Field(None, max_length=2000) + category_id: Optional[UUID] = None + tags: list[str] = Field(default_factory=list) + + +# ── Responses ── + + +class AIChatStartResponse(BaseModel): + """Response after creating a chat session.""" + + session_id: UUID + greeting: str + current_phase: str + + +class AIChatMessageResponse(BaseModel): + """Response after sending a message.""" + + content: str + current_phase: str + working_tree: Optional[dict[str, Any]] = None + tree_metadata: Optional[dict[str, Any]] = None + + +class AIChatSessionResponse(BaseModel): + """Full session state for resume.""" + + session_id: UUID + status: str + current_phase: str + flow_type: str + conversation_history: list[dict[str, Any]] + working_tree: Optional[dict[str, Any]] = None + tree_metadata: Optional[dict[str, Any]] = None + message_count: int + generated_tree: Optional[dict[str, Any]] = None + + +class AIChatGenerateResponse(BaseModel): + """Response with the final generated tree.""" + + tree_structure: dict[str, Any] + tree_metadata: dict[str, Any] + status: str + + +class AIChatImportResponse(BaseModel): + """Response after importing tree to editor.""" + + tree_id: UUID + tree_type: str +``` + +**Step 2: Commit** + +```bash +git add backend/app/schemas/ai_chat.py +git commit -m "feat: add Pydantic schemas for AI chat builder" +``` + +--- + +### Task 3: Chat Service — System Prompt & Conversation Loop + +**Files:** +- Create: `backend/app/core/ai_chat_service.py` + +This is the largest file. It contains the system prompt, response parsing, and the three main functions. + +**Step 1: Create the service file** + +```python +# backend/app/core/ai_chat_service.py +"""AI Chat Builder service. + +Manages the conversational flow builder: system prompt construction, +message exchange with AI provider, and response parsing (extracting +tree updates, phase transitions, and metadata from structured markers). +""" +import json +import logging +import re +import uuid +from datetime import datetime, timezone, timedelta +from typing import Any, Optional + +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from app.core.ai_provider import get_ai_provider +from app.core.ai_tree_validator import validate_generated_tree +from app.core.config import settings +from app.models.ai_chat_session import AIChatSession + +logger = logging.getLogger(__name__) + +# ── Cost estimation ── +COST_PER_INPUT_TOKEN = 1.0 / 1_000_000 +COST_PER_OUTPUT_TOKEN = 5.0 / 1_000_000 + +# ── Max messages per session ── +MAX_MESSAGES_FREE = 10 +MAX_MESSAGES_PAID = 25 + + +# ── System Prompt ── + +ROLE_PERSONA = """You are a senior IT engineer embedded in ResolutionFlow, a troubleshooting platform for MSP (Managed Service Provider) engineers. You have 15+ years of hands-on experience across Windows Server, Active Directory, Entra ID/Azure AD, Microsoft 365, networking (DNS, DHCP, routing, VPN, firewalls), virtualization (Hyper-V, VMware), security, backup/DR, and cloud infrastructure. + +Your job is to help engineers build troubleshooting decision trees by interviewing them about a problem space. You are NOT a generic assistant. You are a colleague who has seen these issues hundreds of times and knows the optimal diagnostic order. + +CRITICAL BEHAVIORS: +- Act as a senior engineer, not a chatbot. Use your domain knowledge to SUGGEST diagnostic steps, not just record what the user says. +- When the user describes a problem area, demonstrate understanding by naming specific sub-categories, common causes, and relevant tools. +- Challenge assumptions constructively: "Before we go down that path, have you considered checking X first? In my experience, that resolves 60% of these cases." +- Capture SPECIFIC commands with exact syntax. Not "check the service" but "Get-Service ADSync | Select-Object Status, StartType". +- Include expected outcomes for every action: what does success look like? +- Surface edge cases proactively: "What about multi-forest environments?" or "Does this change if they have conditional access policies?" +- Explain WHY the diagnostic order matters: "We check connectivity before auth because a network issue masquerades as an auth failure." +- Ask ONE focused question at a time. Do not overwhelm with multiple questions. +- Use plain, collegial language. Sound like a colleague, not a form.""" + +SCHEMA_CONTEXT = """ +TREESTRUCTURE SCHEMA — This is what you are building: + +The tree is a recursive JSON structure. Each node has a "type" field: + +1. decision — A diagnostic question with branching options + Required: id (string), type ("decision"), question (string), options (array), children (array) + Optional: help_text (string) + Each option: { id (string), label (string), next_node_id (string — must match a child's id) } + +2. action — A step the engineer performs + Required: id (string), type ("action"), title (string), description (string) + Optional: commands (string array — exact CLI/PowerShell syntax), expected_outcome (string), help_text (string), next_node_id (string — ID of the next node to navigate to) + +3. solution — A resolution endpoint + Required: id (string), type ("solution"), title (string), description (string) + Optional: resolution_steps (string array) + +STRUCTURAL RULES: +- Root node MUST be type "decision" +- Decision nodes contain their children in the "children" array +- Each decision option's next_node_id must reference a child node's id +- Action nodes use next_node_id to chain to the next step (NOT children) +- Solution nodes are terminal — no next_node_id or children +- All IDs must be unique strings (use descriptive slugs like "check-service-status") +""" + +INTERVIEW_PROTOCOL = """ +INTERVIEW PHASES — Follow this progression: + +PHASE 1 - SCOPING (current_phase: scoping): +Ask broad questions to understand the problem domain and scope: +- What type of issue is this flow for? +- Who is the target audience? (Tier 1 help desk, Tier 2, Tier 3?) +- What environment assumptions? (On-prem, hybrid, specific vendors?) +Demonstrate domain expertise immediately. If the user says "Azure AD Sync failures," show understanding: "Are you primarily seeing password hash sync issues, object attribute sync failures, or full directory sync errors?" +DO NOT emit [TREE_UPDATE] during scoping. You are still understanding the problem. + +PHASE 2 - DISCOVERY (current_phase: discovery): +Work through the troubleshooting logic branch by branch: +- Establish the first diagnostic question (the root decision node) +- For each branch, ask what the engineer would check next +- Suggest checks the user might not have considered +- Capture specific commands, tools, and procedures +EMIT [TREE_UPDATE] ONLY when you and the user have agreed on a concrete node — a decision with clear options, or an action with a specific command. If you are asking a question, you are NOT updating the tree. + +PHASE 3 - ENRICHMENT (current_phase: enrichment): +Circle back to enrich existing nodes: +- Add exact PowerShell/CLI commands with syntax +- Add help text with relevant documentation links +- Add expected outcomes for action nodes +- Suggest edge cases needing additional branches +EMIT [TREE_UPDATE] when enriching existing nodes or adding edge case branches. + +PHASE 4 - REVIEW (current_phase: review): +Present a summary: +- Total node count by type +- Text outline of the flow structure +- Flag any areas of uncertainty +- Offer chance to add/remove/modify branches +EMIT [TREE_UPDATE] only if the user requests structural changes. + +TRANSITION between phases by emitting [PHASE:phase_name] when the conversation naturally moves to the next stage. You decide when enough information has been gathered for each phase. +""" + +RESPONSE_FORMAT = """ +RESPONSE FORMAT: + +Your response is natural conversational text. When the tree structure changes, include structured markers that will be parsed by the system (the user will NOT see these markers): + +1. Tree update (only when structure changes — see phase rules above): +[TREE_UPDATE] +{...valid TreeStructure JSON...} +[/TREE_UPDATE] + +2. Phase transition (when moving to next phase): +[PHASE:discovery] + +3. Metadata capture (when you learn the flow's name, description, or tags): +[METADATA] +{"name": "...", "description": "...", "tags": ["..."]} +[/METADATA] + +IMPORTANT: +- Include [TREE_UPDATE] sparingly. Only when concrete nodes are established or modified. +- The tree update should be the COMPLETE working tree, not a diff. +- Always include conversational text OUTSIDE the markers — never respond with only markers. +""" + + +def _build_system_prompt(flow_type: str) -> str: + """Assemble the full system prompt for the chat builder.""" + flow_context = ( + "The user wants to build a TROUBLESHOOTING flow — a diagnostic decision tree " + "that guides engineers through symptom identification, diagnostic checks, and " + "resolution steps." + if flow_type == "troubleshooting" + else "The user wants to build a PROCEDURAL flow — a step-by-step process guide " + "with phases, checklists, and verification steps." + ) + + return f"{ROLE_PERSONA}\n\n{flow_context}\n\n{SCHEMA_CONTEXT}\n\n{INTERVIEW_PROTOCOL}\n\n{RESPONSE_FORMAT}" + + +def _strip_markdown_fences(text: str) -> str: + """Strip markdown code fences if the model wrapped its JSON response.""" + text = text.strip() + match = re.match(r"^```(?:json)?\s*([\s\S]*?)```$", text) + if match: + return match.group(1).strip() + return text + + +def _parse_ai_response(raw_response: str) -> dict[str, Any]: + """Parse structured markers from AI response. + + Returns dict with: + - content: str (conversational text with markers stripped) + - tree_update: dict | None (parsed TreeStructure JSON) + - phase: str | None (new phase name) + - metadata: dict | None (name, description, tags) + """ + result: dict[str, Any] = { + "content": raw_response, + "tree_update": None, + "phase": None, + "metadata": None, + } + + # Extract [TREE_UPDATE]...[/TREE_UPDATE] + tree_match = re.search( + r"\[TREE_UPDATE\]\s*([\s\S]*?)\s*\[/TREE_UPDATE\]", raw_response + ) + if tree_match: + try: + raw_json = _strip_markdown_fences(tree_match.group(1)) + result["tree_update"] = json.loads(raw_json) + except (json.JSONDecodeError, ValueError) as e: + logger.warning("Failed to parse tree update JSON: %s", e) + result["content"] = raw_response[: tree_match.start()] + raw_response[tree_match.end() :] + + # Extract [PHASE:name] + phase_match = re.search(r"\[PHASE:(\w+)\]", result["content"]) + if phase_match: + result["phase"] = phase_match.group(1) + result["content"] = result["content"][: phase_match.start()] + result["content"][phase_match.end() :] + + # Extract [METADATA]...[/METADATA] + meta_match = re.search( + r"\[METADATA\]\s*([\s\S]*?)\s*\[/METADATA\]", result["content"] + ) + if meta_match: + try: + raw_json = _strip_markdown_fences(meta_match.group(1)) + result["metadata"] = json.loads(raw_json) + except (json.JSONDecodeError, ValueError) as e: + logger.warning("Failed to parse metadata JSON: %s", e) + result["content"] = result["content"][: meta_match.start()] + result["content"][meta_match.end() :] + + # Clean up extra whitespace from marker removal + result["content"] = re.sub(r"\n{3,}", "\n\n", result["content"]).strip() + + return result + + +# ── Main Service Functions ── + + +async def start_chat_session( + flow_type: str, + user_id: uuid.UUID, + account_id: uuid.UUID, + db: AsyncSession, +) -> tuple[AIChatSession, str]: + """Create a chat session and return the AI's opening greeting. + + Returns (session, greeting_text). + """ + session = AIChatSession( + user_id=user_id, + account_id=account_id, + flow_type=flow_type, + expires_at=datetime.now(timezone.utc) + timedelta(hours=settings.AI_CONVERSATION_TTL_HOURS), + ) + db.add(session) + await db.flush() + + # Build system prompt and get opening message + system_prompt = _build_system_prompt(flow_type) + primer = ( + f"I want to build a {flow_type} flow. Help me get started." + ) + + provider = get_ai_provider() + provider_name = settings.AI_PROVIDER + + messages = [{"role": "user", "content": primer}] + response_text, input_tokens, output_tokens = await provider.generate_json( + system_prompt=system_prompt, + messages=messages, + max_tokens=1500, + ) + + # Parse response (greeting shouldn't have tree updates, but handle gracefully) + parsed = _parse_ai_response(response_text) + + # Store conversation history + now_iso = datetime.now(timezone.utc).isoformat() + session.conversation_history = [ + {"role": "user", "content": primer, "timestamp": now_iso, "hidden": True}, + {"role": "assistant", "content": parsed["content"], "timestamp": now_iso}, + ] + session.provider_used = provider_name + session.message_count = 1 + session.total_input_tokens = input_tokens + session.total_output_tokens = output_tokens + + if parsed["metadata"]: + session.tree_metadata = parsed["metadata"] + + return session, parsed["content"] + + +async def send_message( + session: AIChatSession, + user_message: str, + db: AsyncSession, +) -> tuple[str, Optional[dict], Optional[str], Optional[dict]]: + """Send a user message and get AI response. + + Returns (ai_content, working_tree_update, new_phase, metadata_update). + """ + system_prompt = _build_system_prompt(session.flow_type) + + # Build messages array from conversation history + now_iso = datetime.now(timezone.utc).isoformat() + history = list(session.conversation_history) + history.append({"role": "user", "content": user_message, "timestamp": now_iso}) + + # Convert to provider format (just role + content) + provider_messages = [ + {"role": msg["role"], "content": msg["content"]} + for msg in history + ] + + provider = get_ai_provider() + response_text, input_tokens, output_tokens = await provider.generate_json( + system_prompt=system_prompt, + messages=provider_messages, + max_tokens=2000, + ) + + parsed = _parse_ai_response(response_text) + + # Validate tree update if present + tree_update = parsed["tree_update"] + if tree_update: + errors = validate_generated_tree(tree_update) + if errors: + logger.warning("AI tree update failed validation: %s", errors) + tree_update = None # Silently discard invalid updates + + # Update session state + history.append({"role": "assistant", "content": parsed["content"], "timestamp": now_iso}) + session.conversation_history = history + session.message_count = session.message_count + 1 + session.total_input_tokens = session.total_input_tokens + input_tokens + session.total_output_tokens = session.total_output_tokens + output_tokens + + if tree_update: + session.working_tree = tree_update + + if parsed["phase"]: + valid_phases = {"scoping", "discovery", "enrichment", "review", "generation"} + if parsed["phase"] in valid_phases: + session.current_phase = parsed["phase"] + + if parsed["metadata"]: + merged = dict(session.tree_metadata) + merged.update(parsed["metadata"]) + session.tree_metadata = merged + + session.updated_at = datetime.now(timezone.utc) + + return parsed["content"], tree_update, parsed["phase"], parsed["metadata"] + + +async def generate_final_tree( + session: AIChatSession, + db: AsyncSession, +) -> tuple[dict[str, Any], dict[str, Any]]: + """Generate the final validated TreeStructure from the conversation. + + Returns (tree_structure, metadata). + Raises ValueError if generation fails after retry. + """ + system_prompt = _build_system_prompt(session.flow_type) + + # Build generation prompt from full conversation + provider_messages = [ + {"role": msg["role"], "content": msg["content"]} + for msg in session.conversation_history + ] + + generation_instruction = """Based on our entire conversation, generate the COMPLETE and FINAL TreeStructure JSON for this flow. + +Requirements: +- Include ALL branches, steps, and solutions we discussed +- Use descriptive node IDs (slugs, not UUIDs) +- Root node must be type "decision" +- Every decision option must have a valid next_node_id pointing to a child +- Every action node should have commands with exact syntax where discussed +- Every action node should have expected_outcome where discussed +- Solution nodes should have resolution_steps +- Respond with ONLY the JSON — no conversational text, no markdown fences + +Also provide metadata as a separate JSON object after the tree: +[METADATA] +{"name": "...", "description": "...", "tags": ["..."]} +[/METADATA]""" + + provider_messages.append({"role": "user", "content": generation_instruction}) + + provider = get_ai_provider() + + for attempt in range(2): # One try + one retry + response_text, input_tokens, output_tokens = await provider.generate_json( + system_prompt=system_prompt, + messages=provider_messages, + max_tokens=8000, + ) + + session.total_input_tokens = session.total_input_tokens + input_tokens + session.total_output_tokens = session.total_output_tokens + output_tokens + + # Extract metadata first + parsed = _parse_ai_response(response_text) + metadata = parsed["metadata"] or dict(session.tree_metadata) + + # Parse tree JSON — could be in tree_update marker or raw + tree = parsed["tree_update"] + if not tree: + try: + raw = _strip_markdown_fences(parsed["content"]) + tree = json.loads(raw) + except (json.JSONDecodeError, ValueError): + pass + + if not tree: + if attempt == 0: + provider_messages.append({"role": "assistant", "content": response_text}) + provider_messages.append({ + "role": "user", + "content": "That response was not valid JSON. Please respond with ONLY the TreeStructure JSON object, starting with { and ending with }. No markdown fences, no explanatory text.", + }) + continue + raise ValueError("AI failed to produce valid JSON after retry") + + errors = validate_generated_tree(tree) + if errors: + if attempt == 0: + provider_messages.append({"role": "assistant", "content": response_text}) + correction = ( + f"The tree has validation errors: {'; '.join(errors)}. " + "Please fix these issues and respond with the corrected JSON only." + ) + provider_messages.append({"role": "user", "content": correction}) + continue + raise ValueError(f"Generated tree failed validation: {'; '.join(errors)}") + + # Success + session.working_tree = tree + session.tree_metadata = metadata + session.current_phase = "generation" + session.updated_at = datetime.now(timezone.utc) + + return tree, metadata + + raise ValueError("AI failed to generate a valid tree") + + +async def get_chat_session( + session_id: uuid.UUID, + user_id: uuid.UUID, + db: AsyncSession, +) -> AIChatSession: + """Get a chat session, validating ownership and expiry. + + Raises HTTPException on not found, forbidden, or expired. + """ + from fastapi import HTTPException, status + + result = await db.execute( + select(AIChatSession).where(AIChatSession.id == session_id) + ) + session = result.scalar_one_or_none() + + if not session: + raise HTTPException(status_code=404, detail="Chat session not found") + + if session.user_id != user_id: + raise HTTPException(status_code=403, detail="Access denied") + + if session.expires_at < datetime.now(timezone.utc): + session.status = "abandoned" + await db.flush() + raise HTTPException(status_code=410, detail="Chat session has expired") + + return session +``` + +**Step 2: Commit** + +```bash +git add backend/app/core/ai_chat_service.py +git commit -m "feat: add AI chat builder service with system prompt and conversation loop" +``` + +--- + +### Task 4: API Endpoints + +**Files:** +- Create: `backend/app/api/endpoints/ai_chat.py` +- Modify: `backend/app/api/router.py` (lines 9, 40) + +**Step 1: Create the endpoints file** + +```python +# backend/app/api/endpoints/ai_chat.py +"""AI Chat Builder endpoints. + +Conversational flow builder: + POST /ai/chat/sessions — Start session, get AI greeting + POST /ai/chat/sessions/{id}/messages — Send message, get AI response + GET /ai/chat/sessions/{id} — Get session state (for resume) + POST /ai/chat/sessions/{id}/generate — Generate final TreeStructure + POST /ai/chat/sessions/{id}/import — Create Tree from generated structure + DELETE /ai/chat/sessions/{id} — Abandon session +""" +import logging +from typing import Annotated +from uuid import UUID + +from fastapi import APIRouter, Depends, HTTPException, Request, status +from sqlalchemy.ext.asyncio import AsyncSession + +from app.core.rate_limit import limiter +from app.api.deps import get_current_active_user, get_db, require_engineer_or_admin +from app.core.config import settings +from app.core.ai_chat_service import ( + start_chat_session, + send_message, + generate_final_tree, + get_chat_session, + MAX_MESSAGES_FREE, + MAX_MESSAGES_PAID, +) +from app.core.ai_quota_service import check_ai_quota, record_ai_usage, get_user_plan +from app.models.user import User +from app.models.tree import Tree +from app.schemas.ai_chat import ( + AIChatStartRequest, + AIChatStartResponse, + AIChatMessageRequest, + AIChatMessageResponse, + AIChatSessionResponse, + AIChatGenerateResponse, + AIChatImportRequest, + AIChatImportResponse, +) + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/ai/chat", tags=["ai-chat-builder"]) + + +def _require_ai_enabled() -> None: + if not settings.ai_enabled: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="AI is not configured. Set GOOGLE_AI_API_KEY or ANTHROPIC_API_KEY.", + ) + + +@router.post("/sessions", response_model=AIChatStartResponse, status_code=201) +@limiter.limit("10/minute") +async def create_session( + request: Request, + data: AIChatStartRequest, + current_user: Annotated[User, Depends(get_current_active_user)], + db: Annotated[AsyncSession, Depends(get_db)], + _: None = Depends(require_engineer_or_admin), +): + """Start a new AI chat builder session.""" + _require_ai_enabled() + + # Check quota + allowed, quota_status = await check_ai_quota( + user_id=current_user.id, + account_id=current_user.account_id, + db=db, + billing_anchor=current_user.ai_billing_cycle_anchor_at, + is_super_admin=current_user.is_super_admin, + ) + if not allowed: + reset_key = ( + "daily_reset_at" + if quota_status.get("deny_reason") == "daily" + else "monthly_reset_at" + ) + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail={ + "message": f"AI build limit exceeded ({quota_status['deny_reason']})", + "reset_at": quota_status.get(reset_key), + "quota": quota_status, + }, + ) + + plan = await get_user_plan(current_user.account_id, db) + + try: + session, greeting = await start_chat_session( + flow_type=data.flow_type, + user_id=current_user.id, + account_id=current_user.account_id, + db=db, + ) + except Exception as e: + logger.exception("AI chat session start failed: %s", e) + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail=f"AI provider error ({type(e).__name__}). Please try again.", + ) + + # Record usage for the greeting + await record_ai_usage( + user_id=current_user.id, + account_id=current_user.account_id, + conversation_id=session.id, + generation_type="chat_message", + tier=plan, + input_tokens=session.total_input_tokens, + output_tokens=session.total_output_tokens, + estimated_cost=( + session.total_input_tokens * 1.0 / 1_000_000 + + session.total_output_tokens * 5.0 / 1_000_000 + ), + succeeded=True, + counts_toward_quota=False, + error_code=None, + extra_data={"phase": "scoping"}, + db=db, + ) + + await db.commit() + + return AIChatStartResponse( + session_id=session.id, + greeting=greeting, + current_phase=session.current_phase, + ) + + +@router.post("/sessions/{session_id}/messages", response_model=AIChatMessageResponse) +@limiter.limit("10/minute") +async def post_message( + request: Request, + session_id: UUID, + data: AIChatMessageRequest, + current_user: Annotated[User, Depends(get_current_active_user)], + db: Annotated[AsyncSession, Depends(get_db)], + _: None = Depends(require_engineer_or_admin), +): + """Send a user message and get AI response.""" + _require_ai_enabled() + + session = await get_chat_session(session_id, current_user.id, db) + + if session.status != "active": + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Session is {session.status}, cannot send messages", + ) + + # Per-session message limit + plan = await get_user_plan(current_user.account_id, db) + max_messages = MAX_MESSAGES_PAID if plan != "free" else MAX_MESSAGES_FREE + if current_user.is_super_admin: + max_messages = 999 + + if session.message_count >= max_messages: + raise HTTPException( + status_code=status.HTTP_429_TOO_MANY_REQUESTS, + detail=f"Maximum messages per session reached ({max_messages}). Generate your tree or start a new session.", + ) + + prev_input = session.total_input_tokens + prev_output = session.total_output_tokens + + try: + ai_content, tree_update, new_phase, metadata = await send_message( + session, data.content, db + ) + except Exception as e: + logger.exception("AI chat message failed: %s", e) + await record_ai_usage( + user_id=current_user.id, + account_id=current_user.account_id, + conversation_id=session.id, + generation_type="chat_message", + tier=plan, + input_tokens=0, + output_tokens=0, + estimated_cost=0, + succeeded=False, + counts_toward_quota=False, + error_code=type(e).__name__, + extra_data=None, + db=db, + ) + await db.commit() + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail=f"AI provider error ({type(e).__name__}). Please try again.", + ) + + # Record usage + input_delta = session.total_input_tokens - prev_input + output_delta = session.total_output_tokens - prev_output + await record_ai_usage( + user_id=current_user.id, + account_id=current_user.account_id, + conversation_id=session.id, + generation_type="chat_message", + tier=plan, + input_tokens=input_delta, + output_tokens=output_delta, + estimated_cost=( + input_delta * 1.0 / 1_000_000 + + output_delta * 5.0 / 1_000_000 + ), + succeeded=True, + counts_toward_quota=False, + error_code=None, + extra_data={"phase": session.current_phase}, + db=db, + ) + + await db.commit() + + return AIChatMessageResponse( + content=ai_content, + current_phase=session.current_phase, + working_tree=session.working_tree, + tree_metadata=session.tree_metadata if session.tree_metadata else None, + ) + + +@router.get("/sessions/{session_id}", response_model=AIChatSessionResponse) +async def get_session( + session_id: UUID, + current_user: Annotated[User, Depends(get_current_active_user)], + db: Annotated[AsyncSession, Depends(get_db)], +): + """Get full session state for resume after page reload.""" + session = await get_chat_session(session_id, current_user.id, db) + + # Filter out hidden messages (like the primer) + visible_history = [ + msg for msg in session.conversation_history + if not msg.get("hidden") + ] + + return AIChatSessionResponse( + session_id=session.id, + status=session.status, + current_phase=session.current_phase, + flow_type=session.flow_type, + conversation_history=visible_history, + working_tree=session.working_tree, + tree_metadata=session.tree_metadata if session.tree_metadata else None, + message_count=session.message_count, + generated_tree=session.working_tree if session.status == "completed" else None, + ) + + +@router.post("/sessions/{session_id}/generate", response_model=AIChatGenerateResponse) +@limiter.limit("10/minute") +async def generate_tree( + request: Request, + session_id: UUID, + current_user: Annotated[User, Depends(get_current_active_user)], + db: Annotated[AsyncSession, Depends(get_db)], + _: None = Depends(require_engineer_or_admin), +): + """Generate final TreeStructure JSON from conversation.""" + _require_ai_enabled() + + session = await get_chat_session(session_id, current_user.id, db) + + # If already generated, return cached result + if session.status == "completed" and session.working_tree: + return AIChatGenerateResponse( + tree_structure=session.working_tree, + tree_metadata=session.tree_metadata, + status="completed", + ) + + if session.status != "active": + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Session is {session.status}, cannot generate", + ) + + plan = await get_user_plan(current_user.account_id, db) + prev_input = session.total_input_tokens + prev_output = session.total_output_tokens + + try: + tree_structure, metadata = await generate_final_tree(session, db) + except ValueError as e: + await record_ai_usage( + user_id=current_user.id, + account_id=current_user.account_id, + conversation_id=session.id, + generation_type="chat_generate", + tier=plan, + input_tokens=session.total_input_tokens - prev_input, + output_tokens=session.total_output_tokens - prev_output, + estimated_cost=0, + succeeded=False, + counts_toward_quota=False, + error_code="invalid_output", + extra_data={"error": str(e)}, + db=db, + ) + await db.commit() + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail=f"Tree generation failed: {e}", + ) + except Exception as e: + logger.exception("AI chat generate failed: %s", e) + input_delta = session.total_input_tokens - prev_input + output_delta = session.total_output_tokens - prev_output + await record_ai_usage( + user_id=current_user.id, + account_id=current_user.account_id, + conversation_id=session.id, + generation_type="chat_generate", + tier=plan, + input_tokens=input_delta, + output_tokens=output_delta, + estimated_cost=0, + succeeded=False, + counts_toward_quota=False, + error_code=type(e).__name__, + extra_data={"error": str(e)}, + db=db, + ) + await db.commit() + + # Distinguish timeout from other errors + error_name = type(e).__name__ + if "timeout" in error_name.lower() or "Timeout" in str(e): + raise HTTPException( + status_code=status.HTTP_504_GATEWAY_TIMEOUT, + detail="Tree generation timed out. Please try again.", + ) + raise HTTPException( + status_code=status.HTTP_502_BAD_GATEWAY, + detail=f"AI provider error ({error_name}). Please try again.", + ) + + # Record successful quota-consuming usage + input_delta = session.total_input_tokens - prev_input + output_delta = session.total_output_tokens - prev_output + await record_ai_usage( + user_id=current_user.id, + account_id=current_user.account_id, + conversation_id=session.id, + generation_type="chat_generate", + tier=plan, + input_tokens=input_delta, + output_tokens=output_delta, + estimated_cost=( + input_delta * 1.0 / 1_000_000 + + output_delta * 5.0 / 1_000_000 + ), + succeeded=True, + counts_toward_quota=True, + error_code=None, + extra_data=None, + db=db, + ) + + session.status = "completed" + await db.commit() + + return AIChatGenerateResponse( + tree_structure=tree_structure, + tree_metadata=metadata, + status="completed", + ) + + +@router.post("/sessions/{session_id}/import", response_model=AIChatImportResponse) +@limiter.limit("10/minute") +async def import_tree( + request: Request, + session_id: UUID, + data: AIChatImportRequest, + current_user: Annotated[User, Depends(get_current_active_user)], + db: Annotated[AsyncSession, Depends(get_db)], + _: None = Depends(require_engineer_or_admin), +): + """Create a Tree record from the generated tree structure.""" + session = await get_chat_session(session_id, current_user.id, db) + + if session.status != "completed" or not session.working_tree: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Session must be completed with a generated tree before importing", + ) + + # Already imported? + if session.generated_tree_id: + return AIChatImportResponse( + tree_id=session.generated_tree_id, + tree_type=session.flow_type, + ) + + metadata = session.tree_metadata or {} + tree = Tree( + name=data.name or metadata.get("name", "AI-Generated Flow"), + description=data.description or metadata.get("description", ""), + tree_type=session.flow_type, + tree_structure=session.working_tree, + author_id=current_user.id, + account_id=current_user.account_id, + category_id=data.category_id, + is_public=False, + ) + db.add(tree) + await db.flush() + + session.generated_tree_id = tree.id + await db.commit() + + return AIChatImportResponse( + tree_id=tree.id, + tree_type=session.flow_type, + ) + + +@router.delete("/sessions/{session_id}", status_code=204) +async def abandon_session( + session_id: UUID, + current_user: Annotated[User, Depends(get_current_active_user)], + db: Annotated[AsyncSession, Depends(get_db)], +): + """Abandon a chat session.""" + session = await get_chat_session(session_id, current_user.id, db) + session.status = "abandoned" + await db.commit() +``` + +**Step 2: Register in router** + +In `backend/app/api/router.py`, add after line 9 (`from app.api.endpoints import ai_fix`): +```python +from app.api.endpoints import ai_chat +``` + +Add after line 40 (`api_router.include_router(ai_fix.router)`): +```python +api_router.include_router(ai_chat.router) +``` + +**Step 3: Update quota service daily limit** + +In `backend/app/core/ai_quota_service.py`, change line 118: +```python +# Before: + AIUsage.generation_type.in_(["scaffold", "branch_detail"]), +# After: + AIUsage.generation_type.in_(["scaffold", "branch_detail", "chat_message", "chat_generate"]), +``` + +**Step 4: Commit** + +```bash +git add backend/app/api/endpoints/ai_chat.py backend/app/api/router.py backend/app/core/ai_quota_service.py +git commit -m "feat: add AI chat builder endpoints and update quota service" +``` + +--- + +### Task 5: Backend Tests + +**Files:** +- Create: `backend/tests/test_ai_chat.py` + +**Step 1: Write integration tests** + +```python +# backend/tests/test_ai_chat.py +"""Integration tests for AI Chat Builder endpoints. + +These tests mock the AI provider to avoid real API calls. +""" +import pytest +from unittest.mock import AsyncMock, patch + +pytestmark = pytest.mark.asyncio + + +@pytest.fixture +def mock_ai_provider(): + """Mock AI provider that returns realistic responses.""" + provider = AsyncMock() + provider.generate_json = AsyncMock(return_value=( + "Great question! Let's build a troubleshooting flow for DNS resolution issues. " + "To start, I need to understand the scope.\n\n" + "Who is the target audience for this flow? Are we targeting:\n" + "- Tier 1 help desk (basic checks only)\n" + "- Tier 2 desktop support (intermediate diagnostics)\n" + "- Tier 3 systems engineers (deep DNS troubleshooting)\n\n" + "[PHASE:scoping]", + 500, # input tokens + 200, # output tokens + )) + return provider + + +async def test_create_chat_session(client, auth_headers, mock_ai_provider): + """POST /ai/chat/sessions creates a session and returns AI greeting.""" + with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): + resp = await client.post( + "/api/v1/ai/chat/sessions", + json={"flow_type": "troubleshooting"}, + headers=auth_headers, + ) + + assert resp.status_code == 201 + data = resp.json() + assert "session_id" in data + assert "greeting" in data + assert data["current_phase"] == "scoping" + assert len(data["greeting"]) > 0 + + +async def test_send_message(client, auth_headers, mock_ai_provider): + """POST /ai/chat/sessions/{id}/messages returns AI response.""" + # Create session first + with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): + create_resp = await client.post( + "/api/v1/ai/chat/sessions", + json={"flow_type": "troubleshooting"}, + headers=auth_headers, + ) + session_id = create_resp.json()["session_id"] + + # Mock response with tree update + mock_ai_provider.generate_json = AsyncMock(return_value=( + "Good, targeting Tier 2 support. Let's start with the first diagnostic question.\n\n" + "The root question should be: 'What DNS symptom is the user experiencing?'\n\n" + "[TREE_UPDATE]\n" + '{"id": "root", "type": "decision", "question": "What DNS symptom is the user experiencing?", ' + '"options": [{"id": "opt-1", "label": "Cannot resolve any domains", "next_node_id": "check-dns-service"}], ' + '"children": [{"id": "check-dns-service", "type": "action", "title": "Check DNS Client Service", ' + '"description": "Verify the DNS Client service is running", ' + '"commands": ["Get-Service Dnscache | Select-Object Status, StartType"], ' + '"expected_outcome": "Service should show Running status"}]}\n' + "[/TREE_UPDATE]\n\n" + "[PHASE:discovery]", + 800, + 400, + )) + + with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): + resp = await client.post( + f"/api/v1/ai/chat/sessions/{session_id}/messages", + json={"content": "This is for Tier 2 support, hybrid environment with on-prem AD."}, + headers=auth_headers, + ) + + assert resp.status_code == 200 + data = resp.json() + assert "content" in data + assert data["current_phase"] == "discovery" + assert data["working_tree"] is not None + assert data["working_tree"]["type"] == "decision" + # Markers should be stripped from content + assert "[TREE_UPDATE]" not in data["content"] + assert "[PHASE:" not in data["content"] + + +async def test_get_session(client, auth_headers, mock_ai_provider): + """GET /ai/chat/sessions/{id} returns full session state.""" + with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): + create_resp = await client.post( + "/api/v1/ai/chat/sessions", + json={"flow_type": "troubleshooting"}, + headers=auth_headers, + ) + session_id = create_resp.json()["session_id"] + + resp = await client.get( + f"/api/v1/ai/chat/sessions/{session_id}", + headers=auth_headers, + ) + + assert resp.status_code == 200 + data = resp.json() + assert data["session_id"] == session_id + assert data["status"] == "active" + assert data["flow_type"] == "troubleshooting" + # Hidden primer message should be filtered out + assert all(msg.get("role") == "assistant" or not msg.get("hidden") for msg in data["conversation_history"]) + + +async def test_abandon_session(client, auth_headers, mock_ai_provider): + """DELETE /ai/chat/sessions/{id} sets status to abandoned.""" + with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): + create_resp = await client.post( + "/api/v1/ai/chat/sessions", + json={"flow_type": "troubleshooting"}, + headers=auth_headers, + ) + session_id = create_resp.json()["session_id"] + + resp = await client.delete( + f"/api/v1/ai/chat/sessions/{session_id}", + headers=auth_headers, + ) + assert resp.status_code == 204 + + # Verify session is abandoned + get_resp = await client.get( + f"/api/v1/ai/chat/sessions/{session_id}", + headers=auth_headers, + ) + assert get_resp.json()["status"] == "abandoned" + + +async def test_message_limit_enforced(client, auth_headers, mock_ai_provider): + """Per-session message limit should return 429.""" + with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): + create_resp = await client.post( + "/api/v1/ai/chat/sessions", + json={"flow_type": "troubleshooting"}, + headers=auth_headers, + ) + session_id = create_resp.json()["session_id"] + + # Patch the session to have max messages already + with patch("app.api.endpoints.ai_chat.MAX_MESSAGES_FREE", 1): + with patch("app.core.ai_chat_service.get_ai_provider", return_value=mock_ai_provider): + resp = await client.post( + f"/api/v1/ai/chat/sessions/{session_id}/messages", + json={"content": "test message"}, + headers=auth_headers, + ) + assert resp.status_code == 429 + + +async def test_session_not_found(client, auth_headers): + """Accessing nonexistent session returns 404.""" + import uuid + fake_id = str(uuid.uuid4()) + resp = await client.get( + f"/api/v1/ai/chat/sessions/{fake_id}", + headers=auth_headers, + ) + assert resp.status_code == 404 + + +async def test_ai_disabled_returns_503(client, auth_headers): + """When AI is not configured, endpoints return 503.""" + with patch("app.api.endpoints.ai_chat.settings") as mock_settings: + mock_settings.ai_enabled = False + resp = await client.post( + "/api/v1/ai/chat/sessions", + json={"flow_type": "troubleshooting"}, + headers=auth_headers, + ) + assert resp.status_code == 503 +``` + +**Step 2: Run tests** + +Run: `cd backend && python -m pytest tests/test_ai_chat.py -v --override-ini="addopts="` + +Expected: All tests pass. If fixture issues, check `conftest.py` for available fixtures (CLAUDE.md lesson #21). + +**Step 3: Commit** + +```bash +git add backend/tests/test_ai_chat.py +git commit -m "test: add integration tests for AI chat builder endpoints" +``` + +--- + +## Phase 2: Frontend Chat UI + +### Task 6: TypeScript Types + +**Files:** +- Create: `frontend/src/types/ai-chat.ts` +- Modify: `frontend/src/types/index.ts` (after line 54) + +**Step 1: Create the types file** + +```typescript +// frontend/src/types/ai-chat.ts + +export type InterviewPhase = 'scoping' | 'discovery' | 'enrichment' | 'review' | 'generation' + +export interface ChatMessage { + role: 'user' | 'assistant' + content: string + timestamp: string +} + +export interface AIChatStartResponse { + session_id: string + greeting: string + current_phase: InterviewPhase +} + +export interface AIChatMessageResponse { + content: string + current_phase: InterviewPhase + working_tree: Record | null + tree_metadata: Record | null +} + +export interface AIChatSessionResponse { + session_id: string + status: 'active' | 'completed' | 'abandoned' + current_phase: InterviewPhase + flow_type: 'troubleshooting' | 'procedural' + conversation_history: ChatMessage[] + working_tree: Record | null + tree_metadata: Record | null + message_count: number + generated_tree: Record | null +} + +export interface AIChatGenerateResponse { + tree_structure: Record + tree_metadata: Record + status: string +} + +export interface AIChatImportResponse { + tree_id: string + tree_type: string +} +``` + +**Step 2: Export from types/index.ts** + +Add after line 54 (the `ai-fix` export block): +```typescript + +export type { + InterviewPhase, + ChatMessage, + AIChatStartResponse, + AIChatMessageResponse, + AIChatSessionResponse, + AIChatGenerateResponse, + AIChatImportResponse, +} from './ai-chat' +``` + +**Step 3: Commit** + +```bash +git add frontend/src/types/ai-chat.ts frontend/src/types/index.ts +git commit -m "feat: add TypeScript types for AI chat builder" +``` + +--- + +### Task 7: API Client Module + +**Files:** +- Create: `frontend/src/api/aiChat.ts` +- Modify: `frontend/src/api/index.ts` (after line 19) + +**Step 1: Create the API client** + +```typescript +// frontend/src/api/aiChat.ts +import { apiClient } from './client' +import type { + AIChatStartResponse, + AIChatMessageResponse, + AIChatSessionResponse, + AIChatGenerateResponse, + AIChatImportResponse, +} from '@/types' + +export const aiChatApi = { + startSession: async (flowType: 'troubleshooting' | 'procedural'): Promise => { + const { data } = await apiClient.post('/ai/chat/sessions', { flow_type: flowType }) + return data + }, + + sendMessage: async (sessionId: string, content: string): Promise => { + const { data } = await apiClient.post(`/ai/chat/sessions/${sessionId}/messages`, { content }) + return data + }, + + getSession: async (sessionId: string): Promise => { + const { data } = await apiClient.get(`/ai/chat/sessions/${sessionId}`) + return data + }, + + generateTree: async (sessionId: string): Promise => { + const { data } = await apiClient.post(`/ai/chat/sessions/${sessionId}/generate`) + return data + }, + + importTree: async ( + sessionId: string, + params?: { name?: string; description?: string; category_id?: string; tags?: string[] } + ): Promise => { + const { data } = await apiClient.post(`/ai/chat/sessions/${sessionId}/import`, params || {}) + return data + }, + + abandonSession: async (sessionId: string): Promise => { + await apiClient.delete(`/ai/chat/sessions/${sessionId}`) + }, +} + +export default aiChatApi +``` + +**Step 2: Export from api/index.ts** + +Add after line 19 (`export { default as aiBuilderApi } from './aiBuilder'`): +```typescript +export { default as aiChatApi } from './aiChat' +``` + +**Step 3: Commit** + +```bash +git add frontend/src/api/aiChat.ts frontend/src/api/index.ts +git commit -m "feat: add API client for AI chat builder" +``` + +--- + +### Task 8: Zustand Store + +**Files:** +- Create: `frontend/src/store/aiChatStore.ts` + +**Step 1: Create the store** + +```typescript +// frontend/src/store/aiChatStore.ts +import { create } from 'zustand' +import { aiChatApi } from '@/api/aiChat' +import type { + ChatMessage, + InterviewPhase, + TreeStructure, +} from '@/types' + +interface TreeMetadata { + name?: string + description?: string + tags?: string[] + category_id?: string +} + +interface AIChatState { + // Session + sessionId: string | null + status: 'idle' | 'active' | 'completed' | 'abandoned' + currentPhase: InterviewPhase + flowType: 'troubleshooting' | 'procedural' | null + + // Conversation + messages: ChatMessage[] + isResponding: boolean + + // Progressive tree + workingTree: TreeStructure | null + treeMetadata: TreeMetadata | null + + // Final generation + generatedTree: TreeStructure | null + isGenerating: boolean + importedTreeId: string | null + + // Error + error: string | null + + // Actions + startSession: (flowType: 'troubleshooting' | 'procedural') => Promise + sendMessage: (content: string) => Promise + generateTree: () => Promise + importToEditor: (params?: { name?: string; description?: string; category_id?: string; tags?: string[] }) => Promise + abandonSession: () => Promise + resumeSession: (sessionId: string) => Promise + reset: () => void +} + +const initialState = { + sessionId: null, + status: 'idle' as const, + currentPhase: 'scoping' as InterviewPhase, + flowType: null, + messages: [], + isResponding: false, + workingTree: null, + treeMetadata: null, + generatedTree: null, + isGenerating: false, + importedTreeId: null, + error: null, +} + +export const useAIChatStore = create((set, get) => ({ + ...initialState, + + startSession: async (flowType) => { + set({ ...initialState, status: 'active', flowType, isResponding: true, error: null }) + + try { + const response = await aiChatApi.startSession(flowType) + set({ + sessionId: response.session_id, + currentPhase: response.current_phase, + messages: [{ + role: 'assistant', + content: response.greeting, + timestamp: new Date().toISOString(), + }], + isResponding: false, + }) + } catch (e: unknown) { + const message = e instanceof Error ? e.message : 'Failed to start session' + set({ error: message, isResponding: false, status: 'idle' }) + } + }, + + sendMessage: async (content) => { + const { sessionId, messages } = get() + if (!sessionId) return + + const userMessage: ChatMessage = { + role: 'user', + content, + timestamp: new Date().toISOString(), + } + + set({ + messages: [...messages, userMessage], + isResponding: true, + error: null, + }) + + try { + const response = await aiChatApi.sendMessage(sessionId, content) + const aiMessage: ChatMessage = { + role: 'assistant', + content: response.content, + timestamp: new Date().toISOString(), + } + + set((state) => ({ + messages: [...state.messages, aiMessage], + currentPhase: response.current_phase, + workingTree: response.working_tree as TreeStructure | null ?? state.workingTree, + treeMetadata: response.tree_metadata as TreeMetadata | null ?? state.treeMetadata, + isResponding: false, + })) + } catch (e: unknown) { + const message = e instanceof Error ? e.message : 'Failed to send message' + set({ error: message, isResponding: false }) + } + }, + + generateTree: async () => { + const { sessionId } = get() + if (!sessionId) return + + set({ isGenerating: true, error: null }) + + try { + const response = await aiChatApi.generateTree(sessionId) + set({ + generatedTree: response.tree_structure as TreeStructure, + workingTree: response.tree_structure as TreeStructure, + treeMetadata: response.tree_metadata as TreeMetadata, + status: 'completed', + isGenerating: false, + }) + } catch (e: unknown) { + const message = e instanceof Error ? e.message : 'Failed to generate tree' + set({ error: message, isGenerating: false }) + } + }, + + importToEditor: async (params) => { + const { sessionId } = get() + if (!sessionId) throw new Error('No active session') + + const response = await aiChatApi.importTree(sessionId, params) + set({ importedTreeId: response.tree_id }) + return response.tree_id + }, + + abandonSession: async () => { + const { sessionId } = get() + if (!sessionId) return + + try { + await aiChatApi.abandonSession(sessionId) + } catch { + // Best effort — session may have already expired + } + set({ ...initialState }) + }, + + resumeSession: async (sessionId) => { + set({ isResponding: true, error: null }) + + try { + const session = await aiChatApi.getSession(sessionId) + set({ + sessionId: session.session_id, + status: session.status === 'active' ? 'active' : session.status as 'completed' | 'abandoned', + currentPhase: session.current_phase, + flowType: session.flow_type, + messages: session.conversation_history as ChatMessage[], + workingTree: session.working_tree as TreeStructure | null, + treeMetadata: session.tree_metadata as TreeMetadata | null, + generatedTree: session.generated_tree as TreeStructure | null, + isResponding: false, + }) + } catch (e: unknown) { + const message = e instanceof Error ? e.message : 'Failed to resume session' + set({ error: message, isResponding: false }) + } + }, + + reset: () => set({ ...initialState }), +})) +``` + +**Step 2: Commit** + +```bash +git add frontend/src/store/aiChatStore.ts +git commit -m "feat: add Zustand store for AI chat builder" +``` + +--- + +### Task 9: Chat Components + +**Files:** +- Create: `frontend/src/components/ai-chat/ChatMessage.tsx` +- Create: `frontend/src/components/ai-chat/ChatInput.tsx` +- Create: `frontend/src/components/ai-chat/ChatPanel.tsx` +- Create: `frontend/src/components/ai-chat/PhaseIndicator.tsx` +- Create: `frontend/src/components/ai-chat/ChatToolbar.tsx` +- Create: `frontend/src/components/ai-chat/EmptyPreview.tsx` +- Create: `frontend/src/components/ai-chat/StaticTreePreview.tsx` + +This task creates all 7 components. Each is small and focused. + +**Step 1: Create `ChatMessage.tsx`** + +```typescript +// frontend/src/components/ai-chat/ChatMessage.tsx +import { Bot, User } from 'lucide-react' +import { MarkdownContent } from '@/components/ui/MarkdownContent' +import { cn } from '@/lib/utils' +import type { ChatMessage as ChatMessageType } from '@/types' + +interface ChatMessageProps { + message: ChatMessageType +} + +export function ChatMessage({ message }: ChatMessageProps) { + const isAI = message.role === 'assistant' + + return ( +
+ {isAI && ( +
+ +
+ )} +
+ {isAI ? ( + + ) : ( +

{message.content}

+ )} +
+ {!isAI && ( +
+ +
+ )} +
+ ) +} +``` + +**Step 2: Create `ChatInput.tsx`** + +```typescript +// frontend/src/components/ai-chat/ChatInput.tsx +import { useState, useRef, useCallback } from 'react' +import { Send } from 'lucide-react' +import { cn } from '@/lib/utils' + +interface ChatInputProps { + onSend: (content: string) => void + disabled?: boolean + placeholder?: string +} + +export function ChatInput({ onSend, disabled, placeholder = 'Type a message...' }: ChatInputProps) { + const [value, setValue] = useState('') + const textareaRef = useRef(null) + + const handleSend = useCallback(() => { + const trimmed = value.trim() + if (!trimmed || disabled) return + onSend(trimmed) + setValue('') + if (textareaRef.current) { + textareaRef.current.style.height = 'auto' + } + }, [value, disabled, onSend]) + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === 'Enter' && !e.shiftKey) { + e.preventDefault() + handleSend() + } + } + + const handleInput = () => { + if (textareaRef.current) { + textareaRef.current.style.height = 'auto' + textareaRef.current.style.height = Math.min(textareaRef.current.scrollHeight, 160) + 'px' + } + } + + return ( +
+