Newer
Older
cortex-hub / ai-hub / app / api / schemas.py
from pydantic import BaseModel, Field, ConfigDict
from typing import List, Literal, Optional
from datetime import datetime

# --- User Schemas ---

class OIDCLogin(BaseModel):
    """Schema for OIDC user data received from a login callback."""
    oidc_id: str = Field(..., description="The unique ID from the OIDC provider.")
    email: str = Field(..., description="The user's email address.")
    username: str = Field(..., description="The user's username or display name.")

class UserStatus(BaseModel):
    """Schema for the response when checking a user's status."""
    id: str = Field(..., description="The internal user ID.")
    email: str = Field(..., description="The user's email address.")
    is_logged_in: bool = Field(True, description="Indicates if the user is currently authenticated.")
    is_anonymous: bool = Field(False, description="Indicates if the user is an anonymous user.")

class UserProfile(BaseModel):
    id: str
    email: str
    username: Optional[str] = None
    full_name: Optional[str] = None
    role: str = "user"
    group_id: Optional[str] = None
    group_name: Optional[str] = None
    avatar_url: Optional[str] = None
    created_at: datetime
    last_login_at: Optional[datetime] = None
    model_config = ConfigDict(from_attributes=True)

class UserProfileUpdate(BaseModel):
    username: Optional[str] = None
    full_name: Optional[str] = None
    avatar_url: Optional[str] = None

class UserRoleUpdate(BaseModel):
    role: str

class UserGroupUpdate(BaseModel):
    group_id: str

# --- Group Schemas ---
class GroupBase(BaseModel):
    name: str
    description: Optional[str] = None
    # Policy: {"llm": ["openai"], "tts": ["gcloud"], "stt": ["google"]}
    policy: dict = Field(default_factory=dict)

class GroupCreate(GroupBase):
    pass

class GroupUpdate(BaseModel):
    name: Optional[str] = None
    description: Optional[str] = None
    policy: Optional[dict] = None

class GroupInfo(GroupBase):
    id: str
    created_at: Optional[datetime] = None
    model_config = ConfigDict(from_attributes=True)

# --- General Schemas ---
class UserPreferences(BaseModel):
    """Schema for user-specific LLM, TTS, STT preferences."""
    llm: dict = Field(default_factory=dict)
    tts: dict = Field(default_factory=dict)
    stt: dict = Field(default_factory=dict)
    statuses: Optional[dict] = Field(default_factory=dict)
    
class ConfigResponse(BaseModel):
    """Schema for returning user preferences alongside effective settings."""
    preferences: UserPreferences
    effective: dict = Field(default_factory=dict)

# --- Chat Schemas ---
class ChatRequest(BaseModel):
    """Defines the shape of a request to the /chat endpoint."""
    prompt: str = Field(..., min_length=1)
    # The 'model' can now be specified in the request body to switch models mid-conversation.
    provider_name: str = Field("gemini")
    # Add a new optional boolean field to control the retriever
    load_faiss_retriever: Optional[bool] = Field(False, description="Whether to use the FAISS DB retriever for the chat.")


class ChatResponse(BaseModel):
    """Defines the shape of a successful response from the /chat endpoint."""
    answer: str
    provider_used: str
    message_id: Optional[int] = None

# --- Document Schemas ---
class DocumentCreate(BaseModel):
    title: str
    text: str
    source_url: Optional[str] = None
    author: Optional[str] = None
    user_id: str = "default_user"

class DocumentResponse(BaseModel):
    message: str

class SessionBase(BaseModel):
    user_id: str
    title: Optional[str] = None
    provider_name: Optional[str] = None
    feature_name: Optional[str] = "default"
    status: str
    created_at: datetime
    model_config = ConfigDict(from_attributes=True)

class DocumentInfo(BaseModel):
    id: int
    title: str
    source_url: Optional[str] = None
    status: str
    created_at: datetime
    model_config = ConfigDict(from_attributes=True)

class DocumentListResponse(BaseModel):
    documents: List[DocumentInfo]

class DocumentDeleteResponse(BaseModel):
    message: str
    document_id: int

# --- Session Schemas ---
class SessionCreate(BaseModel):
    """Defines the shape for starting a new conversation session."""
    user_id: str
    provider_name: str = "deepseek"
    stt_provider_name: Optional[str] = None
    tts_provider_name: Optional[str] = None
    feature_name: Optional[str] = "default"

class SessionUpdate(BaseModel):
    title: Optional[str] = None
    provider_name: Optional[str] = None
    stt_provider_name: Optional[str] = None
    tts_provider_name: Optional[str] = None

class Session(BaseModel):
    """Defines the shape of a session object returned by the API."""
    id: int
    user_id: str
    title: Optional[str] = None
    provider_name: Optional[str] = None
    stt_provider_name: Optional[str] = None
    tts_provider_name: Optional[str] = None
    feature_name: str
    created_at: datetime
    # M3: Node attachment
    sync_workspace_id: Optional[str] = None
    attached_node_ids: List[str] = []
    node_sync_status: dict = {}
    model_config = ConfigDict(from_attributes=True)

# --- M3: Session Node Attachment Schemas ---

class NodeAttachRequest(BaseModel):
    """Attach one or more nodes to a session."""
    node_ids: List[str]

class NodeSyncStatusEntry(BaseModel):
    """Per-node sync status within a session."""
    node_id: str
    status: str           # 'pending' | 'syncing' | 'synced' | 'error'
    last_sync: Optional[str] = None
    error: Optional[str] = None

class SessionNodeStatusResponse(BaseModel):
    """Response showing all attached nodes and their sync state."""
    session_id: int
    sync_workspace_id: Optional[str] = None
    nodes: List[NodeSyncStatusEntry] = []

# --- M4: Node Config YAML ---

class NodeConfigYamlResponse(BaseModel):
    """The generated config YAML content an admin downloads to set up a node."""
    node_id: str
    config_yaml: str   # Full YAML string ready to save as agent_config.yaml

class Message(BaseModel):
    """Defines the shape of a single message within a session's history."""
    id: int
    # The sender can only be one of two roles.
    sender: Literal["user", "assistant"]
    # The text content of the message.
    content: str
    # The timestamp for when the message was created.
    created_at: datetime
    # URL to the saved audio file
    audio_url: Optional[str] = None
    # Whether audio exists for this message
    has_audio: bool = False

    # Enables creating this schema from a SQLAlchemy database object.
    model_config = ConfigDict(from_attributes=True)

class MessageHistoryResponse(BaseModel):
    """Defines the response for retrieving a session's chat history."""
    session_id: int
    messages: List[Message]

class SessionTokenUsageResponse(BaseModel):
    """Defines the response for retrieving a session's token usage."""
    token_count: int
    token_limit: int
    percentage: float

class SpeechRequest(BaseModel):
    text: str

# --- STT Schemas ---
class STTResponse(BaseModel):
    """Defines the shape of a successful response from the /stt endpoint."""
    transcript: str

class VerifyProviderRequest(BaseModel):
    provider_name: str
    provider_type: Optional[str] = None
    api_key: Optional[str] = None
    model: Optional[str] = None
    voice: Optional[str] = None

class VerifyProviderResponse(BaseModel):
    success: bool
    message: str

class ModelInfoResponse(BaseModel):
    model_name: str
    max_tokens: Optional[int] = None
    max_input_tokens: Optional[int] = None


# ---------------------------------------------------------------------------
# Agent Node Schemas
# ---------------------------------------------------------------------------

# --- Skill Toggles (admin-configured) ---

class SkillConfig(BaseModel):
    """Per-skill enable/disable with optional config."""
    enabled: bool = True
    cwd_jail: Optional[str] = None          # shell only: restrict working directory
    max_file_size_mb: Optional[int] = None  # sync only: file size cap

class NodeSkillConfig(BaseModel):
    """Admin-controlled skill configuration for a node."""
    shell: SkillConfig = SkillConfig(enabled=True)
    browser: SkillConfig = SkillConfig(enabled=True)
    sync: SkillConfig = SkillConfig(enabled=True)

# --- Admin Create / Update ---

class AgentNodeCreate(BaseModel):
    """Payload for admin creating a new node registration."""
    node_id: str = Field(..., description="Stable identifier used in the node's config YAML, e.g. 'dev-macbook-m3'")
    display_name: str = Field(..., description="Human-readable name shown in the UI")
    description: Optional[str] = None
    skill_config: NodeSkillConfig = NodeSkillConfig()

class AgentNodeUpdate(BaseModel):
    """Payload for admin updating node configuration."""
    display_name: Optional[str] = None
    description: Optional[str] = None
    skill_config: Optional[NodeSkillConfig] = None
    is_active: Optional[bool] = None

# --- Group Access ---

class NodeAccessGrant(BaseModel):
    """Admin grants a group access to a node."""
    group_id: str
    access_level: str = Field("use", description="'view', 'use', or 'admin'")

class NodeAccessResponse(BaseModel):
    id: int
    node_id: str
    group_id: str
    access_level: str
    granted_at: datetime
    model_config = ConfigDict(from_attributes=True)

# --- Live Stats ---

class AgentNodeStats(BaseModel):
    """Live performance stats reported via heartbeat."""
    active_worker_count: int = 0
    cpu_usage_percent: float = 0.0
    memory_usage_percent: float = 0.0
    running: List[str] = []

# --- Node Responses ---

class AgentNodeAdminDetail(BaseModel):
    """Full node detail for admin view — includes invite_token and skill config."""
    node_id: str
    display_name: str
    description: Optional[str] = None
    skill_config: dict = {}
    capabilities: dict = {}
    invite_token: Optional[str] = None
    is_active: bool = True
    last_status: str
    last_seen_at: Optional[datetime] = None
    created_at: datetime
    registered_by: str
    group_access: List[NodeAccessResponse] = []
    stats: AgentNodeStats = AgentNodeStats()
    model_config = ConfigDict(from_attributes=True)

class AgentNodeUserView(BaseModel):
    """Node as seen by a user — no invite_token, no admin config details."""
    node_id: str
    display_name: str
    description: Optional[str] = None
    capabilities: dict = {}
    # Which skills are available to this user (derived from skill_config.enabled)
    available_skills: List[str] = []
    last_status: str    # 'online' | 'offline' | 'stale'
    last_seen_at: Optional[datetime] = None
    model_config = ConfigDict(from_attributes=True)

class AgentNodeStatusResponse(BaseModel):
    """Full live status of an agent node (used internally)."""
    node_id: str
    display_name: Optional[str] = None
    stats: AgentNodeStats = AgentNodeStats()
    status: str
    connected_at: Optional[str] = None
    last_heartbeat_at: Optional[str] = None

# --- User Node Preferences ---

class NodeDataSourceConfig(BaseModel):
    """How a node should seed its workspace for a session."""
    source: str = Field("empty", description="'empty' | 'server' | 'node_local'")
    path: Optional[str] = None  # root path on node when source='node_local'

class UserNodePreferences(BaseModel):
    """Stored in User.preferences['nodes']."""
    default_node_ids: List[str] = Field(
        default_factory=list,
        description="Node IDs auto-attached when starting a new session"
    )
    data_source: NodeDataSourceConfig = NodeDataSourceConfig()

# --- Task Dispatch ---

class NodeDispatchRequest(BaseModel):
    """Dispatch a shell or browser action to a specific node."""
    command: str = ""
    browser_action: Optional[dict] = None
    session_id: Optional[str] = None
    timeout_ms: int = 30000

class NodeDispatchResponse(BaseModel):
    task_id: str
    status: str   # 'accepted' | 'rejected'
    reason: Optional[str] = None

# Keep backward-compat alias
AgentNodeSummary = AgentNodeUserView