From ccd178b02e7e9f174a036d3729120223e72157bf Mon Sep 17 00:00:00 2001 From: chihlasm Date: Fri, 27 Feb 2026 03:42:10 -0500 Subject: [PATCH 01/19] feat: add ai_chat_sessions database model and migration Co-Authored-By: Claude Opus 4.6 --- ...e2d81e82ea5e_add_ai_chat_sessions_table.py | 46 ++++++++++ backend/app/models/__init__.py | 2 + backend/app/models/ai_chat_session.py | 88 +++++++++++++++++++ 3 files changed, 136 insertions(+) create mode 100644 backend/alembic/versions/e2d81e82ea5e_add_ai_chat_sessions_table.py create mode 100644 backend/app/models/ai_chat_session.py diff --git a/backend/alembic/versions/e2d81e82ea5e_add_ai_chat_sessions_table.py b/backend/alembic/versions/e2d81e82ea5e_add_ai_chat_sessions_table.py new file mode 100644 index 00000000..42d2264c --- /dev/null +++ b/backend/alembic/versions/e2d81e82ea5e_add_ai_chat_sessions_table.py @@ -0,0 +1,46 @@ +"""add ai_chat_sessions table + +Revision ID: e2d81e82ea5e +Revises: 1490781700bc +Create Date: 2026-02-27 03:41:33.832260 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.dialects.postgresql import UUID, JSONB + + +# revision identifiers, used by Alembic. +revision: str = 'e2d81e82ea5e' +down_revision: Union[str, None] = '1490781700bc' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +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") diff --git a/backend/app/models/__init__.py b/backend/app/models/__init__.py index 3731740b..3748d9c8 100644 --- a/backend/app/models/__init__.py +++ b/backend/app/models/__init__.py @@ -28,6 +28,7 @@ from .maintenance_schedule import MaintenanceSchedule from .feedback import Feedback from .ai_conversation import AIConversation from .ai_usage import AIUsage +from .ai_chat_session import AIChatSession __all__ = [ "User", @@ -67,4 +68,5 @@ __all__ = [ "Feedback", "AIConversation", "AIUsage", + "AIChatSession", ] diff --git a/backend/app/models/ai_chat_session.py b/backend/app/models/ai_chat_session.py new file mode 100644 index 00000000..8fbde1c6 --- /dev/null +++ b/backend/app/models/ai_chat_session.py @@ -0,0 +1,88 @@ +"""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), + ) From 5c6745519069fda78ccde09f71a0dcb01205ce7a Mon Sep 17 00:00:00 2001 From: chihlasm Date: Fri, 27 Feb 2026 03:44:01 -0500 Subject: [PATCH 02/19] feat: add Pydantic schemas for AI chat builder Co-Authored-By: Claude Opus 4.6 --- backend/app/schemas/ai_chat.py | 80 ++++++++++++++++++++++++++++++++++ 1 file changed, 80 insertions(+) create mode 100644 backend/app/schemas/ai_chat.py diff --git a/backend/app/schemas/ai_chat.py b/backend/app/schemas/ai_chat.py new file mode 100644 index 00000000..35ae66b7 --- /dev/null +++ b/backend/app/schemas/ai_chat.py @@ -0,0 +1,80 @@ +"""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 From b7e48fae0e53f0bf3d1ce08a766b643bf2f93878 Mon Sep 17 00:00:00 2001 From: chihlasm Date: Fri, 27 Feb 2026 03:47:56 -0500 Subject: [PATCH 03/19] feat: add AI chat builder service with system prompt and conversation loop Co-Authored-By: Claude Opus 4.6 --- backend/app/core/ai_chat_service.py | 456 ++++++++++++++++++++++++++++ 1 file changed, 456 insertions(+) create mode 100644 backend/app/core/ai_chat_service.py diff --git a/backend/app/core/ai_chat_service.py b/backend/app/core/ai_chat_service.py new file mode 100644 index 00000000..4202ac31 --- /dev/null +++ b/backend/app/core/ai_chat_service.py @@ -0,0 +1,456 @@ +"""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 From babfd0c6d99af5603bd1467eb1d94a690316caad Mon Sep 17 00:00:00 2001 From: chihlasm Date: Fri, 27 Feb 2026 03:49:35 -0500 Subject: [PATCH 04/19] feat: add generate_text method to AIProvider for non-JSON responses The AI Chat Builder needs conversational text responses, not JSON-only. Gemini's generate_json forces response_mime_type='application/json' which is incompatible. The new generate_text method omits this constraint. Co-Authored-By: Claude Opus 4.6 --- backend/app/core/ai_chat_service.py | 6 +-- backend/app/core/ai_provider.py | 78 +++++++++++++++++++++++++++++ 2 files changed, 81 insertions(+), 3 deletions(-) diff --git a/backend/app/core/ai_chat_service.py b/backend/app/core/ai_chat_service.py index 4202ac31..b87a3698 100644 --- a/backend/app/core/ai_chat_service.py +++ b/backend/app/core/ai_chat_service.py @@ -242,7 +242,7 @@ async def start_chat_session( provider_name = settings.AI_PROVIDER messages = [{"role": "user", "content": primer}] - response_text, input_tokens, output_tokens = await provider.generate_json( + response_text, input_tokens, output_tokens = await provider.generate_text( system_prompt=system_prompt, messages=messages, max_tokens=1500, @@ -291,7 +291,7 @@ async def send_message( ] provider = get_ai_provider() - response_text, input_tokens, output_tokens = await provider.generate_json( + response_text, input_tokens, output_tokens = await provider.generate_text( system_prompt=system_prompt, messages=provider_messages, max_tokens=2000, @@ -371,7 +371,7 @@ Also provide metadata as a separate JSON object after the tree: provider = get_ai_provider() for attempt in range(2): # One try + one retry - response_text, input_tokens, output_tokens = await provider.generate_json( + response_text, input_tokens, output_tokens = await provider.generate_text( system_prompt=system_prompt, messages=provider_messages, max_tokens=8000, diff --git a/backend/app/core/ai_provider.py b/backend/app/core/ai_provider.py index cb3f7178..993012c6 100644 --- a/backend/app/core/ai_provider.py +++ b/backend/app/core/ai_provider.py @@ -35,6 +35,25 @@ class AIProvider(ABC): """ ... + @abstractmethod + async def generate_text( + self, + system_prompt: str, + messages: list[dict[str, str]], + max_tokens: int = 4096, + ) -> tuple[str, int, int]: + """Generate a text response from the AI model (no JSON constraint). + + Args: + system_prompt: System-level instruction for the model. + messages: List of message dicts with "role" and "content" keys. + max_tokens: Maximum output tokens. + + Returns: + Tuple of (response_text, input_tokens, output_tokens). + """ + ... + class GeminiProvider(AIProvider): """Google Gemini provider using the google-genai SDK.""" @@ -95,6 +114,56 @@ class GeminiProvider(AIProvider): return text, input_tokens, output_tokens + async def generate_text( + self, + system_prompt: str, + messages: list[dict[str, str]], + max_tokens: int = 4096, + ) -> tuple[str, int, int]: + from google import genai + from google.genai import types as genai_types + + client = genai.Client(api_key=self._api_key) + + contents: list[genai_types.Content] = [] + for msg in messages: + role = "model" if msg["role"] == "assistant" else "user" + contents.append( + genai_types.Content( + role=role, + parts=[genai_types.Part(text=msg["content"])], + ) + ) + + config = genai_types.GenerateContentConfig( + system_instruction=system_prompt, + max_output_tokens=max_tokens, + # No response_mime_type — allow free-form text + ) + + response = await client.aio.models.generate_content( + model=self._model, + contents=contents, + config=config, + ) + + if response.candidates: + finish_reason = getattr(response.candidates[0], "finish_reason", None) + logger.info("Gemini finish_reason=%s model=%s", finish_reason, self._model) + if str(finish_reason) == "MAX_TOKENS": + logger.warning( + "Gemini output truncated (MAX_TOKENS). max_output_tokens=%d", + max_tokens, + ) + + text = response.text or "" + input_tokens = getattr(response.usage_metadata, "prompt_token_count", 0) or 0 + output_tokens = ( + getattr(response.usage_metadata, "candidates_token_count", 0) or 0 + ) + + return text, input_tokens, output_tokens + class AnthropicProvider(AIProvider): """Anthropic Claude provider using the anthropic SDK.""" @@ -130,6 +199,15 @@ class AnthropicProvider(AIProvider): return text, input_tokens, output_tokens + async def generate_text( + self, + system_prompt: str, + messages: list[dict[str, str]], + max_tokens: int = 4096, + ) -> tuple[str, int, int]: + # Anthropic doesn't differentiate between JSON and text mode + return await self.generate_json(system_prompt, messages, max_tokens) + def get_ai_provider() -> AIProvider: """Factory that returns the configured AI provider. From ef96b1a12fbf01328f948210961030bce427544a Mon Sep 17 00:00:00 2001 From: chihlasm Date: Fri, 27 Feb 2026 03:54:08 -0500 Subject: [PATCH 05/19] feat: add AI chat builder endpoints and update quota service 6 endpoints: create session, send message, get session, generate tree, import to editor, abandon. Quota service daily limit updated to include chat builder generation types. Co-Authored-By: Claude Opus 4.6 --- backend/app/api/endpoints/ai_chat.py | 428 +++++++++++++++++++++++++++ backend/app/api/router.py | 2 + backend/app/core/ai_quota_service.py | 2 +- 3 files changed, 431 insertions(+), 1 deletion(-) create mode 100644 backend/app/api/endpoints/ai_chat.py diff --git a/backend/app/api/endpoints/ai_chat.py b/backend/app/api/endpoints/ai_chat.py new file mode 100644 index 00000000..326290d9 --- /dev/null +++ b/backend/app/api/endpoints/ai_chat.py @@ -0,0 +1,428 @@ +"""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() + + 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.", + ) + + 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", + ) + + 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.", + ) + + 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) + + 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 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() + + 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.", + ) + + 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", + ) + + 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() diff --git a/backend/app/api/router.py b/backend/app/api/router.py index 27963a1f..41fdb0b2 100644 --- a/backend/app/api/router.py +++ b/backend/app/api/router.py @@ -7,6 +7,7 @@ from app.api.endpoints import maintenance_schedules from app.api.endpoints import feedback from app.api.endpoints import ai_builder from app.api.endpoints import ai_fix +from app.api.endpoints import ai_chat api_router = APIRouter() @@ -38,3 +39,4 @@ api_router.include_router(maintenance_schedules.router) api_router.include_router(feedback.router) api_router.include_router(ai_builder.router) api_router.include_router(ai_fix.router) +api_router.include_router(ai_chat.router) diff --git a/backend/app/core/ai_quota_service.py b/backend/app/core/ai_quota_service.py index 67eed9e5..49264caa 100644 --- a/backend/app/core/ai_quota_service.py +++ b/backend/app/core/ai_quota_service.py @@ -115,7 +115,7 @@ async def check_ai_quota( select(func.count(AIUsage.id)).where( AIUsage.user_id == user_id, AIUsage.succeeded == True, # noqa: E712 - AIUsage.generation_type.in_(["scaffold", "branch_detail"]), + AIUsage.generation_type.in_(["scaffold", "branch_detail", "chat_message", "chat_generate"]), AIUsage.created_at >= day_start, ) ) or 0 From 0da67586da90f919c9ad226087a89fb91d89670b Mon Sep 17 00:00:00 2001 From: chihlasm Date: Fri, 27 Feb 2026 04:06:47 -0500 Subject: [PATCH 06/19] feat: add backend tests for AI chat builder + fix conversation_id FK issue Tests cover session create, send message with tree update, get session, abandon, 404 on missing session, and 503 when AI disabled. Fixed: ai_usage.conversation_id has FK to ai_conversations, not ai_chat_sessions. Chat builder now passes conversation_id=None and tracks session reference in extra_data.chat_session_id. Co-Authored-By: Claude Opus 4.6 --- backend/app/api/endpoints/ai_chat.py | 24 ++-- backend/tests/test_ai_chat.py | 187 +++++++++++++++++++++++++++ 2 files changed, 199 insertions(+), 12 deletions(-) create mode 100644 backend/tests/test_ai_chat.py diff --git a/backend/app/api/endpoints/ai_chat.py b/backend/app/api/endpoints/ai_chat.py index 326290d9..a1b8c648 100644 --- a/backend/app/api/endpoints/ai_chat.py +++ b/backend/app/api/endpoints/ai_chat.py @@ -106,7 +106,7 @@ async def create_session( await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, - conversation_id=session.id, + conversation_id=None, generation_type="chat_message", tier=plan, input_tokens=session.total_input_tokens, @@ -118,7 +118,7 @@ async def create_session( succeeded=True, counts_toward_quota=False, error_code=None, - extra_data={"phase": "scoping"}, + extra_data={"phase": "scoping", "chat_session_id": str(session.id)}, db=db, ) @@ -175,7 +175,7 @@ async def post_message( await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, - conversation_id=session.id, + conversation_id=None, generation_type="chat_message", tier=plan, input_tokens=0, @@ -184,7 +184,7 @@ async def post_message( succeeded=False, counts_toward_quota=False, error_code=type(e).__name__, - extra_data=None, + extra_data={"chat_session_id": str(session.id)}, db=db, ) await db.commit() @@ -198,7 +198,7 @@ async def post_message( await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, - conversation_id=session.id, + conversation_id=None, generation_type="chat_message", tier=plan, input_tokens=input_delta, @@ -210,7 +210,7 @@ async def post_message( succeeded=True, counts_toward_quota=False, error_code=None, - extra_data={"phase": session.current_phase}, + extra_data={"phase": session.current_phase, "chat_session_id": str(session.id)}, db=db, ) @@ -288,7 +288,7 @@ async def generate_tree( await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, - conversation_id=session.id, + conversation_id=None, generation_type="chat_generate", tier=plan, input_tokens=session.total_input_tokens - prev_input, @@ -297,7 +297,7 @@ async def generate_tree( succeeded=False, counts_toward_quota=False, error_code="invalid_output", - extra_data={"error": str(e)}, + extra_data={"error": str(e), "chat_session_id": str(session.id)}, db=db, ) await db.commit() @@ -312,7 +312,7 @@ async def generate_tree( await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, - conversation_id=session.id, + conversation_id=None, generation_type="chat_generate", tier=plan, input_tokens=input_delta, @@ -321,7 +321,7 @@ async def generate_tree( succeeded=False, counts_toward_quota=False, error_code=type(e).__name__, - extra_data={"error": str(e)}, + extra_data={"error": str(e), "chat_session_id": str(session.id)}, db=db, ) await db.commit() @@ -342,7 +342,7 @@ async def generate_tree( await record_ai_usage( user_id=current_user.id, account_id=current_user.account_id, - conversation_id=session.id, + conversation_id=None, generation_type="chat_generate", tier=plan, input_tokens=input_delta, @@ -354,7 +354,7 @@ async def generate_tree( succeeded=True, counts_toward_quota=True, error_code=None, - extra_data=None, + extra_data={"chat_session_id": str(session.id)}, db=db, ) diff --git a/backend/tests/test_ai_chat.py b/backend/tests/test_ai_chat.py new file mode 100644 index 00000000..0b52f650 --- /dev/null +++ b/backend/tests/test_ai_chat.py @@ -0,0 +1,187 @@ +"""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_text = 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 — must pass validate_generated_tree (min 5 nodes) + import json + tree_obj = { + "id": "root", "type": "decision", + "question": "What DNS symptom is the user experiencing?", + "options": [ + {"id": "opt-1", "label": "Cannot resolve any domains", "next_node_id": "dns-check"}, + {"id": "opt-2", "label": "Intermittent failures", "next_node_id": "dns-cache-fix"}, + ], + "children": [ + { + "id": "dns-check", "type": "decision", + "question": "Is the DNS Client service running?", + "options": [ + {"id": "dc-1", "label": "Yes", "next_node_id": "dns-fwd-fix"}, + {"id": "dc-2", "label": "No", "next_node_id": "dns-svc-fix"}, + ], + "children": [ + {"id": "dns-fwd-fix", "type": "solution", "title": "Check DNS Forwarders", + "description": "DNS forwarders may be misconfigured", + "resolution_steps": ["Check forwarder config"]}, + {"id": "dns-svc-fix", "type": "solution", "title": "Restart DNS Service", + "description": "DNS Client service is stopped", + "resolution_steps": ["Start-Service Dnscache"]}, + ], + }, + {"id": "dns-cache-fix", "type": "solution", "title": "Stale DNS Cache", + "description": "DNS cache has stale entries", + "resolution_steps": ["ipconfig /flushdns"]}, + ], + } + tree_json = json.dumps(tree_obj) + mock_ai_provider.generate_text = 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" + f"[TREE_UPDATE]\n{tree_json}\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_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 From 596153085a4d6b1b610fb9ae07065037d4dca5dc Mon Sep 17 00:00:00 2001 From: chihlasm Date: Fri, 27 Feb 2026 07:20:04 -0500 Subject: [PATCH 07/19] =?UTF-8?q?feat:=20add=20AI=20chat=20builder=20front?= =?UTF-8?q?end=20=E2=80=94=20types,=20API=20client,=20store,=20components,?= =?UTF-8?q?=20page,=20routing?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - TypeScript types for chat session, messages, and responses - API client module with all 6 endpoints - Zustand store with session management, message sending, tree generation, import, resume - 7 chat components: ChatMessage, ChatInput, ChatPanel, PhaseIndicator, ChatToolbar, EmptyPreview, StaticTreePreview - AIChatBuilderPage with split-panel layout (60% chat / 40% preview) - Route at /ai/chat with lazy loading - "Build with AI" button on TreeLibraryPage - Session resume via URL search params Co-Authored-By: Claude Opus 4.6 --- frontend/src/api/aiChat.ts | 44 ++++ frontend/src/api/index.ts | 1 + frontend/src/components/ai-chat/ChatInput.tsx | 72 +++++++ .../src/components/ai-chat/ChatMessage.tsx | 41 ++++ frontend/src/components/ai-chat/ChatPanel.tsx | 47 +++++ .../src/components/ai-chat/ChatToolbar.tsx | 76 +++++++ .../src/components/ai-chat/EmptyPreview.tsx | 13 ++ .../src/components/ai-chat/PhaseIndicator.tsx | 50 +++++ .../components/ai-chat/StaticTreePreview.tsx | 80 ++++++++ frontend/src/pages/AIChatBuilderPage.tsx | 151 ++++++++++++++ frontend/src/pages/TreeLibraryPage.tsx | 23 ++- frontend/src/router.tsx | 9 + frontend/src/store/aiChatStore.ts | 190 ++++++++++++++++++ frontend/src/types/ai-chat.ts | 43 ++++ frontend/src/types/index.ts | 10 + 15 files changed, 844 insertions(+), 6 deletions(-) create mode 100644 frontend/src/api/aiChat.ts create mode 100644 frontend/src/components/ai-chat/ChatInput.tsx create mode 100644 frontend/src/components/ai-chat/ChatMessage.tsx create mode 100644 frontend/src/components/ai-chat/ChatPanel.tsx create mode 100644 frontend/src/components/ai-chat/ChatToolbar.tsx create mode 100644 frontend/src/components/ai-chat/EmptyPreview.tsx create mode 100644 frontend/src/components/ai-chat/PhaseIndicator.tsx create mode 100644 frontend/src/components/ai-chat/StaticTreePreview.tsx create mode 100644 frontend/src/pages/AIChatBuilderPage.tsx create mode 100644 frontend/src/store/aiChatStore.ts create mode 100644 frontend/src/types/ai-chat.ts diff --git a/frontend/src/api/aiChat.ts b/frontend/src/api/aiChat.ts new file mode 100644 index 00000000..1722b272 --- /dev/null +++ b/frontend/src/api/aiChat.ts @@ -0,0 +1,44 @@ +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 diff --git a/frontend/src/api/index.ts b/frontend/src/api/index.ts index fba65d24..016efcbe 100644 --- a/frontend/src/api/index.ts +++ b/frontend/src/api/index.ts @@ -17,3 +17,4 @@ export { targetListsApi } from './targetLists' export { maintenanceSchedulesApi, batchLaunchApi } from './maintenanceSchedules' export { default as feedbackApi } from './feedback' export { default as aiBuilderApi } from './aiBuilder' +export { default as aiChatApi } from './aiChat' diff --git a/frontend/src/components/ai-chat/ChatInput.tsx b/frontend/src/components/ai-chat/ChatInput.tsx new file mode 100644 index 00000000..c10f4970 --- /dev/null +++ b/frontend/src/components/ai-chat/ChatInput.tsx @@ -0,0 +1,72 @@ +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 ( +
+