225 lines
8.6 KiB
Python
225 lines
8.6 KiB
Python
"""
|
|
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 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
|
|
|
|
|
|
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
|