From 5b101f26c671d3cfe7d239e3b7fa5f1655b653d6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Thu, 7 Aug 2025 15:56:14 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E8=B0=83=E6=95=B4=E7=BB=93=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/base_rag/core.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/base_rag/core.py b/src/base_rag/core.py index 7256bed..f839ff9 100644 --- a/src/base_rag/core.py +++ b/src/base_rag/core.py @@ -173,13 +173,14 @@ class ModelManager: class BaseRAG(ABC): + def __init__( self, vector_store_name: str = "default", - embedding_config: Optional[Dict] = None, retriever_top_k: int = 3, - llm: Optional[BaseLLM] = None, persist_directory: str = "./chroma_db", + llm: Optional[BaseLLM] = None, + embedding_config: Optional[Dict] = None, rerank_config: Optional[Dict] = None, ): """ @@ -271,7 +272,7 @@ class BaseRAG(ABC): # 根据分数排序 doc_scores = list(zip(documents, scores)) doc_scores.sort(key=lambda x: x[1], reverse=True) - + # 返回top_k个文档 result = [doc for doc, score in doc_scores[:top_k]] print(f"重排输出: {len(result)} 个文档")