191 lines
6.0 KiB
Python
191 lines
6.0 KiB
Python
"""RAG 平台 Pydantic 数据模型。
|
||
|
||
这些是领域模型（非 ORM），ORM 模型定义在 store.py（U2）。
|
||
与 memory/knowledge_base.py 的 KnowledgeBase protocol 分离 —
|
||
rag_platform 服务企业知识库场景，memory/ 服务 Agent 运行时记忆。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import uuid
|
||
from datetime import datetime, timezone
|
||
from enum import Enum
|
||
|
||
from pydantic import BaseModel, ConfigDict, Field
|
||
|
||
|
||
def _utcnow() -> datetime:
|
||
return datetime.now(timezone.utc)
|
||
|
||
|
||
def _uuid() -> str:
|
||
return str(uuid.uuid4())
|
||
|
||
|
||
class KBStatus(str, Enum):
|
||
active = "active"
|
||
archived = "archived"
|
||
|
||
|
||
class DocumentStatus(str, Enum):
|
||
"""文档处理状态机：pending → parsing → segmenting → vectorizing → indexed | failed。"""
|
||
|
||
pending = "pending"
|
||
parsing = "parsing"
|
||
segmenting = "segmenting"
|
||
vectorizing = "vectorizing"
|
||
indexed = "indexed"
|
||
failed = "failed"
|
||
|
||
|
||
class QueryMode(str, Enum):
|
||
"""检索模式：embedding 语义 / keywords 全文 / blend 双索引合并。"""
|
||
|
||
embedding = "embedding"
|
||
keywords = "keywords"
|
||
blend = "blend"
|
||
|
||
|
||
class KnowledgeBase(BaseModel):
|
||
"""知识库领域模型。"""
|
||
|
||
model_config = ConfigDict(from_attributes=True)
|
||
|
||
id: str = Field(default_factory=_uuid)
|
||
name: str
|
||
description: str = ""
|
||
owner: str
|
||
status: KBStatus = KBStatus.active
|
||
# 检索与命中处理默认配置（Agent 运行时可覆盖）
|
||
default_query_mode: QueryMode = QueryMode.blend
|
||
default_hit_processing: str = "model_opt" # model_opt | direct
|
||
caching_disabled: bool = False
|
||
created_at: datetime = Field(default_factory=_utcnow)
|
||
updated_at: datetime = Field(default_factory=_utcnow)
|
||
|
||
|
||
class Document(BaseModel):
|
||
"""知识库文档领域模型。"""
|
||
|
||
model_config = ConfigDict(from_attributes=True)
|
||
|
||
id: str = Field(default_factory=_uuid)
|
||
kb_id: str
|
||
filename: str
|
||
file_type: str
|
||
file_size: int
|
||
status: DocumentStatus = DocumentStatus.pending
|
||
error_message: str | None = None
|
||
created_at: datetime = Field(default_factory=_utcnow)
|
||
updated_at: datetime = Field(default_factory=_utcnow)
|
||
|
||
|
||
class Chunk(BaseModel):
|
||
"""文档分段后的文本块。"""
|
||
|
||
model_config = ConfigDict(from_attributes=True)
|
||
|
||
id: str = Field(default_factory=_uuid)
|
||
document_id: str
|
||
kb_id: str
|
||
content: str
|
||
metadata: dict = Field(default_factory=dict)
|
||
embedding: list[float] | None = None
|
||
|
||
|
||
class QueryResult(BaseModel):
|
||
"""检索结果条目。"""
|
||
|
||
model_config = ConfigDict(from_attributes=True)
|
||
|
||
chunk_id: str
|
||
content: str
|
||
score: float
|
||
metadata: dict = Field(default_factory=dict)
|
||
document_id: str
|
||
kb_id: str
|
||
|
||
|
||
class RetrievalRequest(BaseModel):
|
||
"""检索请求 — 支持覆盖 KB 默认配置。
|
||
|
||
retrieval_mode / hit_processing_mode 为 None 时使用 KB 默认值。
|
||
"""
|
||
|
||
model_config = ConfigDict()
|
||
|
||
query: str
|
||
kb_ids: list[str]
|
||
retrieval_mode: QueryMode | None = None # None = 使用 KB 默认
|
||
hit_processing_mode: str | None = None # None = 使用 KB 默认
|
||
top_k: int = 5
|
||
user_id: str | None = None # 用于 ACL 过滤
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# ORM Models (SQLAlchemy 2 DeclarativeBase + Mapped)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
from sqlalchemy import DateTime, ForeignKey, String, Text, UniqueConstraint # noqa: E402
|
||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column # noqa: E402
|
||
|
||
|
||
class ORMBase(DeclarativeBase):
|
||
"""rag_platform ORM 基类 — 与 memory/models.py 的 declarative_base() 独立。"""
|
||
|
||
pass
|
||
|
||
|
||
class KBModel(ORMBase):
|
||
"""知识库 ORM 模型。"""
|
||
|
||
__tablename__ = "rag_platform_kbs"
|
||
|
||
id: Mapped[str] = mapped_column(String, primary_key=True, default=_uuid)
|
||
name: Mapped[str] = mapped_column(String(255))
|
||
description: Mapped[str] = mapped_column(Text, default="")
|
||
owner: Mapped[str] = mapped_column(String(255), index=True)
|
||
status: Mapped[str] = mapped_column(String(32), default="active", index=True)
|
||
default_query_mode: Mapped[str] = mapped_column(String(32), default="blend")
|
||
default_hit_processing: Mapped[str] = mapped_column(String(32), default="model_opt")
|
||
caching_disabled: Mapped[bool] = mapped_column(default=False)
|
||
created_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)
|
||
updated_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)
|
||
|
||
|
||
class DocumentModel(ORMBase):
|
||
"""知识库文档 ORM 模型。"""
|
||
|
||
__tablename__ = "rag_platform_documents"
|
||
|
||
id: Mapped[str] = mapped_column(String, primary_key=True, default=_uuid)
|
||
kb_id: Mapped[str] = mapped_column(
|
||
String, ForeignKey("rag_platform_kbs.id", ondelete="CASCADE"), index=True
|
||
)
|
||
filename: Mapped[str] = mapped_column(String(512))
|
||
file_type: Mapped[str] = mapped_column(String(64))
|
||
file_size: Mapped[int] = mapped_column(default=0)
|
||
status: Mapped[str] = mapped_column(String(32), default="pending", index=True)
|
||
error_message: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||
created_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)
|
||
updated_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)
|
||
|
||
|
||
class KBAclModel(ORMBase):
|
||
"""Per-KB 访问控制 ORM 模型。
|
||
|
||
kb_id FK → rag_platform_kbs.id ON DELETE CASCADE（KTD5 不变量）。
|
||
ACL 条目与 KB 元数据共享事务边界（同一 PG DB）。
|
||
"""
|
||
|
||
__tablename__ = "rag_platform_kb_acl"
|
||
__table_args__ = (UniqueConstraint("kb_id", "user_id", name="uq_kb_acl_kb_user"),)
|
||
|
||
id: Mapped[str] = mapped_column(String, primary_key=True, default=_uuid)
|
||
kb_id: Mapped[str] = mapped_column(
|
||
String, ForeignKey("rag_platform_kbs.id", ondelete="CASCADE"), index=True
|
||
)
|
||
user_id: Mapped[str] = mapped_column(String(255), index=True)
|
||
role: Mapped[str] = mapped_column(String(32), default="viewer") # owner | viewer
|
||
created_at: Mapped[datetime] = mapped_column(DateTime, default=_utcnow)
|