fix: embbeding and chunking

.gitignore

@@ -43,3 +43,4 @@ test.py
 /web_ui
 .venv/
 uv.lock
+/test

pkg/entity/persistence/rag.py

@@ -2,6 +2,7 @@ from sqlalchemy import create_engine, Column, String, Text, DateTime, LargeBinar
 from sqlalchemy.orm import declarative_base, sessionmaker
 from datetime import datetime
 import os
+import uuid
 
 Base = declarative_base()
 DATABASE_URL = os.getenv('DATABASE_URL', 'sqlite:///./rag_knowledge.db')
@@ -35,11 +36,11 @@ class File(Base):
     path = Column(String)
     created_at = Column(DateTime, default=datetime.utcnow)
     file_type = Column(String)
-    status = Column(String, default='0')
+    status = Column(Integer, default=0)  # 0: uploaded and processing, 1: completed, 2: failed
 
 class Chunk(Base):
     __tablename__ = 'chunks'
-    id = Column(String, primary_key=True, index=True)
+    id = Column(String, primary_key=True, default=lambda: str(uuid.uuid4()))
     file_id = Column(String, nullable=True)
     text = Column(Text)
 

pkg/rag/knowledge/mgr.py

@@ -7,6 +7,9 @@ from pkg.rag.knowledge.services.parser import FileParser
 from pkg.rag.knowledge.services.chunker import Chunker
 from pkg.rag.knowledge.services.database import create_db_and_tables, SessionLocal, KnowledgeBase, File, Chunk
 from pkg.core import app
+from pkg.rag.knowledge.services.embedder import Embedder
+from pkg.rag.knowledge.services.retriever import Retriever
+from pkg.rag.knowledge.services.chroma_manager import ChromaIndexManager
 
 
 class RAGManager:
@@ -14,11 +17,12 @@ class RAGManager:
 
     def __init__(self, ap: app.Application):
         self.ap = ap
-        self.chroma_manager = None
+        self.chroma_manager = ChromaIndexManager()
         self.parser = FileParser()
         self.chunker = Chunker()
-        self.embedder = None
+        # Initialize Embedder with targeted model type and name
-        self.retriever = None
+        self.embedder = Embedder(model_type='third_party_api', model_name_key='bge-m3', chroma_manager=self.chroma_manager)
+        self.retriever = Retriever(model_type='third_party_api', model_name_key='bge-m3', chroma_manager=self.chroma_manager)
 
     async def initialize_rag_system(self):
         """Initializes the RAG system by creating database tables."""
@@ -140,7 +144,7 @@ class RAGManager:
             return []
 
     async def store_data(
-        self, file_path: str, kb_name: str, file_type: str, kb_description: str = 'Default knowledge base'
+        self, file_path: str, kb_id: str, file_type: str, file_id: str = None
     ):
         """
         Parses, chunks, embeds, and stores data from a given file into the RAG system.
@@ -151,58 +155,35 @@ class RAGManager:
         file_obj = None
 
         try:
-            kb = session.query(KnowledgeBase).filter_by(name=kb_name).first()
+            kb = session.query(KnowledgeBase).filter_by(id=kb_id).first()
             if not kb:
-                kb = KnowledgeBase(name=kb_name, description=kb_description)
+                self.ap.logger.info(f'Knowledge Base "{kb_id}" does not exist. ')
-                session.add(kb)
+                self.ap.logger.info(f'Created Knowledge Base with ID: {kb_id}')
-                session.commit()
-                session.refresh(kb)
-                self.ap.logger.info(f"Knowledge Base '{kb_name}' created during store_data.")
             else:
-                self.ap.logger.info(f"Knowledge Base '{kb_name}' already exists.")
+                self.ap.logger.info(f"Knowledge Base '{kb_id}' already exists.")
 
             file_name = os.path.basename(file_path)
-            existing_file = session.query(File).filter_by(kb_id=kb.id, file_name=file_name).first()
-            if existing_file:
-                self.ap.logger.warning(
-                    f"File '{file_name}' already exists in knowledge base '{kb_name}'. Skipping storage."
-                )
-                return
-
-            file_obj = File(kb_id=kb.id, file_name=file_name, path=file_path, file_type=file_type)
-            session.add(file_obj)
-            session.commit()
-            session.refresh(file_obj)
-            self.ap.logger.info(
-                f"File record '{file_name}' added to database with ID: {file_obj.id}, associated with KB ID: {kb.id}"
-            )
-
             text = await self.parser.parse(file_path)
             if not text:
                 self.ap.logger.warning(
-                    f'No text extracted from file {file_path}. Deleting file record ID: {file_obj.id}.'
+                    f'No text extracted from file {file_path}. '
                 )
-                session.delete(file_obj)
-                session.commit()
                 return
 
             chunks_texts = await self.chunker.chunk(text)
             self.ap.logger.info(f"Chunked file '{file_name}' into {len(chunks_texts)} chunks.")
-            await self.embedder.embed_and_store(file_id=file_obj.id, chunks=chunks_texts)
+            await self.embedder.embed_and_store(file_id=file_id, chunks=chunks_texts)
             self.ap.logger.info(f'Data storage process completed for file: {file_path}')
 
         except Exception as e:
             session.rollback()
             self.ap.logger.error(f'Error in store_data for file {file_path}: {str(e)}', exc_info=True)
-            if file_obj and file_obj.id:
-                try:
-                    await asyncio.to_thread(self.chroma_manager.delete_by_file_id_sync, file_obj.id)
-                except Exception as chroma_e:
-                    self.ap.logger.warning(
-                        f'Could not clean up ChromaDB entries for file_id {file_obj.id} after store_data failure: {chroma_e}'
-                    )
             raise
         finally:
+            if file_id:
+                file_obj = session.query(File).filter_by(id=file_id).first()
+            if file_obj:
+                file_obj.status = 1
             session.close()
 
     async def retrieve_data(self, query: str):
@@ -245,7 +226,6 @@ class RAGManager:
                 self.ap.logger.warning(
                     f'File with ID {file_id} not found in database. Skipping deletion of file record.'
                 )
-
             session.commit()
             self.ap.logger.info(f'Successfully completed data deletion for file_id: {file_id}')
         except Exception as e:
@@ -338,13 +318,13 @@ class RAGManager:
                 self.ap.logger.error(f'Knowledge Base with UUID {knowledge_base_uuid} not found.')
                 return
 
-            # 更新文件的 kb_id
+            if not self.ap.storage_mgr.storage_provider.exists(file_id):
-            file_to_update = session.query(File).filter_by(id=file_id).first()
+                self.ap.logger.error(f'File with ID {file_id} does not exist.')
-            if not file_to_update:
-                self.ap.logger.error(f'File with ID {file_id} not found.')
                 return
-
+            self.ap.logger.info(f'File with ID {file_id} exists, proceeding with association.')
-            file_to_update.kb_id = kb.id
+            # add new file record
+            file_to_update = File(id=file_id, kb_id=kb.id)
+            session.add(file_to_update)
             session.commit()
             self.ap.logger.info(
                 f'Successfully associated file ID {file_id} with knowledge base UUID {knowledge_base_uuid}'
@@ -356,4 +336,20 @@ class RAGManager:
                 exc_info=True,
             )
         finally:
+            # 进行文件解析
+            try:
+                await self.store_data(
+                    file_path = os.path.join('data', 'storage', file_id),
+                    kb_id=knowledge_base_uuid,
+                    file_type=os.path.splitext(file_id)[1].lstrip('.'),
+                    file_id=file_id
+                )
+            except Exception as store_e:
+                # 如果存储数据时出错，更新文件状态为失败
+                file_obj = session.query(File).filter_by(id=file_id).first()
+                if file_obj:
+                    file_obj.status = 2
+                session.commit()
+                self.ap.logger.error(f'Error storing data for file ID {file_id}', exc_info=True)
+
             session.close()

pkg/rag/knowledge/services/chunker.py

@@ -24,33 +24,28 @@ class Chunker(BaseService):
         """
         if not text:
             return []
+        # words = text.split()
+        # chunks = []
+        # current_chunk = []
         
-        # Simple whitespace-based splitting for demonstration
+        # for word in words:
-        # For more advanced chunking, consider libraries like LangChain's text splitters
+        #     current_chunk.append(word)
-        words = text.split()
+        #     if len(current_chunk) > self.chunk_size:
-        chunks = []
+        #         chunks.append(" ".join(current_chunk[:self.chunk_size]))
-        current_chunk = []
+        #         current_chunk = current_chunk[self.chunk_size - self.chunk_overlap:]
         
-        for word in words:
+        # if current_chunk:
-            current_chunk.append(word)
+        #     chunks.append(" ".join(current_chunk))
-            if len(current_chunk) > self.chunk_size:
-                chunks.append(" ".join(current_chunk[:self.chunk_size]))
-                current_chunk = current_chunk[self.chunk_size - self.chunk_overlap:]
-        
-        if current_chunk:
-            chunks.append(" ".join(current_chunk))
         
         # A more robust chunking strategy (e.g., using recursive character text splitter)
-        # from langchain.text_splitter import RecursiveCharacterTextSplitter
+        from langchain.text_splitter import RecursiveCharacterTextSplitter
-        # text_splitter = RecursiveCharacterTextSplitter(
+        text_splitter = RecursiveCharacterTextSplitter(
-        #     chunk_size=self.chunk_size,
+            chunk_size=self.chunk_size,
-        #     chunk_overlap=self.chunk_overlap,
+            chunk_overlap=self.chunk_overlap,
-        #     length_function=len,
+            length_function=len,
-        #     is_separator_regex=False,
+            is_separator_regex=False,
-        # )
+        )
-        # return text_splitter.split_text(text)
+        return text_splitter.split_text(text)
-        
-        return [chunk for chunk in chunks if chunk.strip()] # Filter out empty chunks
 
     async def chunk(self, text: str) -> List[str]:
         """

pkg/rag/knowledge/services/embedder.py

@@ -12,7 +12,7 @@ from pkg.rag.knowledge.services.chroma_manager import ChromaIndexManager # Impor
 logger = logging.getLogger(__name__)
 
 class Embedder(BaseService):
-    def __init__(self, model_type: str, model_name_key: str, chroma_manager: ChromaIndexManager):
+    def __init__(self, model_type: str, model_name_key: str,  chroma_manager: ChromaIndexManager = None):
         super().__init__()
         self.logger = logging.getLogger(self.__class__.__name__)
         self.model_type = model_type