feat: 调整结构

This commit is contained in:
李如威 2025-08-07 15:56:14 +08:00
parent 148474de9b
commit 5b101f26c6
1 changed files with 4 additions and 3 deletions

View File

@ -173,13 +173,14 @@ class ModelManager:
class BaseRAG(ABC): class BaseRAG(ABC):
def __init__( def __init__(
self, self,
vector_store_name: str = "default", vector_store_name: str = "default",
embedding_config: Optional[Dict] = None,
retriever_top_k: int = 3, retriever_top_k: int = 3,
llm: Optional[BaseLLM] = None,
persist_directory: str = "./chroma_db", persist_directory: str = "./chroma_db",
llm: Optional[BaseLLM] = None,
embedding_config: Optional[Dict] = None,
rerank_config: Optional[Dict] = None, rerank_config: Optional[Dict] = None,
): ):
""" """
@ -271,7 +272,7 @@ class BaseRAG(ABC):
# 根据分数排序 # 根据分数排序
doc_scores = list(zip(documents, scores)) doc_scores = list(zip(documents, scores))
doc_scores.sort(key=lambda x: x[1], reverse=True) doc_scores.sort(key=lambda x: x[1], reverse=True)
# 返回top_k个文档 # 返回top_k个文档
result = [doc for doc, score in doc_scores[:top_k]] result = [doc for doc, score in doc_scores[:top_k]]
print(f"重排输出: {len(result)} 个文档") print(f"重排输出: {len(result)} 个文档")