api/app/api/routes/agent_setup.py

"""Chatbot Journey endpoints — guided conversation to build an agent prompt_template.

Endpoints:
  POST /agents/journey/start    — start a new journey session
  POST /agents/journey/message  — continue the conversation

Sessions are stored in-memory with a 30-minute TTL.  Stale entries are
cleaned up lazily on access.  Upgrade to Redis for multi-instance deployments.

Journey flow:
  1. Client sends ``{ agent_type, agent_id? }`` to ``/start``.
  2. Server creates a session, calls the LLM with a contextual system prompt,
     and returns the first question.
  3. Client sends follow-up messages to ``/message``.
  4. After 3-5 turns the LLM wraps up by emitting a ``prompt_template`` block
     delimited by ``PROMPT_TEMPLATE_START`` / ``PROMPT_TEMPLATE_END``.
  5. Server parses the block, sets ``done=True``, and returns the template.

The ``prompt_template`` from the final response is meant to be stored in
``LocalAgentConfig.prompt_template`` or ``CloudAgentConfig.prompt_template``
by the Electron client (via the agent CRUD endpoints).
"""

from __future__ import annotations

import logging
import time
import uuid
from dataclasses import dataclass, field
from typing import Any

from fastapi import APIRouter, Depends, HTTPException, status
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.deps import get_current_user
from app.core.llm import get_llm
from app.db import get_session
from app.models import CloudAgentConfig, LocalAgentConfig
from app.schemas import (
    JourneyMessageRequest,
    JourneyResponse,
    JourneyStartRequest,
    UserProfile,
)

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/agents/journey", tags=["agents"])

# ── Session TTL ───────────────────────────────────────────────────────────

_SESSION_TTL_SECONDS: int = 1800  # 30 minutes

# Sentinel strings used to delimit the LLM-produced prompt_template.
_TEMPLATE_START = "PROMPT_TEMPLATE_START"
_TEMPLATE_END = "PROMPT_TEMPLATE_END"

# Maximum number of conversation turns before the LLM is nudged to wrap up.
_MAX_TURNS: int = 5

# ── In-memory session store ───────────────────────────────────────────────


@dataclass
class _JourneySession:
    session_id: str
    user_id: str
    agent_type: str  # "local" | "cloud"
    history: list[dict[str, Any]] = field(default_factory=list)
    created_at: float = field(default_factory=time.monotonic)

    def is_expired(self) -> bool:
        return (time.monotonic() - self.created_at) > _SESSION_TTL_SECONDS


# session_id → session
_sessions: dict[str, _JourneySession] = {}


def _get_session(session_id: str, user_id: str) -> _JourneySession:
    """Retrieve session; raise 404 on missing, expired, or wrong owner."""
    s = _sessions.get(session_id)
    if s is None or s.is_expired():
        _sessions.pop(session_id, None)
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Journey session not found or expired")
    if s.user_id != user_id:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Journey session not found or expired")
    return s


# ── System prompt builder ─────────────────────────────────────────────────

_LOCAL_PREAMBLE = """\
What kind of files are in the directories you want to monitor? \
(for example: emails saved as .eml, documents in .pdf or .txt, markdown notes, etc.)"""

_CLOUD_PREAMBLE = """\
What kind of emails or messages should I look for? \
(for example: client communications, invoices, meeting notes, project updates, etc.)"""

_SYSTEM_PROMPT_TEMPLATE = """\
You are a friendly assistant helping a freelancer configure a data-extraction agent.
Your job is to understand exactly what data the user wants to extract from their {source_description} \
and produce a detailed prompt_template that a separate AI will use as its instruction set.

Ask concise, focused questions one at a time.  Cover these topics (not necessarily in this order):
  1. The type and format of the source content.
  2. Which data types to extract: tasks, notes, timelines, and/or projects.
  3. How fields should be mapped (e.g. email subject → task title).
  4. Priority or status rules (e.g. "urgent" keyword → high priority).
  5. Any special handling, date extraction, or exclusions.

After 3-5 questions (when you have enough information), output the final prompt_template between \
these exact markers on their own lines:

{template_start}
<the complete extraction prompt here>
{template_end}

The prompt_template must be a self-contained instruction for an AI that receives a document/email/message \
and must return a JSON array of records in this shape:
  [{{ "table": "<tasks|notes|timelines|projects>", "data": {{ <field: value> }} }}, ...]

Rules for the generated template:
  - Be explicit about field names (camelCase: title, status, priority, dueDate, projectId, content, etc.).
  - Include concrete examples of mappings.
  - Mention that Electron adds id/createdAt/updatedAt automatically.
  - Set isAiSuggested: true and isApproved: false on every record.
{existing_section}\
Do not ask more than {max_turns} questions total. Start with your first question now.\
"""


def _build_system_prompt(agent_type: str, existing_template: str | None) -> str:
    source_description = (
        "files in local directories" if agent_type == "local" else "emails and messages from cloud providers"
    )
    existing_section = (
        f"\nThe user already has the following prompt_template — refine it based on their answers:\n"
        f"---\n{existing_template}\n---\n"
        if existing_template
        else ""
    )
    return _SYSTEM_PROMPT_TEMPLATE.format(
        source_description=source_description,
        template_start=_TEMPLATE_START,
        template_end=_TEMPLATE_END,
        existing_section=existing_section,
        max_turns=_MAX_TURNS,
    )


def _first_question(agent_type: str) -> str:
    return _LOCAL_PREAMBLE if agent_type == "local" else _CLOUD_PREAMBLE


# ── Template extraction ───────────────────────────────────────────────────


def _extract_template(text: str) -> str | None:
    """Return the text between PROMPT_TEMPLATE_START and PROMPT_TEMPLATE_END, or None."""
    if _TEMPLATE_START not in text or _TEMPLATE_END not in text:
        return None
    start_idx = text.index(_TEMPLATE_START) + len(_TEMPLATE_START)
    end_idx = text.index(_TEMPLATE_END)
    return text[start_idx:end_idx].strip() or None


# ── LLM call ─────────────────────────────────────────────────────────────


async def _call_llm(system_prompt: str, history: list[dict[str, Any]]) -> str:
    """Build LangChain messages from history and invoke the LLM."""
    messages: list[Any] = [SystemMessage(content=system_prompt)]
    for turn in history:
        if turn["role"] == "user":
            messages.append(HumanMessage(content=turn["content"]))
        else:
            messages.append(AIMessage(content=turn["content"]))

    llm = get_llm(model=None, temperature=0.4)
    response = await llm.ainvoke(messages)
    return response.content  # type: ignore[return-value]


# ── Existing-config loader ────────────────────────────────────────────────


async def _load_existing_template(
    agent_id: str,
    user_id: str,
    db: AsyncSession,
) -> str | None:
    """Return the prompt_template of an existing agent config, or None."""
    # Try local first, then cloud.
    local_result = await db.execute(
        select(LocalAgentConfig).where(
            LocalAgentConfig.id == agent_id,
            LocalAgentConfig.user_id == user_id,
        )
    )
    local = local_result.scalar_one_or_none()
    if local is not None:
        return local.prompt_template

    cloud_result = await db.execute(
        select(CloudAgentConfig).where(
            CloudAgentConfig.id == agent_id,
            CloudAgentConfig.user_id == user_id,
        )
    )
    cloud = cloud_result.scalar_one_or_none()
    return cloud.prompt_template if cloud is not None else None


# ── Routes ────────────────────────────────────────────────────────────────


@router.post("/start", response_model=JourneyResponse, status_code=status.HTTP_200_OK)
async def start_journey(
    body: JourneyStartRequest,
    current_user: UserProfile = Depends(get_current_user),
    db: AsyncSession = Depends(get_session),
) -> JourneyResponse:
    """Start a new Chatbot Journey session.

    If ``agent_id`` is provided the session is pre-seeded with the existing
    agent's ``prompt_template`` so the user can refine it.
    """
    # Load existing template (may be None).
    existing_template: str | None = None
    if body.agent_id:
        existing_template = await _load_existing_template(body.agent_id, current_user.id, db)
        # If agent_id was given but not found, proceed without seeding (don't 404 —
        # the user may be starting a fresh journey for a not-yet-persisted config).

    system_prompt = _build_system_prompt(body.agent_type, existing_template)
    first_question = _first_question(body.agent_type)

    session_id = str(uuid.uuid4())
    session = _JourneySession(
        session_id=session_id,
        user_id=current_user.id,
        agent_type=body.agent_type,
        # Seed history with the AI's first question so it stays consistent.
        history=[{"role": "assistant", "content": first_question}],
    )
    # Store the system prompt inside the session for reuse in /message.
    session.__dict__["_system_prompt"] = system_prompt  # type: ignore[index]
    _sessions[session_id] = session

    logger.info("Journey session %s started for user %s (agent_type=%s)", session_id, current_user.id, body.agent_type)
    return JourneyResponse(session_id=session_id, message=first_question, done=False)


@router.post("/message", response_model=JourneyResponse, status_code=status.HTTP_200_OK)
async def send_journey_message(
    body: JourneyMessageRequest,
    current_user: UserProfile = Depends(get_current_user),
    db: AsyncSession = Depends(get_session),
) -> JourneyResponse:
    """Send a message in an existing Chatbot Journey session.

    The server appends the user's message to the conversation history,
    calls the LLM, and appends the AI reply.  When the LLM wraps up with a
    ``prompt_template`` block the response includes ``done=True`` and the
    extracted template.
    """
    session = _get_session(body.session_id, current_user.id)
    system_prompt: str = session.__dict__.get("_system_prompt", _build_system_prompt(session.agent_type, None))  # type: ignore[assignment]

    # Append user turn to history.
    session.history.append({"role": "user", "content": body.message})

    # Call the LLM with the full conversation so far.
    ai_reply = await _call_llm(system_prompt, session.history)

    # Append AI turn.
    session.history.append({"role": "assistant", "content": ai_reply})

    # Check if the LLM produced the final template.
    prompt_template = _extract_template(ai_reply)
    done = prompt_template is not None

    # Strip the sentinel markers from the message shown to the user.
    display_message = ai_reply
    if done:
        display_message = (
            ai_reply[: ai_reply.index(_TEMPLATE_START)].strip()
            or "Here is your agent configuration. You can save it or continue refining."
        )

    if done:
        logger.info("Journey session %s completed for user %s", body.session_id, current_user.id)
        # Clean up the session immediately on completion.
        _sessions.pop(body.session_id, None)
    else:
        # Nudge the LLM to wrap up after max turns.
        turns = sum(1 for t in session.history if t["role"] == "user")
        if turns >= _MAX_TURNS:
            # Add a system-level nudge as a hidden user message.
            session.history.append({
                "role": "user",
                "content": (
                    "[System: You have enough information. Please generate the final "
                    f"prompt_template now, wrapped in {_TEMPLATE_START} / {_TEMPLATE_END} markers.]"
                ),
            })

    return JourneyResponse(
        session_id=body.session_id,
        message=display_message,
        done=done,
        prompt_template=prompt_template,
    )