langchain-learning-kit/app/db/models.py

"""
LangChain Learning Kit 的 SQLAlchemy 数据库模型。
"""
from datetime import datetime
from typing import Optional

from sqlalchemy import (
    Boolean,
    BigInteger,
    Column,
    DateTime,
    ForeignKey,
    Integer,
    String,
    Text,
    JSON,
    Index,
)
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import relationship

Base = declarative_base()


class Model(Base):
    """
    LLM 和 Embedding 模型的配置表。
    """

    __tablename__ = "models"

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String(128), nullable=False, unique=True)
    type = Column(String(32), nullable=False)  # 'llm' 或 'embedding'
    config = Column(JSON, nullable=False)
    is_default = Column(Boolean, default=False)
    status = Column(String(32), default="active")
    created_at = Column(DateTime, default=datetime.utcnow)

    __table_args__ = (
        Index("idx_models_type", "type"),
        Index("idx_models_status", "status"),
    )

    def __repr__(self) -> str:
        return f"<Model(id={self.id}, name='{self.name}', type='{self.type}')>"


class KnowledgeBase(Base):
    """
    知识库表。
    """

    __tablename__ = "knowledge_bases"

    id = Column(Integer, primary_key=True, autoincrement=True)
    name = Column(String(128), unique=True, nullable=False)
    description = Column(Text)
    created_at = Column(DateTime, default=datetime.utcnow)

    # 关系
    documents = relationship("Document", back_populates="kb", cascade="all, delete-orphan")

    def __repr__(self) -> str:
        return f"<KnowledgeBase(id={self.id}, name='{self.name}')>"


class Document(Base):
    """
    知识库的文档元数据表。
    """

    __tablename__ = "documents"

    id = Column(BigInteger, primary_key=True, autoincrement=True)
    kb_id = Column(Integer, ForeignKey("knowledge_bases.id"), nullable=False)
    title = Column(String(512))
    content = Column(Text, nullable=False)
    source = Column(String(512))
    doc_metadata = Column(JSON)  # 从 'metadata' 重命名以避免 SQLAlchemy 保留字
    embedding_id = Column(String(128))  # 链接到向量存储 ID
    created_at = Column(DateTime, default=datetime.utcnow)

    # 关系
    kb = relationship("KnowledgeBase", back_populates="documents")

    __table_args__ = (Index("idx_documents_kb_id", "kb_id"),)

    def __repr__(self) -> str:
        return f"<Document(id={self.id}, kb_id={self.kb_id}, title='{self.title}')>"


class Conversation(Base):
    """
    多轮对话的会话表。
    """

    __tablename__ = "conversations"

    id = Column(BigInteger, primary_key=True, autoincrement=True)
    user_id = Column(Integer)  # 为将来的认证实现保留
    title = Column(String(255))
    created_at = Column(DateTime, default=datetime.utcnow)

    # 关系
    messages = relationship("Message", back_populates="conversation", cascade="all, delete-orphan")

    def __repr__(self) -> str:
        return f"<Conversation(id={self.id}, title='{self.title}')>"


class Message(Base):
    """
    对话历史的消息表。
    """

    __tablename__ = "messages"

    id = Column(BigInteger, primary_key=True, autoincrement=True)
    conversation_id = Column(BigInteger, ForeignKey("conversations.id"), nullable=False)
    role = Column(String(32), nullable=False)  # 'user', 'assistant', 'system', 'tool'
    content = Column(Text, nullable=False)
    msg_metadata = Column(JSON)  # 从 'metadata' 重命名以避免 SQLAlchemy 保留字
    created_at = Column(DateTime, default=datetime.utcnow)

    # 关系
    conversation = relationship("Conversation", back_populates="messages")

    __table_args__ = (Index("idx_messages_conv_id", "conversation_id"),)

    def __repr__(self) -> str:
        return f"<Message(id={self.id}, conversation_id={self.conversation_id}, role='{self.role}')>"


class ToolCall(Base):
    """
    智能体执行的工具调用审计日志。
    """

    __tablename__ = "tool_calls"

    id = Column(BigInteger, primary_key=True, autoincrement=True)
    agent_id = Column(String(128), nullable=False)  # 运行 ID 或智能体标识符
    tool_name = Column(String(128), nullable=False)
    call_input = Column(JSON)
    call_output = Column(JSON)
    created_at = Column(DateTime, default=datetime.utcnow)

    __table_args__ = (Index("idx_tool_calls_agent_id", "agent_id"),)

    def __repr__(self) -> str:
        return f"<ToolCall(id={self.id}, agent_id='{self.agent_id}', tool='{self.tool_name}')>"