ai-robot-core/ai-service/app/models/entities.py

"""
Memory layer entities for AI Service.
[AC-AISVC-13] SQLModel entities for chat sessions and messages with tenant isolation.
"""

import uuid
from datetime import datetime
from enum import Enum
from typing import Any

from sqlalchemy import Column, JSON
from sqlmodel import Field, Index, SQLModel


class ChatSession(SQLModel, table=True):
    """
    [AC-AISVC-13] Chat session entity with tenant isolation.
    Primary key: (tenant_id, session_id) composite unique constraint.
    """

    __tablename__ = "chat_sessions"
    __table_args__ = (
        Index("ix_chat_sessions_tenant_session", "tenant_id", "session_id", unique=True),
        Index("ix_chat_sessions_tenant_id", "tenant_id"),
    )

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    tenant_id: str = Field(..., description="Tenant ID for multi-tenant isolation", index=True)
    session_id: str = Field(..., description="Session ID for conversation tracking")
    channel_type: str | None = Field(default=None, description="Channel type: wechat, douyin, jd")
    metadata_: dict[str, Any] | None = Field(
        default=None,
        sa_column=Column("metadata", JSON, nullable=True),
        description="Session metadata"
    )
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Session creation time")
    updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last update time")


class ChatMessage(SQLModel, table=True):
    """
    [AC-AISVC-13] Chat message entity with tenant isolation.
    Messages are scoped by (tenant_id, session_id) for multi-tenant security.
    """

    __tablename__ = "chat_messages"
    __table_args__ = (
        Index("ix_chat_messages_tenant_session", "tenant_id", "session_id"),
        Index("ix_chat_messages_tenant_session_created", "tenant_id", "session_id", "created_at"),
    )

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    tenant_id: str = Field(..., description="Tenant ID for multi-tenant isolation", index=True)
    session_id: str = Field(..., description="Session ID for conversation tracking", index=True)
    role: str = Field(..., description="Message role: user or assistant")
    content: str = Field(..., description="Message content")
    prompt_tokens: int | None = Field(default=None, description="Number of prompt tokens used")
    completion_tokens: int | None = Field(default=None, description="Number of completion tokens used")
    total_tokens: int | None = Field(default=None, description="Total tokens used")
    latency_ms: int | None = Field(default=None, description="Response latency in milliseconds")
    first_token_ms: int | None = Field(default=None, description="Time to first token in milliseconds (for streaming)")
    is_error: bool = Field(default=False, description="Whether this message is an error response")
    error_message: str | None = Field(default=None, description="Error message if any")
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Message creation time")


class ChatSessionCreate(SQLModel):
    """Schema for creating a new chat session."""

    tenant_id: str
    session_id: str
    channel_type: str | None = None
    metadata_: dict[str, Any] | None = None


class ChatMessageCreate(SQLModel):
    """Schema for creating a new chat message."""

    tenant_id: str
    session_id: str
    role: str
    content: str


class DocumentStatus(str, Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


class IndexJobStatus(str, Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"


class SessionStatus(str, Enum):
    ACTIVE = "active"
    CLOSED = "closed"
    EXPIRED = "expired"


class KnowledgeBase(SQLModel, table=True):
    """
    [AC-ASA-01] Knowledge base entity with tenant isolation.
    """

    __tablename__ = "knowledge_bases"
    __table_args__ = (
        Index("ix_knowledge_bases_tenant_id", "tenant_id"),
    )

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    tenant_id: str = Field(..., description="Tenant ID for multi-tenant isolation", index=True)
    name: str = Field(..., description="Knowledge base name")
    description: str | None = Field(default=None, description="Knowledge base description")
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation time")
    updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last update time")


class Document(SQLModel, table=True):
    """
    [AC-ASA-01, AC-ASA-08] Document entity with tenant isolation.
    """

    __tablename__ = "documents"
    __table_args__ = (
        Index("ix_documents_tenant_kb", "tenant_id", "kb_id"),
        Index("ix_documents_tenant_status", "tenant_id", "status"),
    )

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    tenant_id: str = Field(..., description="Tenant ID for multi-tenant isolation", index=True)
    kb_id: str = Field(..., description="Knowledge base ID")
    file_name: str = Field(..., description="Original file name")
    file_path: str | None = Field(default=None, description="Storage path")
    file_size: int | None = Field(default=None, description="File size in bytes")
    file_type: str | None = Field(default=None, description="File MIME type")
    status: str = Field(default=DocumentStatus.PENDING.value, description="Document status")
    error_msg: str | None = Field(default=None, description="Error message if failed")
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Upload time")
    updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last update time")


class IndexJob(SQLModel, table=True):
    """
    [AC-ASA-02] Index job entity for tracking document indexing progress.
    """

    __tablename__ = "index_jobs"
    __table_args__ = (
        Index("ix_index_jobs_tenant_doc", "tenant_id", "doc_id"),
        Index("ix_index_jobs_tenant_status", "tenant_id", "status"),
    )

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    tenant_id: str = Field(..., description="Tenant ID for multi-tenant isolation", index=True)
    doc_id: uuid.UUID = Field(..., description="Document ID being indexed")
    status: str = Field(default=IndexJobStatus.PENDING.value, description="Job status")
    progress: int = Field(default=0, ge=0, le=100, description="Progress percentage")
    error_msg: str | None = Field(default=None, description="Error message if failed")
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Job creation time")
    updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last update time")


class KnowledgeBaseCreate(SQLModel):
    """Schema for creating a new knowledge base."""

    tenant_id: str
    name: str
    description: str | None = None


class DocumentCreate(SQLModel):
    """Schema for creating a new document."""

    tenant_id: str
    kb_id: str
    file_name: str
    file_path: str | None = None
    file_size: int | None = None
    file_type: str | None = None