chore: stash

pkg/rag/knowledge/mgr.py

@@ -13,10 +13,9 @@ from pkg.rag.knowledge.services.database import (
 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
-from pkg.core import taskmgr
-from ...entity.persistence import rag as persistence_rag
 import sqlalchemy
+from ...entity.persistence import rag as persistence_rag
+from pkg.core import taskmgr
 
 
 class RuntimeKnowledgeBase:
@@ -24,8 +23,6 @@ class RuntimeKnowledgeBase:
 
     knowledge_base_entity: persistence_rag.KnowledgeBase
 
-    chroma_manager: ChromaIndexManager
-
     parser: FileParser
 
     chunker: Chunker
@@ -37,11 +34,12 @@ class RuntimeKnowledgeBase:
     def __init__(self, ap: app.Application, knowledge_base_entity: persistence_rag.KnowledgeBase):
         self.ap = ap
         self.knowledge_base_entity = knowledge_base_entity
-        self.chroma_manager = ChromaIndexManager(ap=self.ap)
         self.parser = FileParser(ap=self.ap)
         self.chunker = Chunker(ap=self.ap)
-        self.embedder = Embedder(ap=self.ap, chroma_manager=self.chroma_manager)
-        self.retriever = Retriever(ap=self.ap, chroma_manager=self.chroma_manager)
+        self.embedder = Embedder(ap=self.ap)
+        self.retriever = Retriever(ap=self.ap)
+        # 传递kb_id给retriever
+        self.retriever.kb_id = knowledge_base_entity.uuid
 
     async def initialize(self):
         pass

pkg/rag/knowledge/services/chroma_manager.py

@@ -1,67 +0,0 @@
-import numpy as np
-import logging
-from chromadb import PersistentClient
-from pkg.core import app
-
-logger = logging.getLogger(__name__)
-
-
-class ChromaIndexManager:
-    def __init__(self, ap: app.Application, collection_name: str = 'default_collection'):
-        self.ap = ap
-        chroma_data_path = './data/chroma'
-        self.client = PersistentClient(path=chroma_data_path)
-        self._collection_name = collection_name
-        self._collection = None
-
-        self.ap.logger.info(f'ChromaIndexManager initialized. Collection name: {self._collection_name}')
-
-    @property
-    def collection(self):
-        if self._collection is None:
-            self._collection = self.client.get_or_create_collection(name=self._collection_name)
-            self.ap.logger.info(f"Chroma collection '{self._collection_name}' accessed/created.")
-        return self._collection
-
-    def add_embeddings_sync(
-        self, file_ids: list[int], chunk_ids: list[int], embeddings: np.ndarray, documents: list[str]
-    ):
-        if (
-            embeddings.shape[0] != len(chunk_ids)
-            or embeddings.shape[0] != len(file_ids)
-            or embeddings.shape[0] != len(documents)
-        ):
-            raise ValueError('Embedding, file_id, chunk_id, and document count mismatch.')
-
-        chroma_ids = [f'{file_id}_{chunk_id}' for file_id, chunk_id in zip(file_ids, chunk_ids)]
-        metadatas = [{'file_id': fid, 'chunk_id': cid} for fid, cid in zip(file_ids, chunk_ids)]
-
-        self.logger.debug(f"Adding {len(embeddings)} embeddings to Chroma collection '{self._collection_name}'.")
-        self.collection.add(embeddings=embeddings.tolist(), ids=chroma_ids, metadatas=metadatas, documents=documents)
-        self.logger.info(f"Added {len(embeddings)} embeddings to Chroma collection '{self._collection_name}'.")
-
-    def search_sync(self, query_embedding: np.ndarray, k: int = 5):
-        """
-        Searches the Chroma collection for the top-k nearest neighbors.
-        Args:
-            query_embedding: A numpy array of the query embedding.
-            k: The number of results to return.
-        Returns:
-            A dictionary containing query results from Chroma.
-        """
-        self.logger.debug(f"Searching Chroma collection '{self._collection_name}' with k={k}.")
-        results = self.collection.query(
-            query_embeddings=query_embedding.tolist(),
-            n_results=k,
-            # REMOVE 'ids' from the include list. It's returned by default.
-            include=['metadatas', 'distances', 'documents'],
-        )
-        self.logger.debug(f'Chroma search returned {len(results.get("ids", [[]])[0])} results.')
-        return results
-
-    def delete_by_file_id_sync(self, file_id: int):
-        self.logger.info(
-            f"Deleting embeddings for file_id: {file_id} from Chroma collection '{self._collection_name}'."
-        )
-        self.collection.delete(where={'file_id': file_id})
-        self.logger.info(f'Deleted embeddings for file_id: {file_id} from Chroma.')

pkg/rag/knowledge/services/embedder.py

@@ -1,21 +1,17 @@
 from __future__ import annotations
 import asyncio
-import logging
 import numpy as np
 from typing import List
 from sqlalchemy.orm import Session
 from pkg.rag.knowledge.services.base_service import BaseService
 from pkg.rag.knowledge.services.database import Chunk, SessionLocal
-from pkg.rag.knowledge.services.chroma_manager import ChromaIndexManager
 from ....core import app
 from ....provider.modelmgr.requester import RuntimeEmbeddingModel
 
 
 class Embedder(BaseService):
-    def __init__(self, ap: app.Application, chroma_manager: ChromaIndexManager = None) -> None:
+    def __init__(self, ap: app.Application) -> None:
         super().__init__()
-        self.logger = logging.getLogger(self.__class__.__name__)
-        self.chroma_manager = chroma_manager
         self.ap = ap
 
     def _db_save_chunks_sync(self, session: Session, file_id: int, chunks_texts: List[str]):
@@ -24,22 +20,19 @@ class Embedder(BaseService):
         This function assumes it's called within a context where the session
         will be committed/rolled back and closed by the caller.
         """
-        self.logger.debug(f'Saving {len(chunks_texts)} chunks for file_id {file_id} to DB (sync).')
+        self.ap.logger.debug(f'Saving {len(chunks_texts)} chunks for file_id {file_id} to DB (sync).')
         chunk_objects = []
         for text in chunks_texts:
             chunk = Chunk(file_id=file_id, text=text)
             session.add(chunk)
             chunk_objects.append(chunk)
         session.flush()  # This populates the .id attribute for each new chunk object
-        self.logger.debug(f'Successfully added {len(chunk_objects)} chunk entries to DB.')
+        self.ap.logger.debug(f'Successfully added {len(chunk_objects)} chunk entries to DB.')
         return chunk_objects
 
     async def embed_and_store(
         self, file_id: int, chunks: List[str], embedding_model: RuntimeEmbeddingModel
     ) -> List[Chunk]:
-        if not embedding_model:
-            raise RuntimeError('Embedding model not loaded. Please check Embedder initialization.')
-
         session = SessionLocal()  # Start a session that will live for the whole operation
         chunk_objects = []
         try:
@@ -50,7 +43,7 @@ class Embedder(BaseService):
             session.commit()  # Commit chunks to make their IDs permanent and accessible
 
             if not chunk_objects:
-                self.logger.warning(
+                self.ap.logger.warning(
                     f'No chunk objects created for file_id {file_id}. Skipping embedding and Chroma storage.'
                 )
                 return []
@@ -67,23 +60,28 @@ class Embedder(BaseService):
 
             embeddings_np = np.array(embeddings, dtype=np.float32)
 
-            self.logger.info('Saving embeddings to Chroma...')
             chunk_ids = [c.id for c in chunk_objects]
-            file_ids_for_chroma = [file_id] * len(chunk_ids)
-
-            await self._run_sync(  # Use _run_sync for the Chroma operation, as it's a sync call
-                self.chroma_manager.add_embeddings_sync,
-                file_ids_for_chroma,
-                chunk_ids,
+            # collection名用kb_id（file对象有kb_id字段）
+            kb_id = session.query(Chunk).filter_by(id=chunk_ids[0]).first().file.kb_id if chunk_ids else None
+            if not kb_id:
+                self.ap.logger.warning('无法获取kb_id，向量存储失败')
+                return chunk_objects
+            chroma_ids = [f'{file_id}_{cid}' for cid in chunk_ids]
+            metadatas = [{'file_id': file_id, 'chunk_id': cid} for cid in chunk_ids]
+            await self._run_sync(
+                self.ap.vector_db_mgr.vector_db.add_embeddings,
+                kb_id,
+                chroma_ids,
                 embeddings_np,
-                chunks,  # Pass original chunks texts for documents
+                metadatas,
+                chunks,
             )
-            self.logger.info(f'Successfully saved {len(chunk_objects)} embeddings to Chroma.')
+            self.ap.logger.info(f'Successfully saved {len(chunk_objects)} embeddings to VectorDB.')
             return chunk_objects
 
         except Exception as e:
             session.rollback()  # Rollback on any error
-            self.logger.error(f'Failed to process and store data for file_id {file_id}: {e}', exc_info=True)
+            self.ap.logger.error(f'Failed to process and store data for file_id {file_id}: {e}', exc_info=True)
             raise  # Re-raise the exception to propagate it
         finally:
             session.close()  # Ensure the session is always closed

pkg/rag/knowledge/services/retriever.py

@@ -5,18 +5,18 @@ from typing import List, Dict, Any
 from sqlalchemy.orm import Session
 from pkg.rag.knowledge.services.base_service import BaseService
 from pkg.rag.knowledge.services.database import Chunk, SessionLocal
-from pkg.rag.knowledge.services.chroma_manager import ChromaIndexManager
+from pkg.vector.vdb import VectorDatabase
 from ....core import app
 
 logger = logging.getLogger(__name__)
 
 
 class Retriever(BaseService):
-    def __init__(self, ap:app.Application, chroma_manager: ChromaIndexManager):
+    def __init__(self, ap: app.Application):
         super().__init__()
         self.logger = logging.getLogger(self.__class__.__name__)
-        self.chroma_manager = chroma_manager
         self.ap = ap
+        self.vector_db: VectorDatabase = ap.vector_db_mgr.vector_db
 
     async def retrieve(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
         if not self.embedding_model:
@@ -27,7 +27,12 @@ class Retriever(BaseService):
         query_embedding: List[float] = await self.embedding_model.embed_query(query)
         query_embedding_np = np.array([query_embedding], dtype=np.float32)
 
-        chroma_results = await self._run_sync(self.chroma_manager.search_sync, query_embedding_np, k)
+        # collection名用kb_id（假设retriever有kb_id属性或通过ap传递）
+        kb_id = getattr(self, 'kb_id', None)
+        if not kb_id:
+            self.logger.warning('无法获取kb_id，向量检索失败')
+            return []
+        chroma_results = await self._run_sync(self.vector_db.search, kb_id, query_embedding_np, k)
 
         # 'ids' is always returned by ChromaDB, even if not explicitly in 'include'
         matched_chroma_ids = chroma_results.get('ids', [[]])[0]