dify/api/tests/integration_tests/vdb/__mock/tcvectordb.py

import os
from typing import Optional

import pytest
from _pytest.monkeypatch import MonkeyPatch
from requests.adapters import HTTPAdapter
from tcvectordb import VectorDBClient  # type: ignore
from tcvectordb.model.database import Collection, Database  # type: ignore
from tcvectordb.model.document import Document, Filter  # type: ignore
from tcvectordb.model.enum import ReadConsistency  # type: ignore
from tcvectordb.model.index import Index  # type: ignore
from xinference_client.types import Embedding  # type: ignore


class MockTcvectordbClass:
    def mock_vector_db_client(
        self,
        url=None,
        username="",
        key="",
        read_consistency: ReadConsistency = ReadConsistency.EVENTUAL_CONSISTENCY,
        timeout=5,
        adapter: HTTPAdapter = None,
    ):
        self._conn = None
        self._read_consistency = read_consistency

    def list_databases(self) -> list[Database]:
        return [
            Database(
                conn=self._conn,
                read_consistency=self._read_consistency,
                name="dify",
            )
        ]

    def list_collections(self, timeout: Optional[float] = None) -> list[Collection]:
        return []

    def drop_collection(self, name: str, timeout: Optional[float] = None):
        return {"code": 0, "msg": "operation success"}

    def create_collection(
        self,
        name: str,
        shard: int,
        replicas: int,
        description: str,
        index: Index,
        embedding: Embedding = None,
        timeout: Optional[float] = None,
    ) -> Collection:
        return Collection(
            self,
            name,
            shard,
            replicas,
            description,
            index,
            embedding=embedding,
            read_consistency=self._read_consistency,
            timeout=timeout,
        )

    def describe_collection(self, name: str, timeout: Optional[float] = None) -> Collection:
        collection = Collection(self, name, shard=1, replicas=2, description=name, timeout=timeout)
        return collection

    def collection_upsert(
        self, documents: list[Document], timeout: Optional[float] = None, build_index: bool = True, **kwargs
    ):
        return {"code": 0, "msg": "operation success"}

    def collection_search(
        self,
        vectors: list[list[float]],
        filter: Filter = None,
        params=None,
        retrieve_vector: bool = False,
        limit: int = 10,
        output_fields: Optional[list[str]] = None,
        timeout: Optional[float] = None,
    ) -> list[list[dict]]:
        return [[{"metadata": '{"doc_id":"foo1"}', "text": "text", "doc_id": "foo1", "score": 0.1}]]

    def collection_query(
        self,
        document_ids: Optional[list] = None,
        retrieve_vector: bool = False,
        limit: Optional[int] = None,
        offset: Optional[int] = None,
        filter: Optional[Filter] = None,
        output_fields: Optional[list[str]] = None,
        timeout: Optional[float] = None,
    ) -> list[dict]:
        return [{"metadata": '{"doc_id":"foo1"}', "text": "text", "doc_id": "foo1", "score": 0.1}]

    def collection_delete(
        self,
        document_ids: Optional[list[str]] = None,
        filter: Filter = None,
        timeout: Optional[float] = None,
    ):
        return {"code": 0, "msg": "operation success"}


MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"


@pytest.fixture
def setup_tcvectordb_mock(request, monkeypatch: MonkeyPatch):
    if MOCK:
        monkeypatch.setattr(VectorDBClient, "__init__", MockTcvectordbClass.mock_vector_db_client)
        monkeypatch.setattr(VectorDBClient, "list_databases", MockTcvectordbClass.list_databases)
        monkeypatch.setattr(Database, "collection", MockTcvectordbClass.describe_collection)
        monkeypatch.setattr(Database, "list_collections", MockTcvectordbClass.list_collections)
        monkeypatch.setattr(Database, "drop_collection", MockTcvectordbClass.drop_collection)
        monkeypatch.setattr(Database, "create_collection", MockTcvectordbClass.create_collection)
        monkeypatch.setattr(Collection, "upsert", MockTcvectordbClass.collection_upsert)
        monkeypatch.setattr(Collection, "search", MockTcvectordbClass.collection_search)
        monkeypatch.setattr(Collection, "query", MockTcvectordbClass.collection_query)
        monkeypatch.setattr(Collection, "delete", MockTcvectordbClass.collection_delete)

    yield

    if MOCK:
        monkeypatch.undo()