factory.py 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. from open_webui.retrieval.vector.main import VectorDBBase
  2. from open_webui.retrieval.vector.type import VectorType
  3. from open_webui.config import VECTOR_DB, ENABLE_QDRANT_MULTITENANCY_MODE
  4. class Vector:
  5. @staticmethod
  6. def get_vector(vector_type: str) -> VectorDBBase:
  7. """
  8. get vector db instance by vector type
  9. """
  10. match vector_type:
  11. case VectorType.MILVUS:
  12. from open_webui.retrieval.vector.dbs.milvus import MilvusClient
  13. return MilvusClient()
  14. case VectorType.QDRANT:
  15. if ENABLE_QDRANT_MULTITENANCY_MODE:
  16. from open_webui.retrieval.vector.dbs.qdrant_multitenancy import QdrantClient
  17. return QdrantClient()
  18. else:
  19. from open_webui.retrieval.vector.dbs.qdrant import QdrantClient
  20. return QdrantClient()
  21. case VectorType.PINECONE:
  22. from open_webui.retrieval.vector.dbs.pinecone import PineconeClient
  23. return PineconeClient()
  24. case VectorType.OPENSEARCH:
  25. from open_webui.retrieval.vector.dbs.opensearch import OpenSearchClient
  26. return OpenSearchClient()
  27. case VectorType.PGVECTOR:
  28. from open_webui.retrieval.vector.dbs.pgvector import PgvectorClient
  29. return PgvectorClient()
  30. case VectorType.ELASTICSEARCH:
  31. from open_webui.retrieval.vector.dbs.elasticsearch import (
  32. ElasticsearchClient,
  33. )
  34. return ElasticsearchClient()
  35. case VectorType.CHROMA:
  36. from open_webui.retrieval.vector.dbs.chroma import ChromaClient
  37. return ChromaClient()
  38. case _:
  39. raise ValueError(f"Unsupported vector type: {vector_type}")
  40. VECTOR_DB_CLIENT = Vector.get_vector(VECTOR_DB)