init

backend/app/api/admin_routes/retrieve/models.py

@@ -0,0 +1,28 @@
+from app.models import Chunk
+from typing import List
+from pydantic import BaseModel
+
+
+class AppConfig(BaseModel):
+    pass
+
+
+class RetrievalConfig(BaseModel):
+    top_k: int = 10
+    similarity_top_k: int = None
+    metadata_filters: dict = {}
+    oversampling_factor: int = 5
+
+
+class RetrieveRequest(BaseModel):
+    query: str
+    retrieval_config: RetrievalConfig = RetrievalConfig()
+
+
+class RetrievedChunk(BaseModel):
+    chunk: Chunk
+    score: float
+
+
+class RetrieveResponse(BaseModel):
+    chunks: List[RetrievedChunk]

backend/app/api/admin_routes/retrieve/routes.py

@@ -0,0 +1,35 @@
+import logging
+from typing import List
+
+from fastapi import APIRouter
+from app.api.admin_routes.models import ChatEngineBasedRetrieveRequest
+from app.api.deps import SessionDep, CurrentSuperuserDep
+from llama_index.core.schema import NodeWithScore
+from app.rag.retrieve import retrieve_service
+
+from app.exceptions import InternalServerError, KBNotFound
+
+router = APIRouter()
+logger = logging.getLogger(__name__)
+
+
+@router.get("/admin/retrieve/chunks")
+def retrieve_chunks(
+    session: SessionDep,
+    user: CurrentSuperuserDep,
+    request: ChatEngineBasedRetrieveRequest,
+) -> List[NodeWithScore]:
+    try:
+        return retrieve_service.chat_engine_retrieve_chunks(
+            session,
+            request.query,
+            top_k=request.top_k,
+            similarity_top_k=request.similarity_top_k,
+            oversampling_factor=request.oversampling_factor,
+            enable_kg_enhance_query_refine=request.enable_kg_enhance_query_refine,
+        )
+    except KBNotFound as e:
+        raise e
+    except Exception as e:
+        logger.exception(e)
+        raise InternalServerError()

backend/app/rag/retrievers/KBVectorSearchRetriever.py

@@ -0,0 +1,57 @@
+import logging
+from llama_index.core.retrievers import BaseRetriever
+from llama_index.core.schema import NodeWithScore, QueryBundle
+from pydantic import BaseModel, List
+
+from app.rag.postprocessors.metadata_post_filter import MetadataFilters
+
+
+logger = logging.getLogger(__name__)
+
+
+class VectorSearchRerankerConfig(BaseModel):
+    enable: bool = True
+
+
+class VectorSearchMetadataFilterConfig(BaseModel):
+    enable: bool = True
+    filters: MetadataFilters = None
+
+
+class VectorSearchRetrieverConfig(BaseModel):
+    enable: bool = True
+    top_k: int = 10
+    similarity_top_k: int = None
+    oversampling_factor: int = 5
+    reranker: VectorSearchRerankerConfig = None
+    metadata_filter: VectorSearchMetadataFilterConfig = None
+
+
+class KnowledgeGraphRetrieverConfig(BaseModel):
+    enable: bool = False
+
+
+class KnowledgeBaseConfig(BaseModel):
+    linked_knowledge_base: LinkedKnowledgeBaseConfig
+
+
+class RetrieverConfig(BaseModel):
+    knowledge_base: KnowledgeBaseConfig
+    vector_search: VectorSearchRetrieverConfig
+    knowledge_graph: KnowledgeGraphRetrieverConfig
+
+
+class VectorSearchRetriever(BaseRetriever):
+    def __init__(self, config: VectorSearchRetrieverConfig):
+        pass
+
+
+class AppRetriever(BaseRetriever):
+    def __init__(
+        self,
+        config: RetrieverConfig,
+    ):
+        pass
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        pass