factory.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455
  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 (
  17. QdrantClient,
  18. )
  19. return QdrantClient()
  20. else:
  21. from open_webui.retrieval.vector.dbs.qdrant import QdrantClient
  22. return QdrantClient()
  23. case VectorType.PINECONE:
  24. from open_webui.retrieval.vector.dbs.pinecone import PineconeClient
  25. return PineconeClient()
  26. case VectorType.OPENSEARCH:
  27. from open_webui.retrieval.vector.dbs.opensearch import OpenSearchClient
  28. return OpenSearchClient()
  29. case VectorType.PGVECTOR:
  30. from open_webui.retrieval.vector.dbs.pgvector import PgvectorClient
  31. return PgvectorClient()
  32. case VectorType.ELASTICSEARCH:
  33. from open_webui.retrieval.vector.dbs.elasticsearch import (
  34. ElasticsearchClient,
  35. )
  36. return ElasticsearchClient()
  37. case VectorType.CHROMA:
  38. from open_webui.retrieval.vector.dbs.chroma import ChromaClient
  39. return ChromaClient()
  40. case _:
  41. raise ValueError(f"Unsupported vector type: {vector_type}")
  42. VECTOR_DB_CLIENT = Vector.get_vector(VECTOR_DB)