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."""
task_id: Optional[str] = None # NEW: Support client-side generated task IDs
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