# retrieval.py
# --- 确认 Import 路径已更新 ---
from milvus_haystack import MilvusDocumentStore  # 用于类型提示
from milvus_haystack.milvus_embedding_retriever import (
    MilvusEmbeddingRetriever,
)  # 使用正确的 integration import

# 从配置导入 top_k
from config import RETRIEVER_TOP_K


def initialize_vector_retriever(
    document_store: MilvusDocumentStore,
) -> MilvusEmbeddingRetriever:
    """
    Initializes the MilvusEmbeddingRetriever using milvus-haystack package.
    Requires a correctly initialized MilvusDocumentStore instance.
    """
    print(f"Initializing Milvus Embedding Retriever with top_k={RETRIEVER_TOP_K}")

    # 初始化 MilvusEmbeddingRetriever 实例
    # 它需要 document_store 实例来进行实际的搜索操作
    # top_k 参数控制返回文档的数量
    retriever = MilvusEmbeddingRetriever(
        document_store=document_store,
        top_k=RETRIEVER_TOP_K,
        # 其他可选参数可以根据需要添加，例如 filters_policy
    )
    print("Milvus Embedding Retriever initialized.")
    return retriever