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 Tenant(SQLModel, table=True):
    """
    [AC-AISVC-10] Tenant entity for storing tenant information.
    Tenant ID format: name@ash@year (e.g., szmp@ash@2026)
    """

    __tablename__ = "tenants"

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    tenant_id: str = Field(..., description="Full tenant ID (format: name@ash@year)", unique=True, index=True)
    name: str = Field(..., description="Tenant display name (first part of tenant_id)")
    year: str = Field(..., description="Year part from tenant_id")
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation time")
    updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last update time")


class KBType(str, Enum):
    PRODUCT = "product"
    FAQ = "faq"
    SCRIPT = "script"
    POLICY = "policy"
    GENERAL = "general"


class KnowledgeBase(SQLModel, table=True):
    """
    [AC-ASA-01, AC-AISVC-59] Knowledge base entity with tenant isolation.
    [v0.6.0] Extended with kb_type, priority, is_enabled, doc_count for multi-KB management.
    """

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

    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")
    kb_type: str = Field(default=KBType.GENERAL.value, description="Knowledge base type: product/faq/script/policy/general")
    description: str | None = Field(default=None, description="Knowledge base description")
    priority: int = Field(default=0, ge=0, description="Priority weight, higher value means higher priority")
    is_enabled: bool = Field(default=True, description="Whether the knowledge base is enabled")
    doc_count: int = Field(default=0, ge=0, description="Document count (cached statistic)")
    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."""

    name: str
    kb_type: str = KBType.GENERAL.value
    description: str | None = None
    priority: int = 0


class KnowledgeBaseUpdate(SQLModel):
    """Schema for updating a knowledge base."""

    name: str | None = None
    kb_type: str | None = None
    description: str | None = None
    priority: int | None = None
    is_enabled: bool | 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


class ApiKey(SQLModel, table=True):
    """
    [AC-AISVC-50] API Key entity for lightweight authentication.
    Keys are loaded into memory on startup for fast validation.
    """

    __tablename__ = "api_keys"

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    key: str = Field(..., description="API Key (unique)", unique=True, index=True)
    name: str = Field(..., description="Key name/description for identification")
    is_active: bool = Field(default=True, description="Whether the key is active")
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation time")
    updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last update time")


class ApiKeyCreate(SQLModel):
    """Schema for creating a new API key."""

    key: str
    name: str
    is_active: bool = True


class TemplateVersionStatus(str, Enum):
    DRAFT = "draft"
    PUBLISHED = "published"
    ARCHIVED = "archived"


class PromptTemplate(SQLModel, table=True):
    """
    [AC-AISVC-51, AC-AISVC-52] Prompt template entity with tenant isolation.
    Main table for storing template metadata.
    """

    __tablename__ = "prompt_templates"
    __table_args__ = (
        Index("ix_prompt_templates_tenant_scene", "tenant_id", "scene"),
        Index("ix_prompt_templates_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="Template name, e.g., 'Default Customer Service Persona'")
    scene: str = Field(..., description="Scene tag: chat/rag_qa/greeting/farewell")
    description: str | None = Field(default=None, description="Template description")
    is_default: bool = Field(default=False, description="Whether this is the default template for the scene")
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation time")
    updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last update time")


class PromptTemplateVersion(SQLModel, table=True):
    """
    [AC-AISVC-53] Prompt template version entity.
    Stores versioned content with status management.
    """

    __tablename__ = "prompt_template_versions"
    __table_args__ = (
        Index("ix_template_versions_template_status", "template_id", "status"),
    )

    id: uuid.UUID = Field(default_factory=uuid.uuid4, primary_key=True)
    template_id: uuid.UUID = Field(..., description="Foreign key to prompt_templates.id", foreign_key="prompt_templates.id", index=True)
    version: int = Field(..., description="Version number (auto-incremented per template)")
    status: str = Field(default=TemplateVersionStatus.DRAFT.value, description="Version status: draft/published/archived")
    system_instruction: str = Field(..., description="System instruction content with {{variable}} placeholders")
    variables: list[dict[str, Any]] | None = Field(
        default=None,
        sa_column=Column("variables", JSON, nullable=True),
        description="Variable definitions, e.g., [{'name': 'persona_name', 'default': '小N', 'description': '人设名称'}]"
    )
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation time")


class PromptTemplateCreate(SQLModel):
    """Schema for creating a new prompt template."""

    name: str
    scene: str
    description: str | None = None
    system_instruction: str
    variables: list[dict[str, Any]] | None = None
    is_default: bool = False


class PromptTemplateUpdate(SQLModel):
    """Schema for updating a prompt template."""

    name: str | None = None
    scene: str | None = None
    description: str | None = None
    system_instruction: str | None = None
    variables: list[dict[str, Any]] | None = None
    is_default: bool | None = None


class ResponseType(str, Enum):
    """[AC-AISVC-65] Response type for intent rules."""
    FIXED = "fixed"
    RAG = "rag"
    FLOW = "flow"
    TRANSFER = "transfer"


class IntentRule(SQLModel, table=True):
    """
    [AC-AISVC-65] Intent rule entity with tenant isolation.
    Supports keyword and regex matching for intent recognition.
    """

    __tablename__ = "intent_rules"
    __table_args__ = (
        Index("ix_intent_rules_tenant_enabled_priority", "tenant_id", "is_enabled"),
        Index("ix_intent_rules_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="Intent name, e.g., 'Return Intent'")
    keywords: list[str] | None = Field(
        default=None,
        sa_column=Column("keywords", JSON, nullable=True),
        description="Keyword list for matching, e.g., ['return', 'refund']"
    )
    patterns: list[str] | None = Field(
        default=None,
        sa_column=Column("patterns", JSON, nullable=True),
        description="Regex pattern list for matching, e.g., ['return.*goods', 'how to return']"
    )
    priority: int = Field(default=0, description="Priority (higher value = higher priority)")
    response_type: str = Field(..., description="Response type: fixed/rag/flow/transfer")
    target_kb_ids: list[str] | None = Field(
        default=None,
        sa_column=Column("target_kb_ids", JSON, nullable=True),
        description="Target knowledge base IDs for rag type"
    )
    flow_id: uuid.UUID | None = Field(default=None, description="Flow ID for flow type")
    fixed_reply: str | None = Field(default=None, description="Fixed reply content for fixed type")
    transfer_message: str | None = Field(default=None, description="Transfer message for transfer type")
    is_enabled: bool = Field(default=True, description="Whether the rule is enabled")
    hit_count: int = Field(default=0, description="Hit count for statistics")
    created_at: datetime = Field(default_factory=datetime.utcnow, description="Creation time")
    updated_at: datetime = Field(default_factory=datetime.utcnow, description="Last update time")


class IntentRuleCreate(SQLModel):
    """[AC-AISVC-65] Schema for creating a new intent rule."""

    name: str
    keywords: list[str] | None = None
    patterns: list[str] | None = None
    priority: int = 0
    response_type: str
    target_kb_ids: list[str] | None = None
    flow_id: str | None = None
    fixed_reply: str | None = None
    transfer_message: str | None = None


class IntentRuleUpdate(SQLModel):
    """[AC-AISVC-67] Schema for updating an intent rule."""

    name: str | None = None
    keywords: list[str] | None = None
    patterns: list[str] | None = None
    priority: int | None = None
    response_type: str | None = None
    target_kb_ids: list[str] | None = None
    flow_id: str | None = None
    fixed_reply: str | None = None
    transfer_message: str | None = None
    is_enabled: bool | None = None


class IntentMatchResult:
    """
    [AC-AISVC-69] Result of intent matching.
    Contains the matched rule and match details.
    """

    def __init__(
        self,
        rule: IntentRule,
        match_type: str,
        matched: str,
    ):
        self.rule = rule
        self.match_type = match_type
        self.matched = matched

    def to_dict(self) -> dict[str, Any]:
        return {
            "rule_id": str(self.rule.id),
            "rule_name": self.rule.name,
            "match_type": self.match_type,
            "matched": self.matched,
            "response_type": self.rule.response_type,
        }