feat(rag): make embedding and retrieving available

pkg/rag/knowledge/services/retriever.py

@@ -1,99 +1,46 @@
 from __future__ import annotations
-import logging
-import numpy as np  # Make sure numpy is imported
-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.vector.vdb import VectorDatabase
+
+from . import base_service
 from ....core import app
-
-logger = logging.getLogger(__name__)
+from ....provider.modelmgr.requester import RuntimeEmbeddingModel
+from ....entity.rag import retriever as retriever_entities
 
 
-class Retriever(BaseService):
+class Retriever(base_service.BaseService):
     def __init__(self, ap: app.Application):
         super().__init__()
-        self.logger = logging.getLogger(self.__class__.__name__)
         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:
-            raise RuntimeError('Retriever embedding model not loaded. Please check Retriever initialization.')
+    async def retrieve(
+        self, kb_id: str, query: str, embedding_model: RuntimeEmbeddingModel, k: int = 5
+    ) -> list[retriever_entities.RetrieveResultEntry]:
+        self.ap.logger.info(f"Retrieving for query: '{query}' with k={k} using {embedding_model.model_entity.uuid}")
 
-        self.logger.info(f"Retrieving for query: '{query}' with k={k} using {self.model_name_key}")
+        query_embedding: list[float] = await embedding_model.requester.invoke_embedding(
+            model=embedding_model,
+            input_text=[query],
+            extra_args={},  # TODO: add extra args
+        )
 
-        query_embedding: List[float] = await self.embedding_model.embed_query(query)
-        query_embedding_np = np.array([query_embedding], dtype=np.float32)
-
-        # 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)
+        chroma_results = await self._run_sync(self.ap.vector_db_mgr.vector_db.search, kb_id, query_embedding[0], k)
 
         # 'ids' is always returned by ChromaDB, even if not explicitly in 'include'
         matched_chroma_ids = chroma_results.get('ids', [[]])[0]
         distances = chroma_results.get('distances', [[]])[0]
         chroma_metadatas = chroma_results.get('metadatas', [[]])[0]
-        chroma_documents = chroma_results.get('documents', [[]])[0]
 
         if not matched_chroma_ids:
-            self.logger.info('No relevant chunks found in Chroma.')
+            self.ap.logger.info('No relevant chunks found in Chroma.')
             return []
 
-        db_chunk_ids = []
-        for metadata in chroma_metadatas:
-            if 'chunk_id' in metadata:
-                db_chunk_ids.append(metadata['chunk_id'])
-            else:
-                self.logger.warning(f"Metadata missing 'chunk_id': {metadata}. Skipping this entry.")
+        result: list[retriever_entities.RetrieveResultEntry] = []
 
-        if not db_chunk_ids:
-            self.logger.warning('No valid chunk_ids extracted from Chroma results metadata.')
-            return []
-
-        self.logger.info(f'Fetching {len(db_chunk_ids)} chunk details from relational database...')
-        chunks_from_db = await self._run_sync(
-            lambda cids: self._db_get_chunks_sync(
-                SessionLocal(), cids
-            ),  # Ensure SessionLocal is passed correctly for _db_get_chunks_sync
-            db_chunk_ids,
-        )
-
-        chunk_map = {chunk.id: chunk for chunk in chunks_from_db}
-        results_list: List[Dict[str, Any]] = []
-
-        for i, chroma_id in enumerate(matched_chroma_ids):
-            try:
-                # Ensure original_chunk_id is int for DB lookup
-                original_chunk_id = int(chroma_id.split('_')[-1])
-            except (ValueError, IndexError):
-                self.logger.warning(f'Could not parse chunk_id from Chroma ID: {chroma_id}. Skipping.')
-                continue
-
-            chunk_text_from_chroma = chroma_documents[i]
-            distance = float(distances[i])
-            file_id_from_chroma = chroma_metadatas[i].get('file_id')
-
-            chunk_from_db = chunk_map.get(original_chunk_id)
-
-            results_list.append(
-                {
-                    'chunk_id': original_chunk_id,
-                    'text': chunk_from_db.text if chunk_from_db else chunk_text_from_chroma,
-                    'distance': distance,
-                    'file_id': file_id_from_chroma,
-                }
+        for i, id in enumerate(matched_chroma_ids):
+            entry = retriever_entities.RetrieveResultEntry(
+                id=id,
+                metadata=chroma_metadatas[i],
+                distance=distances[i],
             )
+            result.append(entry)
 
-        self.logger.info(f'Retrieved {len(results_list)} chunks.')
-        return results_list
-
-    def _db_get_chunks_sync(self, session: Session, chunk_ids: List[int]) -> List[Chunk]:
-        self.logger.debug(f'Fetching {len(chunk_ids)} chunk details from database (sync).')
-        chunks = session.query(Chunk).filter(Chunk.id.in_(chunk_ids)).all()
-        session.close()
-        return chunks
+        return result