factory.py 2.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. from open_webui.retrieval.vector.main import VectorDBBase
  2. from open_webui.retrieval.vector.type import VectorType
  3. from open_webui.config import (
  4. VECTOR_DB,
  5. ENABLE_QDRANT_MULTITENANCY_MODE,
  6. ENABLE_MILVUS_MULTITENANCY_MODE,
  7. )
  8. class Vector:
  9. @staticmethod
  10. def get_vector(vector_type: str) -> VectorDBBase:
  11. """
  12. get vector db instance by vector type
  13. """
  14. match vector_type:
  15. case VectorType.MILVUS:
  16. if ENABLE_MILVUS_MULTITENANCY_MODE:
  17. from open_webui.retrieval.vector.dbs.milvus_multitenancy import (
  18. MilvusClient,
  19. )
  20. return MilvusClient()
  21. else:
  22. from open_webui.retrieval.vector.dbs.milvus import MilvusClient
  23. return MilvusClient()
  24. case VectorType.QDRANT:
  25. if ENABLE_QDRANT_MULTITENANCY_MODE:
  26. from open_webui.retrieval.vector.dbs.qdrant_multitenancy import (
  27. QdrantClient,
  28. )
  29. return QdrantClient()
  30. else:
  31. from open_webui.retrieval.vector.dbs.qdrant import QdrantClient
  32. return QdrantClient()
  33. case VectorType.PINECONE:
  34. from open_webui.retrieval.vector.dbs.pinecone import PineconeClient
  35. return PineconeClient()
  36. case VectorType.S3VECTOR:
  37. from open_webui.retrieval.vector.dbs.s3vector import S3VectorClient
  38. return S3VectorClient()
  39. case VectorType.OPENSEARCH:
  40. from open_webui.retrieval.vector.dbs.opensearch import OpenSearchClient
  41. return OpenSearchClient()
  42. case VectorType.PGVECTOR:
  43. from open_webui.retrieval.vector.dbs.pgvector import PgvectorClient
  44. return PgvectorClient()
  45. case VectorType.ELASTICSEARCH:
  46. from open_webui.retrieval.vector.dbs.elasticsearch import (
  47. ElasticsearchClient,
  48. )
  49. return ElasticsearchClient()
  50. case VectorType.CHROMA:
  51. from open_webui.retrieval.vector.dbs.chroma import ChromaClient
  52. return ChromaClient()
  53. case VectorType.ORACLE23AI:
  54. from open_webui.retrieval.vector.dbs.oracle23ai import Oracle23aiClient
  55. return Oracle23aiClient()
  56. case _:
  57. raise ValueError(f"Unsupported vector type: {vector_type}")
  58. VECTOR_DB_CLIENT = Vector.get_vector(VECTOR_DB)