Agentic_Rag/tests/test_milvus_store.py

from app.embeddings import HashEmbeddingProvider
from app.milvus_store import IndexedDocument, MilvusVectorStore


class FakeSchema:
    def __init__(self) -> None:
        self.fields = []

    def add_field(self, *args, **kwargs) -> None:
        self.fields.append((args, kwargs))


class FakeIndexParams:
    def __init__(self) -> None:
        self.indexes = []

    def add_index(self, **kwargs) -> None:
        self.indexes.append(kwargs)


class FakeMilvusClient:
    last_instance = None

    def __init__(self, **kwargs) -> None:
        self.kwargs = kwargs
        self.created = None
        self.inserted = []
        self.search_args = None
        self.index_params = FakeIndexParams()
        FakeMilvusClient.last_instance = self

    @classmethod
    def create_schema(cls, **kwargs):
        return FakeSchema()

    def list_collections(self):
        return []

    def has_collection(self, **kwargs):
        return False

    def prepare_index_params(self):
        return self.index_params

    def create_collection(self, **kwargs):
        self.created = kwargs

    def load_collection(self, **kwargs):
        return None

    def insert(self, **kwargs):
        self.inserted = kwargs["data"]

    def flush(self, **kwargs):
        return None

    def search(self, **kwargs):
        self.search_args = kwargs
        return [
            [
                {
                    "id": 1,
                    "distance": 0.88,
                    "entity": {
                        "content": "耳机拆封且影响卫生安全时不适用无理由退货。",
                        "source": "return_policy.md",
                        "doc_type": "commerce_policy",
                        "chunk_index": 1,
                        "policy_type": "return_policy",
                        "category": "all",
                        "metadata": {"title": "商品退货政策"},
                    },
                }
            ]
        ]


def test_milvus_schema_insert_and_filtered_search(monkeypatch) -> None:
    monkeypatch.setattr("app.milvus_store.MilvusClient", FakeMilvusClient)
    embeddings = HashEmbeddingProvider(128)
    store = MilvusVectorStore(
        uri="http://localhost:19530",
        token="",
        collection_name="test_docs",
        embeddings=embeddings,
    )
    store.create_collection(recreate=True)

    client = FakeMilvusClient.last_instance
    assert client.created["collection_name"] == "test_docs"
    assert {item["index_type"] for item in client.index_params.indexes} == {
        "HNSW",
        "INVERTED",
    }

    inserted = store.insert_documents(
        [
            IndexedDocument(
                content="耳机拆封后不适用无理由退货",
                source="return_policy.md",
                doc_type="commerce_policy",
                chunk_index=1,
                policy_type="return_policy",
                category="all",
                metadata={},
            )
        ]
    )
    assert inserted == 1
    assert len(client.inserted[0]["embedding"]) == 128

    evidence = store.search("耳机拆封后能退货吗", top_k=4, policy_type="return_policy")
    assert client.search_args["filter"] == 'policy_type == "return_policy"'
    assert client.search_args["limit"] == 4
    assert evidence[0].source_type == "milvus"
    assert evidence[0].score == 0.88