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)} 个文档")