feat: implement RAG management module with knowledge base and file models, and add processing endpoint

Author: Dallas98

runtime/datamate-python/app/db/models/knowledge_gen.py

@@ -0,0 +1,47 @@
+"""
+Tables of RAG Management Module
+"""
+import uuid
+from sqlalchemy import Column, String, TIMESTAMP, Text, Integer, JSON
+from sqlalchemy.sql import func
+from app.db.session import Base
+
+
+class RagKnowledgeBase(Base):
+    """知识库模型"""
+    __tablename__ = "t_rag_knowledge_base"
+
+    id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
+    name = Column(String(255), nullable=False, comment="知识库名称")
+    type = Column(String(50), nullable=False, comment="知识库类型")
+    description = Column(String(512), nullable=True, comment="知识库描述")
+    embedding_model = Column(String(255), nullable=False, comment="嵌入模型")
+    chat_model = Column(String(255), nullable=True, comment="聊天模型")
+    created_at = Column(TIMESTAMP, server_default=func.current_timestamp(), comment="创建时间")
+    updated_at = Column(TIMESTAMP, server_default=func.current_timestamp(), onupdate=func.current_timestamp(),
+                      comment="更新时间")
+    created_by = Column(String(255), nullable=True, comment="创建者")
+    updated_by = Column(String(255), nullable=True, comment="更新者")
+
+    def __repr__(self):
+        return f"<RagKnowledgeBase(id={self.id}, name={self.name}, type={self.type})>"
+
+
+class RagFile(Base):
+    """知识库文件模型"""
+    __tablename__ = "t_rag_file"
+
+    id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()), comment="UUID")
+    knowledge_base_id = Column(String(36), nullable=False, comment="知识库ID")
+    file_name = Column(String(255), nullable=False, comment="文件名")
+    file_id = Column(String(255), nullable=False, comment="文件ID")
+    chunk_count = Column(Integer, nullable=True, comment="切片数")
+    file_metadata = Column("metadata", JSON, nullable=True, comment="元数据")
+    status = Column(String(50), nullable=True, comment="文件状态")
+    err_msg = Column(Text, nullable=True, comment="错误信息")
+    created_at = Column(TIMESTAMP, server_default=func.current_timestamp(), comment="创建时间")
+    updated_at = Column(TIMESTAMP, server_default=func.current_timestamp(), onupdate=func.current_timestamp(),
+                      comment="更新时间")
+    created_by = Column(String(255), nullable=True, comment="创建者")
+    updated_by = Column(String(255), nullable=True, comment="更新者")
+

runtime/datamate-python/app/module/__init__.py

@@ -6,6 +6,7 @@
 from .generation.interface import router as generation_router
 from .evaluation.interface import router as evaluation_router
 from .collection.interface import router as collection_route
+from .rag.interface.rag_interface import router as rag_router
 
 router = APIRouter(
     prefix="/api"
@@ -17,5 +18,6 @@
 router.include_router(generation_router)
 router.include_router(evaluation_router)
 router.include_router(collection_route)
+router.include_router(rag_router)
 
 __all__ = ["router"]

runtime/datamate-python/app/module/rag/__init__.py

@@ -0,0 +1,24 @@
+from fastapi import APIRouter, Depends, HTTPException
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db.session import get_db
+from app.module.rag.service.rag_service import RAGService
+from app.module.shared.schema import StandardResponse
+
+router = APIRouter(prefix="/rag", tags=["rag"])
+
+@router.post("/process/{knowledge_base_id}")
+async def process_knowledge_base(knowledge_base_id: str, db: AsyncSession = Depends(get_db)):
+    """
+    Process all unprocessed files in a knowledge base.
+    """
+    try:
+        await RAGService(db).init_graph_rag(knowledge_base_id)
+        return StandardResponse(
+            code=200,
+            message="Processing started for knowledge base.",
+            data=None
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+

runtime/datamate-python/app/module/rag/interface/__init__.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+class ProcessRequest(BaseModel):
+    knowledge_base_id: str
+

runtime/datamate-python/app/module/rag/interface/rag_interface.py

@@ -0,0 +1,60 @@
+import asyncio
+import os
+from typing import Awaitable, Callable, Optional
+
+import numpy as np
+from lightrag import LightRAG, QueryParam
+from lightrag.llm.openai import openai_embed, openai_complete_if_cache
+from lightrag.utils import setup_logger, EmbeddingFunc
+
+setup_logger("lightrag", level="DEBUG")
+DEFAULT_WORKING_DIR = "/rag_storage"
+
+
+async def build_llm_model_func(model_name: str, base_url: str, api_key: str) -> Callable[..., Awaitable[str]]:
+    async def _llm_model(
+        prompt, system_prompt=None, history_messages=None, **kwargs
+    ) -> str:
+        history_messages = history_messages or []
+        return await openai_complete_if_cache(
+            model_name,
+            prompt,
+            system_prompt=system_prompt,
+            history_messages=history_messages,
+            api_key=api_key,
+            base_url=base_url,
+            **kwargs,
+        )
+
+    return _llm_model
+
+
+async def build_embedding_func(
+    model_name: str, base_url: str, api_key: str, embedding_dim: int
+) -> EmbeddingFunc:
+    async def _embedding_func(texts: list[str]) -> np.ndarray:
+        return await openai_embed.func(
+            texts,
+            model=model_name,
+            api_key=api_key,
+            base_url=base_url,
+            embedding_dim=embedding_dim,
+        )
+
+    return EmbeddingFunc(embedding_dim=embedding_dim, func=_embedding_func)
+
+
+async def initialize_rag(
+    llm_callable: Callable[..., Awaitable[str]],
+    embedding_callable: EmbeddingFunc,
+    working_dir: Optional[str] = None,
+):
+    target_dir = working_dir or DEFAULT_WORKING_DIR
+    os.makedirs(target_dir, exist_ok=True)
+    rag = LightRAG(
+        working_dir=target_dir,
+        llm_model_func=llm_callable,
+        embedding_func=embedding_callable,
+    )
+    await rag.initialize_storages()
+    return rag

runtime/datamate-python/app/module/rag/schema/rag_schema.py

@@ -0,0 +1,67 @@
+import os
+from typing import Optional
+
+from fastapi import Depends
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from app.db.models.knowledge_gen import RagKnowledgeBase
+from app.db.models.model_config import ModelConfig
+from app.db.session import AsyncSessionLocal
+from .graph_rag import (
+    DEFAULT_WORKING_DIR,
+    build_embedding_func,
+    build_llm_model_func,
+    initialize_rag,
+)
+
+
+class RAGService:
+    def __init__(
+        self,
+        db: AsyncSession = Depends(AsyncSessionLocal),
+    ):
+        self.db = db
+        self.rag = None
+
+
+    async def get_unprocessed_files(self, knowledge_base_id: str) -> list[str]:
+        pass
+
+    async def init_graph_rag(self, knowledge_base_id: str):
+        kb = await self._get_knowledge_base(knowledge_base_id)
+        embedding_model = await self._get_model_config(kb.embedding_model)
+        chat_model = await self._get_model_config(kb.chat_model)
+
+        llm_callable = await build_llm_model_func(
+            chat_model.model_name, chat_model.base_url, chat_model.api_key
+        )
+        embedding_callable = await build_embedding_func(
+            embedding_model.model_name,
+            embedding_model.base_url,
+            embedding_model.api_key,
+            embedding_dim=embedding_model.embedding_dim if hasattr(embedding_model, "embedding_dim") else 1024,
+        )
+
+        kb_working_dir = os.path.join(DEFAULT_WORKING_DIR, kb.name)
+        self.rag = await initialize_rag(llm_callable, embedding_callable, kb_working_dir)
+        return {"status": "initialized", "knowledge_base_id": knowledge_base_id}
+
+    async def _get_knowledge_base(self, knowledge_base_id: str):
+        result = await self.db.execute(
+            select(RagKnowledgeBase).where(RagKnowledgeBase.id == knowledge_base_id)
+        )
+        knowledge_base = result.scalars().first()
+        if not knowledge_base:
+            raise ValueError(f"Knowledge base with ID {knowledge_base_id} not found.")
+        return knowledge_base
+
+    async def _get_model_config(self, model_id: Optional[str]):
+        if not model_id:
+            raise ValueError("Model ID is required for initializing RAG.")
+        result = await self.db.execute(select(ModelConfig).where(ModelConfig.id == model_id))
+        model = result.scalars().first()
+        if not model:
+            raise ValueError(f"Model config with ID {model_id} not found.")
+        return model
+