planet/backend/app/schemas/ai.py

from typing import Any

from pydantic import BaseModel, Field

from app.core.enums import PlaygroundMessageKind, PlaygroundMessageRole, PlaygroundMessageStatus

class AIContentBlock(BaseModel):
    type: str
    text: str | None = None
    thinking: str | None = None
    signature: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)


class SituationalAnalysisRequest(BaseModel):
    title: str = Field(..., min_length=1, max_length=200)
    objective: str = Field(..., min_length=1, max_length=20000)
    context: dict[str, Any] = Field(default_factory=dict)
    observations: list[str] = Field(default_factory=list)
    constraints: list[str] = Field(default_factory=list)
    system_prompt: str | None = Field(default=None, max_length=8000)
    preferred_model: str | None = Field(default=None, max_length=200)
    thinking: dict[str, Any] | None = None


class BGPBriefRequest(BaseModel):
    incident_limit: int = Field(default=5, ge=1, le=10)
    anomaly_limit: int = Field(default=6, ge=1, le=12)
    collector_limit: int = Field(default=5, ge=1, le=10)
    preferred_model: str | None = Field(default=None, max_length=200)
    thinking: dict[str, Any] | None = None


class AlertBriefRequest(BaseModel):
    alert_limit: int = Field(default=8, ge=1, le=20)
    preferred_model: str | None = Field(default=None, max_length=200)
    thinking: dict[str, Any] | None = None


class SituationalAlertBriefRequest(BaseModel):
    preferred_model: str | None = Field(default=None, max_length=200)
    thinking: dict[str, Any] | None = None


class SituationalAnalysisResponse(BaseModel):
    provider: str
    model: str
    content: str
    content_blocks: list[AIContentBlock] = Field(default_factory=list)
    text_blocks: list[str] = Field(default_factory=list)
    thinking_blocks: list[str] = Field(default_factory=list)
    raw_response: dict[str, Any] = Field(default_factory=dict)


class BGPBriefRecordSummary(BaseModel):
    id: str
    title: str
    provider: str
    model: str
    request_id: str | None = None
    generated_at: str


class BGPBriefRecordResponse(BGPBriefRecordSummary):
    content_markdown: str
    facts: list[str] = Field(default_factory=list)
    context: dict[str, Any] = Field(default_factory=dict)


class AlertBriefResponse(SituationalAnalysisResponse):
    title: str
    objective: str
    facts: list[str] = Field(default_factory=list)
    context: dict[str, Any] = Field(default_factory=dict)


class SituationalAlertBriefResponse(SituationalAnalysisResponse):
    title: str
    objective: str
    facts: list[str] = Field(default_factory=list)
    context: dict[str, Any] = Field(default_factory=dict)


class AIProviderStatusResponse(BaseModel):
    provider: str
    api: str | None = None
    enabled: bool
    configured: bool
    model: str | None = None
    base_url: str | None = None


class PlaygroundSessionState(BaseModel):
    messages: list[dict[str, Any]] = Field(default_factory=list)
    selectedPresetKey: str = Field(default="bgp-brief", max_length=100)
    title: str = Field(default="", max_length=200)
    objective: str = Field(default="", max_length=1000)
    constraints: str = Field(default="")
    inputValue: str = Field(default="")
    analysis: dict[str, Any] | None = None
    latestAnalysisMessageId: str | None = Field(default=None, max_length=200)
    analysisMeta: dict[str, Any] = Field(default_factory=dict)
    helpExpanded: bool = True


class PlaygroundSessionUpsertRequest(BaseModel):
    session_key: str = Field(default="default", min_length=1, max_length=100)
    title: str | None = Field(default=None, max_length=200)
    state: PlaygroundSessionState


class PlaygroundMessageRecord(BaseModel):
    id: str
    role: PlaygroundMessageRole
    kind: PlaygroundMessageKind = PlaygroundMessageKind.MESSAGE
    status: PlaygroundMessageStatus = PlaygroundMessageStatus.DONE
    title: str | None = None
    content: str = ""
    thinking_content: str = ""
    meta: list[str] = Field(default_factory=list)
    markdown: bool = True
    provider: str | None = None
    model: str | None = None
    request_id: str | None = None
    raw_response: dict[str, Any] = Field(default_factory=dict)
    content_blocks: list[dict[str, Any]] = Field(default_factory=list)
    text_blocks: list[str] = Field(default_factory=list)
    thinking_blocks: list[str] = Field(default_factory=list)
    parent_message_id: str | None = None
    created_at: str
    updated_at: str


class PlaygroundSessionResponse(BaseModel):
    id: str
    session_key: str
    title: str
    state: PlaygroundSessionState
    created_at: str
    updated_at: str


class PlaygroundThreadResponse(BaseModel):
    session: PlaygroundSessionResponse
    messages: list[PlaygroundMessageRecord] = Field(default_factory=list)


class PlaygroundMessageCreateRequest(BaseModel):
    session_key: str = Field(default="default", min_length=1, max_length=100)
    title: str = Field(..., min_length=1, max_length=200)
    objective: str = Field(..., min_length=1, max_length=1000)
    constraints: str = Field(default="")
    input: str = Field(..., min_length=1)
    selected_preset_key: str = Field(default="bgp-brief", max_length=100)
    help_expanded: bool = True


class PlaygroundMessageActionResponse(BaseModel):
    session: PlaygroundSessionResponse
    messages: list[PlaygroundMessageRecord] = Field(default_factory=list)
    active_message_id: str | None = None


class PlaygroundMessageStopRequest(BaseModel):
    session_key: str = Field(default="default", min_length=1, max_length=100)
    message_id: str = Field(..., min_length=1, max_length=64)


class PlaygroundMessageResendRequest(BaseModel):
    session_key: str = Field(default="default", min_length=1, max_length=100)
    user_message_id: str = Field(..., min_length=1, max_length=64)


class PlaygroundMessageEditRequest(BaseModel):
    session_key: str = Field(default="default", min_length=1, max_length=100)
    user_message_id: str = Field(..., min_length=1, max_length=64)
    content: str = Field(..., min_length=1)