api/app/schemas.py

"""Pydantic schemas — API request/response contracts.

Mirrors the TypeScript types from the Electron app (src/shared/api-types.ts).
"""

from __future__ import annotations

from enum import Enum
from typing import Any, Literal

from pydantic import BaseModel, Field


# ── Billing ──────────────────────────────────────────────────────────

BillingTier = Literal["free", "pro", "power", "team"]


# ── Auth ─────────────────────────────────────────────────────────────

class AuthTokens(BaseModel):
    access_token: str
    refresh_token: str
    expires_at: int


class UserProfile(BaseModel):
    id: str
    email: str
    tier: BillingTier


# ── Chat ─────────────────────────────────────────────────────────────

class ChatContext(BaseModel):
    user_profile: dict[str, Any] = Field(default_factory=dict)
    relevant_documents: list[str] = Field(default_factory=list)
    recent_tasks: list[dict[str, Any]] = Field(default_factory=list)
    conversation_history: list[dict[str, Any]] = Field(default_factory=list)


class PlanAction(BaseModel):
    type: Literal[
        "create_record",
        "update_record",
        "delete_record",
        "index_document",
        "send_notification",
    ]
    table: str | None = None
    data: dict[str, Any] | None = None


class ChatRequest(BaseModel):
    message: str
    context: ChatContext = Field(default_factory=ChatContext)
    execution_mode: Literal["direct", "plan"] = "direct"


class ChatResponse(BaseModel):
    response: str
    actions: list[PlanAction] = Field(default_factory=list)


# ── Execution Plans ──────────────────────────────────────────────────

class PlanStep(BaseModel):
    action: str
    prompt_template: str | None = None
    variables: dict[str, Any] | None = None
    data_from_step: int | None = None


class ExecutionPlan(BaseModel):
    agent: str
    steps: list[PlanStep] = Field(default_factory=list)


# ── Backup ───────────────────────────────────────────────────────────

class BackupMetadata(BaseModel):
    version: int
    timestamp: int
    checksum: str
    chunk_count: int


# ── Cloud Storage (E2E encrypted blobs) ──────────────────────────────

class StorageRecord(BaseModel):
    id: str
    user_id: str
    table: str
    blob: bytes
    checksum: str
    created_at: int
    updated_at: int


class StorageRecordCreate(BaseModel):
    table: str
    blob: bytes
    checksum: str


class StorageRecordUpdate(BaseModel):
    blob: bytes
    checksum: str


# ── Cloud Vector Store (E2E encrypted vectors) ────────────────────────

class VectorItem(BaseModel):
    id: str
    blob: bytes   # encrypted vector + metadata — backend never decrypts
    checksum: str


class VectorUpsertRequest(BaseModel):
    vectors: list[VectorItem]


class VectorSearchRequest(BaseModel):
    query_blob: bytes   # encrypted query — backend never decrypts
    top_k: int = 10


class VectorSearchResult(BaseModel):
    id: str
    score: float
    blob: bytes


class VectorSearchResponse(BaseModel):
    results: list[VectorSearchResult]


# ── Plugin Marketplace ────────────────────────────────────────────────

class PluginManifest(BaseModel):
    id: str
    name: str
    description: str
    version: str
    author: str
    permissions: list[str]
    category: str
    price_cents: int = 0


class PluginListResponse(BaseModel):
    plugins: list[PluginManifest]
    total: int
    page: int


class PluginInstallRequest(BaseModel):
    plugin_id: str


# ── WebSocket Frame Protocol ──────────────────────────────────────────

class WsFrameType(str, Enum):
    chat_request = "chat_request"
    text_chunk = "text_chunk"
    tool_call = "tool_call"
    tool_result = "tool_result"
    final = "final"
    ping = "ping"


class WsToolCall(BaseModel):
    """Server → Client: requests a CRUD/vector operation on the local DB."""

    type: Literal[WsFrameType.tool_call] = WsFrameType.tool_call
    id: str
    action: str
    table: str | None = None
    data: dict[str, Any] | None = None
    filters: dict[str, Any] | None = None
    vector: list[float] | None = None
    limit: int | None = None


class WsToolResult(BaseModel):
    """Client → Server: result of a CRUD/vector operation."""

    type: Literal[WsFrameType.tool_result] = WsFrameType.tool_result
    id: str
    row: dict[str, Any] | None = None
    rows: list[dict[str, Any]] | None = None
    results: list[dict[str, Any]] | None = None
    deleted: bool | None = None
    ok: bool | None = None
    error: str | None = None


class WsTextChunk(BaseModel):
    """Server → Client: incremental LLM response text."""

    type: Literal[WsFrameType.text_chunk] = WsFrameType.text_chunk
    text: str


class WsFinal(BaseModel):
    """Server → Client: signals end of response with the complete text."""

    type: Literal[WsFrameType.final] = WsFrameType.final
    response: str