resolutionflow/backend/app/services/copilot_service.py

"""Copilot service — in-session AI assistant with RAG context.

Builds system prompts with current flow context and RAG results,
manages conversation state, and returns AI responses with flow suggestions.
"""
import logging
from datetime import datetime, timezone, timedelta
from typing import Optional, Any
from uuid import UUID

from sqlalchemy import select, or_
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.core.ai_provider import get_ai_provider
from app.models.tree import Tree
from app.models.copilot_conversation import CopilotConversation
from app.services.rag_service import search as rag_search, build_rag_context, extract_suggested_flows

logger = logging.getLogger(__name__)

COPILOT_SYSTEM_PROMPT = """You are a Senior Systems and Network Engineer with 15+ years of experience working in Managed Service Provider (MSP) environments. You specialize in:
- Windows Server, Active Directory, Group Policy, and Hybrid Identity (Entra ID)
- Networking (TCP/IP, DNS, DHCP, VPN, firewall troubleshooting, Cisco/Fortinet)
- Virtualization (VMware, Hyper-V) and cloud platforms (Azure, AWS, M365)
- Endpoint management, RMM tools, and PSA platforms (ConnectWise, Datto, Kaseya)
- PowerShell scripting and automation

You are acting as an in-session copilot while the user navigates a troubleshooting or procedural flow. You can see the flow context and their current position.

When answering:
- Be direct and actionable — MSP engineers need fast, practical answers
- Include specific commands, paths, and config values when relevant
- Mention potential risks or gotchas before suggesting changes
- If a relevant troubleshooting flow exists in the team's library, reference it
- Keep responses concise but thorough — prefer bullet points and code blocks
"""


def _build_flow_context(tree: Tree, current_node_id: Optional[str]) -> str:
    """Build flow context string for the system prompt."""
    parts = [
        f"\n--- CURRENT FLOW CONTEXT ---",
        f"Flow: {tree.name}",
        f"Type: {tree.tree_type}",
    ]
    if tree.description:
        parts.append(f"Description: {tree.description}")

    if current_node_id and tree.tree_structure:
        node = _find_node(tree.tree_structure, current_node_id)
        if node:
            parts.append(f"Current node type: {node.get('type', 'unknown')}")
            node_label = (
                node.get('title')
                or node.get('question')
                or node.get('description')
                or node.get('content')
                or node.get('label')
                or 'Unknown'
            )
            parts.append(f"Current node: {node_label}")
            # Add options if it's a question/decision node
            children = node.get("children", [])
            if children and isinstance(children, list):
                option_labels = [
                    c.get("label", c.get("content", ""))
                    for c in children if isinstance(c, dict)
                ]
                if option_labels:
                    parts.append(f"Available options: {', '.join(option_labels)}")

    return "\n".join(parts)


def _find_node(structure: dict, node_id: str) -> Optional[dict]:
    """Recursively find a node by ID in tree structure."""
    if structure.get("id") == node_id:
        return structure
    for child in structure.get("children", []):
        if isinstance(child, dict):
            found = _find_node(child, node_id)
            if found:
                return found
    # Check steps array for procedural flows
    for step in structure.get("steps", []):
        if isinstance(step, dict):
            found = _find_node(step, node_id)
            if found:
                return found
    return None


async def start_conversation(
    user_id: UUID,
    account_id: UUID,
    tree_id: UUID,
    session_id: Optional[UUID],
    current_node_id: Optional[str],
    db: AsyncSession,
) -> tuple[CopilotConversation, str]:
    """Start a new copilot conversation.

    Returns (conversation, greeting_message).
    """
    # Load tree — must be accessible to this account.
    # Allows own account's trees, default trees, and public trees.
    # Raises ValueError (caught by endpoint as 404) if not found or not accessible.
    result = await db.execute(
        select(Tree).options(selectinload(Tree.tags)).where(
            Tree.id == tree_id,
            or_(
                Tree.account_id == account_id,
                Tree.author_id == user_id,
                Tree.is_default == True,
                Tree.is_public == True,
            ),
        )
    )
    tree = result.scalar_one_or_none()
    if not tree:
        raise ValueError(f"Tree {tree_id} not found or not accessible")

    conversation = CopilotConversation(
        user_id=user_id,
        account_id=account_id,
        tree_id=tree_id,
        session_id=session_id,
        current_node_id=current_node_id,
        messages=[],
        expires_at=datetime.now(timezone.utc) + timedelta(hours=24),
    )
    db.add(conversation)
    await db.flush()

    greeting = f"I'm your copilot for this **{tree.tree_type}** flow: **{tree.name}**. Ask me anything about the current step, alternative approaches, or related troubleshooting tips."

    conversation.messages = [{"role": "assistant", "content": greeting}]
    conversation.message_count = 1

    return conversation, greeting


async def send_message(
    conversation_id: UUID,
    user_id: UUID,
    message: str,
    current_node_id: Optional[str],
    db: AsyncSession,
) -> tuple[str, list[dict[str, Any]], CopilotConversation]:
    """Send a user message and get AI response.

    Returns (ai_content, suggested_flows, conversation).
    """
    result = await db.execute(
        select(CopilotConversation).where(
            CopilotConversation.id == conversation_id,
            CopilotConversation.user_id == user_id,
        )
    )
    conversation = result.scalar_one_or_none()
    if not conversation:
        raise ValueError("Conversation not found")

    if conversation.expires_at < datetime.now(timezone.utc):
        raise ValueError("Conversation has expired")

    # Load tree for context
    tree_result = await db.execute(
        select(Tree).options(selectinload(Tree.tags)).where(Tree.id == conversation.tree_id)
    )
    tree = tree_result.scalar_one_or_none()
    if not tree:
        raise ValueError("Associated flow not found")

    # Update current node
    if current_node_id:
        conversation.current_node_id = current_node_id

    # RAG search
    rag_results = await rag_search(
        query=message,
        account_id=conversation.account_id,
        db=db,
        limit=8,
    )

    # Build system prompt
    system_prompt = COPILOT_SYSTEM_PROMPT
    system_prompt += _build_flow_context(tree, conversation.current_node_id)
    system_prompt += build_rag_context(rag_results)

    # Inject PSA ticket context if session has a linked ticket
    if conversation.session_id:
        try:
            from app.models.session import Session as SessionModel
            session_result = await db.execute(
                select(SessionModel).where(SessionModel.id == conversation.session_id)
            )
            session = session_result.scalar_one_or_none()
            if session and session.psa_ticket_id:
                try:
                    from app.services.psa.registry import get_provider_for_account
                    from app.services.psa.ticket_context import format_ticket_context_for_prompt

                    provider = await get_provider_for_account(conversation.account_id, db)
                    connection_id = str(session.psa_connection_id) if session.psa_connection_id else None
                    ticket_ctx = await provider.get_ticket_context(
                        ticket_id=int(session.psa_ticket_id),
                        connection_id=connection_id,
                    )
                    system_prompt += "\n\n" + format_ticket_context_for_prompt(ticket_ctx)
                except Exception as psa_err:
                    logger.warning(
                        "Failed to fetch PSA ticket context for copilot (session=%s, ticket=%s): %s",
                        conversation.session_id,
                        session.psa_ticket_id,
                        psa_err,
                    )
        except Exception as session_err:
            logger.warning(
                "Failed to look up session for copilot PSA context (session_id=%s): %s",
                conversation.session_id,
                session_err,
            )

    # Build messages for AI
    ai_messages = []
    for msg in conversation.messages:
        if msg["role"] in ("user", "assistant"):
            ai_messages.append({"role": msg["role"], "content": msg["content"]})
    ai_messages.append({"role": "user", "content": message})

    # Call AI
    provider = get_ai_provider()
    ai_content, input_tokens, output_tokens = await provider.generate_text(
        system_prompt=system_prompt,
        messages=ai_messages,
        max_tokens=2048,
    )

    # Update conversation
    msgs = list(conversation.messages)
    msgs.append({"role": "user", "content": message})
    msgs.append({"role": "assistant", "content": ai_content})
    conversation.messages = msgs
    conversation.message_count += 2
    conversation.total_input_tokens += input_tokens
    conversation.total_output_tokens += output_tokens

    # Extract suggested flows
    suggested_flows = extract_suggested_flows(rag_results, exclude_tree_id=tree.id)

    return ai_content, suggested_flows, conversation