refactor: refine dynamic model creation and remove PatchSQLModel (#660)

- using `@lru_cache` to avoid creating duplicate dynamic
chunk/entity/relationship models.
- using `type` function to dynamic create model class, `PatchSQLModel`
and custom `SQLModelMetaclass` is no more need

Author: Mini256

backend/app/api/admin_routes/knowledge_base/document/routes.py

@@ -8,7 +8,7 @@
 from app.api.admin_routes.knowledge_base.models import ChunkItem
 from app.api.deps import SessionDep, CurrentSuperuserDep
 from app.models import Document
-from app.models.chunk import Chunk, KgIndexStatus, get_kb_chunk_model
+from app.models.chunk import KgIndexStatus, get_kb_chunk_model
 from app.models.document import DocIndexTaskStatus
 from app.models.entity import get_kb_entity_model
 from app.models.relationship import get_kb_relationship_model
@@ -198,7 +198,7 @@ def rebuild_kb_document_index_by_ids(
         build_index_for_document.delay(kb.id, doc.id)
 
     # Retry failed kg index tasks.
-    chunks: list[Chunk] = kb_chunk_repo.fetch_by_document_ids(db_session, document_ids)
+    chunks = kb_chunk_repo.fetch_by_document_ids(db_session, document_ids)
     reindex_chunk_ids = []
     ignore_chunk_ids = []
     for chunk in chunks:

backend/app/api/routes/api_key.py

@@ -21,7 +21,7 @@ class CreateApiKeyResponse(BaseModel):
 async def create_api_key(
     session: AsyncSessionDep, user: CurrentSuperuserDep, request: CreateApiKeyRequest
 ) -> CreateApiKeyResponse:
-    _, raw_api_key = await api_key_manager.create_api_key(
+    _, raw_api_key = await api_key_manager.acreate_api_key(
         session, user, request.description
     )
     return CreateApiKeyResponse(api_key=raw_api_key)

backend/app/auth/api_keys.py

@@ -5,7 +5,7 @@
 from typing import Optional, Tuple
 
 from fastapi import Request
-from sqlmodel import select
+from sqlmodel import Session, select
 from sqlmodel.ext.asyncio.session import AsyncSession
 from fastapi_pagination import Params, Page
 from fastapi_pagination.ext.sqlmodel import paginate
@@ -49,7 +49,7 @@ def encrypt_api_key(api_key: str) -> str:
 
 
 class ApiKeyManager:
-    async def create_api_key(
+    async def acreate_api_key(
         self, session: AsyncSession, user: User, description: str
     ) -> Tuple[ApiKey, str]:
         api_key = generate_api_key()
@@ -65,6 +65,22 @@ async def create_api_key(
         await session.refresh(api_key_obj)
         return api_key_obj, api_key
 
+    def create_api_key(
+        self, session: Session, user: User, description: str
+    ) -> Tuple[ApiKey, str]:
+        api_key = generate_api_key()
+        hashed_api_key = encrypt_api_key(api_key)
+        api_key_obj = ApiKey(
+            hashed_secret=hashed_api_key,
+            api_key_display=api_key[:7] + "...." + api_key[-3:],
+            user_id=user.id,
+            description=description,
+        )
+        session.add(api_key_obj)
+        session.commit()
+        session.refresh(api_key_obj)
+        return api_key_obj, api_key
+
     async def get_active_user_by_raw_api_key(
         self, session: AsyncSession, api_key: str
     ) -> Optional[User]:

backend/app/models/__init__.py

@@ -1,10 +1,10 @@
 # flake8: noqa
 from .entity import (
     EntityType,
-    Entity,
     EntityPublic,
+    get_kb_entity_model,
 )
-from .relationship import Relationship, RelationshipPublic
+from .relationship import RelationshipPublic, get_kb_relationship_model
 from .feedback import (
     Feedback,
     FeedbackType,
@@ -18,7 +18,7 @@
 from .chat import Chat, ChatUpdate, ChatVisibility, ChatFilters, ChatOrigin
 from .chat_message import ChatMessage
 from .document import Document, DocIndexTaskStatus
-from .chunk import Chunk, KgIndexStatus
+from .chunk import KgIndexStatus, get_kb_chunk_model
 from .auth import User, UserSession
 from .api_key import ApiKey, PublicApiKey
 from .site_setting import SiteSetting

backend/app/models/chunk.py

@@ -1,9 +1,7 @@
 import enum
-from datetime import datetime
-from typing import Any, Optional, Type
-from uuid import UUID
+from functools import lru_cache
 
-from sqlalchemy import DateTime, func
+from typing import Optional, Type
 from sqlmodel import (
     Field,
     Column,
@@ -15,17 +13,11 @@
 from tidb_vector.sqlalchemy import VectorType
 from llama_index.core.schema import TextNode
 
-from app.core.config import settings
 from app.models.document import Document
 from app.models.knowledge_base import KnowledgeBase
+from app.models.knowledge_base_scoped.table_naming import get_kb_vector_dims
+from app.utils.namespace import format_namespace
 from .base import UpdatableBaseModel, UUIDBaseModel
-from app.models.knowledge_base_scoped.registry import get_kb_scoped_registry
-from .knowledge_base_scoped.table_naming import (
-    get_kb_chunks_table_name,
-    get_kb_vector_dims,
-)
-from app.models.patch.sql_model import SQLModel as PatchSQLModel
-from ..utils.uuid6 import uuid7
 
 
 class KgIndexStatus(str, enum.Enum):
@@ -36,69 +28,26 @@ class KgIndexStatus(str, enum.Enum):
     FAILED = "failed"
 
 
-# Notice: DO NOT forget to modify the definition in `get_kb_chunk_model` to
-# keep the table structure on both sides consistent.
-class Chunk(UUIDBaseModel, UpdatableBaseModel, table=True):
-    hash: str = Field(max_length=64)
-    text: str = Field(sa_column=Column(Text))
-    meta: dict | list = Field(default={}, sa_column=Column(JSON))
-    embedding: Any = Field(
-        sa_column=Column(
-            VectorType(settings.EMBEDDING_DIMS), comment="hnsw(distance=cosine)"
-        )
-    )
-    document_id: int = Field(foreign_key="documents.id", nullable=True)
-    document: "Document" = SQLRelationship(
-        sa_relationship_kwargs={
-            "lazy": "joined",
-            "primaryjoin": "Chunk.document_id == Document.id",
-        },
-    )
-    relations: dict | list = Field(default={}, sa_column=Column(JSON))
-    source_uri: str = Field(max_length=512, nullable=True)
-
-    # TODO: Add vector_index_status, vector_index_result column, vector index should be optional in the future.
-
-    # TODO: Rename to kg_index_status, kg_index_result column.
-    index_status: KgIndexStatus = KgIndexStatus.NOT_STARTED
-    index_result: str = Field(sa_column=Column(Text, nullable=True))
-
-    __tablename__ = "chunks"
-
-    def to_llama_text_node(self) -> TextNode:
-        return TextNode(
-            id_=self.id.hex,
-            text=self.text,
-            embedding=list(self.embedding),
-            metadata=self.meta,
-        )
-
-
 def get_kb_chunk_model(kb: KnowledgeBase) -> Type[SQLModel]:
     vector_dimension = get_kb_vector_dims(kb)
-    chunks_table_name = get_kb_chunks_table_name(kb)
-    ctx = get_kb_scoped_registry(kb)
+    return get_dynamic_chunk_model(vector_dimension, str(kb.id))
 
-    if ctx.chunk_model:
-        return ctx.chunk_model
 
-    class KBChunk(PatchSQLModel, table=True, registry=ctx.registry):
-        __tablename__ = chunks_table_name
-        __table_args__ = {"extend_existing": True}
+@lru_cache(maxsize=None)
+def get_dynamic_chunk_model(
+    vector_dimension: int,
+    namespace: Optional[str] = None,
+) -> Type[SQLModel]:
+    namespace = format_namespace(namespace)
+    chunk_table_name = f"chunks_{namespace}"
+    chunk_model_name = f"Chunk_{namespace}_{vector_dimension}"
 
-        id: UUID = Field(
-            primary_key=True, index=True, nullable=False, default_factory=uuid7
-        )
+    class Chunk(UUIDBaseModel, UpdatableBaseModel):
         hash: str = Field(max_length=64)
         text: str = Field(sa_column=Column(Text))
-        meta: dict | list = Field(default={}, sa_column=Column(JSON))
-        embedding: Any = Field(
-            sa_column=Column(
-                VectorType(vector_dimension), comment="hnsw(distance=cosine)"
-            )
-        )
+        meta: dict = Field(default={}, sa_column=Column(JSON))
+        embedding: list[float] = Field(sa_type=VectorType(vector_dimension))
         document_id: int = Field(foreign_key="documents.id", nullable=True)
-        document: "Document" = SQLRelationship()
         relations: dict | list = Field(default={}, sa_column=Column(JSON))
         source_uri: str = Field(max_length=512, nullable=True)
 
@@ -108,16 +57,6 @@ class KBChunk(PatchSQLModel, table=True, registry=ctx.registry):
         index_status: KgIndexStatus = KgIndexStatus.NOT_STARTED
         index_result: str = Field(sa_column=Column(Text, nullable=True))
 
-        created_at: Optional[datetime] = Field(
-            default=None,
-            sa_column=Column(DateTime(timezone=True), server_default=func.now()),
-        )
-        updated_at: Optional[datetime] = Field(
-            default=None,
-            sa_type=DateTime(timezone=True),
-            sa_column_kwargs={"server_default": func.now(), "onupdate": func.now()},
-        )
-
         def to_llama_text_node(self) -> TextNode:
             return TextNode(
                 id_=self.id.hex,
@@ -126,5 +65,23 @@ def to_llama_text_node(self) -> TextNode:
                 metadata=self.meta,
             )
 
-    ctx.chunk_model = KBChunk
-    return KBChunk
+    chunk_model = type(
+        chunk_model_name,
+        (Chunk,),
+        {
+            "__tablename__": chunk_table_name,
+            "__table_args__": {"extend_existing": True},
+            "__annotations__": {
+                "document": Document,
+            },
+            "document": SQLRelationship(
+                sa_relationship_kwargs={
+                    "lazy": "joined",
+                    "primaryjoin": f"{chunk_model_name}.document_id == Document.id",
+                },
+            ),
+        },
+        table=True,
+    )
+
+    return chunk_model