ss-tools/backend/src/api/routes/assistant/_schemas.py

# #region AssistantSchemas [C:2] [TYPE Module] [SEMANTICS assistant, pydantic, schema, store, permission]
# @defgroup AssistantApi Module group.
# @BRIEF Pydantic models, in-memory stores, and permission mappings for the assistant API.
# @LAYER API
# @RELATION CALLED_BY -> [AssistantHistory]
# @RELATION CALLED_BY -> [AssistantHistory]
# @INVARIANT In-memory stores are module-level singletons shared across the assistant package.
# @RATIONALE In-memory stores documented with NOTE about restart loss. ASSISTANT_ARCHIVE_AFTER_DAYS and ASSISTANT_MESSAGE_TTL_DAYS kept as module-level constants with TODO for config migration — Pydantic schemas module should not depend on ConfigManager for architectural purity.

from __future__ import annotations

from datetime import datetime
from typing import Any

from pydantic import BaseModel, Field


# #region AssistantMessageRequest [C:1] [TYPE Class]
# @BRIEF Input payload for assistant message endpoint.
# @DATA_CONTRACT Input[conversation_id?:str, message:str(1..4000)] -> Output[AssistantMessageRequest]
# @RELATION CALLED_BY -> [send_message]
# @SIDE_EFFECT None (schema declaration only).
# @PRE message length is within accepted bounds.
# @POST Request object provides message text and optional conversation binding.
# @INVARIANT message is always non-empty and no longer than 4000 characters.
class AssistantMessageRequest(BaseModel):
    conversation_id: str | None = None
    message: str = Field(..., min_length=1, max_length=4000)
    dataset_review_session_id: str | None = None


# #endregion AssistantMessageRequest


# #region AssistantAction [C:1] [TYPE Class]
# @BRIEF UI action descriptor returned with assistant responses.
# @DATA_CONTRACT Input[type:str, label:str, target?:str] -> Output[AssistantAction]
# @RELATION CALLED_BY -> [AssistantMessageResponse]
# @SIDE_EFFECT None (schema declaration only).
# @PRE type and label are provided by orchestration logic.
# @POST Action can be rendered as button on frontend.
# @INVARIANT type and label are required for every UI action.
class AssistantAction(BaseModel):
    type: str
    label: str
    target: str | None = None


# #endregion AssistantAction


# #region AssistantMessageResponse [C:1] [TYPE Class]
# @BRIEF Output payload contract for assistant interaction endpoints.
# @DATA_CONTRACT Input[conversation_id,response_id,state,text,intent?,confirmation_id?,task_id?,actions[],created_at] -> Output[AssistantMessageResponse]
# @RELATION CALLED_BY -> [send_message]
# @RELATION CALLED_BY -> [confirm_operation]
# @RELATION CALLED_BY -> [cancel_operation]
# @SIDE_EFFECT None (schema declaration only).
# @PRE Response includes deterministic state and text.
# @POST Payload may include task_id/confirmation_id/actions for UI follow-up.
# @INVARIANT created_at and state are always present in endpoint responses.
class AssistantMessageResponse(BaseModel):
    conversation_id: str
    response_id: str
    state: str
    text: str
    intent: dict[str, Any] | None = None
    confirmation_id: str | None = None
    task_id: str | None = None
    actions: list[AssistantAction] = Field(default_factory=list)
    created_at: datetime


# #endregion AssistantMessageResponse


# #region ConfirmationRecord [C:1] [TYPE Class]
# @BRIEF In-memory confirmation token model for risky operation dispatch.
# @DATA_CONTRACT Input[id,user_id,conversation_id,intent,dispatch,expires_at,state?,created_at] -> Output[ConfirmationRecord]
# @RELATION CALLED_BY -> [send_message]
# @RELATION CALLED_BY -> [confirm_operation]
# @RELATION CALLED_BY -> [cancel_operation]
# @SIDE_EFFECT None (schema declaration only).
# @PRE intent/dispatch/user_id are populated at confirmation request time.
# @POST Record tracks lifecycle state and expiry timestamp.
# @INVARIANT state defaults to "pending" and expires_at bounds confirmation validity.
class ConfirmationRecord(BaseModel):
    id: str
    user_id: str
    conversation_id: str
    intent: dict[str, Any]
    dispatch: dict[str, Any]
    expires_at: datetime
    state: str = "pending"
    created_at: datetime


# #endregion ConfirmationRecord


# --- In-memory stores ---

# #region CONVERSATIONS [C:1] [TYPE Constant]
# @BRIEF In-memory conversation history store: {(user_id, conversation_id) -> [message_records]}
# #endregion CONVERSATIONS
CONVERSATIONS: dict[tuple[str, str], list[dict[str, Any]]] = {}

# #region ASSISTANT_AUDIT [C:1] [TYPE Constant]
# @BRIEF In-memory assistant audit trail: {user_id -> [audit_records]}. Volatile — lost on restart.
# #endregion ASSISTANT_AUDIT
USER_ACTIVE_CONVERSATION: dict[str, str] = {}
CONFIRMATIONS: dict[str, ConfirmationRecord] = {}
ASSISTANT_AUDIT: dict[str, list[dict[str, Any]]] = {}
ASSISTANT_ARCHIVE_AFTER_DAYS = 14
ASSISTANT_MESSAGE_TTL_DAYS = 90

# Operations that are read-only and do not require confirmation.
_SAFE_OPS = {
    "show_capabilities",
    "get_task_status",
    "get_health_summary",
    "dataset_review_answer_context",
}

_DATASET_REVIEW_OPS = {
    "dataset_review_approve_mappings",
    "dataset_review_set_field_semantics",
    "dataset_review_generate_sql_preview",
}

INTENT_PERMISSION_CHECKS: dict[str, list[tuple[str, str]]] = {
    "get_task_status": [("tasks", "READ")],
    "create_branch": [("plugin:git", "EXECUTE")],
    "commit_changes": [("plugin:git", "EXECUTE")],
    "deploy_dashboard": [("plugin:git", "EXECUTE")],
    "execute_migration": [
        ("plugin:migration", "EXECUTE"),
        ("plugin:superset-migration", "EXECUTE"),
    ],
    "run_backup": [("plugin:superset-backup", "EXECUTE"), ("plugin:backup", "EXECUTE")],
    "run_llm_validation": [("plugin:llm_dashboard_validation", "EXECUTE")],
    "run_llm_documentation": [("plugin:llm_documentation", "EXECUTE")],
    "get_health_summary": [("plugin:migration", "READ")],
    "dataset_review_answer_context": [("dataset:session", "READ")],
    "dataset_review_approve_mappings": [("dataset:session", "MANAGE")],
    "dataset_review_set_field_semantics": [("dataset:session", "MANAGE")],
    "dataset_review_generate_sql_preview": [("dataset:session", "MANAGE")],
}


# #endregion AssistantSchemas