Coverage for open_webui/retrieval/vector/factory.py: 35%

1from threading import Lock 

2 

3from fastapi import HTTPException 

4from open_webui.config import ( 

5 ENABLE_MILVUS_MULTITENANCY_MODE, 

6 ENABLE_QDRANT_MULTITENANCY_MODE, 

7 VECTOR_DB, 

8) 

9from open_webui.env import USE_SLIM 

10from open_webui.retrieval.vector.main import VectorDBBase 

11from open_webui.retrieval.vector.type import VectorType 

12 

13 

14class Vector: 

15 @staticmethod 

16 def get_vector(vector_type: str) -> VectorDBBase: 

17 """ 

18 get vector db instance by vector type 

19 """ 

20 if USE_SLIM and vector_type != VectorType.PGVECTOR: 20 ↛ 21line 20 didn't jump to line 21 because the condition on line 20 was never true

21 raise HTTPException( 

22 503, 

23 'Slim requires PostgreSQL/pgvector for vector storage. Set VECTOR_DB=pgvector and PGVECTOR_DB_URL, or use the standard image.', 

24 ) 

25 match vector_type: 

26 case VectorType.MILVUS: 26 ↛ 27line 26 didn't jump to line 27 because the pattern on line 26 never matched

27 if ENABLE_MILVUS_MULTITENANCY_MODE: 

28 from open_webui.retrieval.vector.dbs.milvus_multitenancy import ( 

29 MilvusClient, 

30 ) 

31 

32 return MilvusClient() 

33 else: 

34 from open_webui.retrieval.vector.dbs.milvus import MilvusClient 

35 

36 return MilvusClient() 

37 case VectorType.QDRANT: 37 ↛ 38line 37 didn't jump to line 38 because the pattern on line 37 never matched

38 if ENABLE_QDRANT_MULTITENANCY_MODE: 

39 from open_webui.retrieval.vector.dbs.qdrant_multitenancy import ( 

40 QdrantClient, 

41 ) 

42 

43 return QdrantClient() 

44 else: 

45 from open_webui.retrieval.vector.dbs.qdrant import QdrantClient 

46 

47 return QdrantClient() 

48 case VectorType.PINECONE: 48 ↛ 49line 48 didn't jump to line 49 because the pattern on line 48 never matched

49 from open_webui.retrieval.vector.dbs.pinecone import PineconeClient 

50 

51 return PineconeClient() 

52 case VectorType.S3VECTOR: 52 ↛ 53line 52 didn't jump to line 53 because the pattern on line 52 never matched

53 from open_webui.retrieval.vector.dbs.s3vector import S3VectorClient 

54 

55 return S3VectorClient() 

56 case VectorType.OPENSEARCH: 56 ↛ 57line 56 didn't jump to line 57 because the pattern on line 56 never matched

57 from open_webui.retrieval.vector.dbs.opensearch import OpenSearchClient 

58 

59 return OpenSearchClient() 

60 case VectorType.PGVECTOR: 60 ↛ 61line 60 didn't jump to line 61 because the pattern on line 60 never matched

61 from open_webui.retrieval.vector.dbs.pgvector import PgvectorClient 

62 

63 return PgvectorClient() 

64 case VectorType.OPENGAUSS: 64 ↛ 65line 64 didn't jump to line 65 because the pattern on line 64 never matched

65 from open_webui.retrieval.vector.dbs.opengauss import OpenGaussClient 

66 

67 return OpenGaussClient() 

68 case VectorType.MARIADB_VECTOR: 68 ↛ 69line 68 didn't jump to line 69 because the pattern on line 68 never matched

69 from open_webui.retrieval.vector.dbs.mariadb_vector import ( 

70 MariaDBVectorClient, 

71 ) 

72 

73 return MariaDBVectorClient() 

74 case VectorType.ELASTICSEARCH: 74 ↛ 75line 74 didn't jump to line 75 because the pattern on line 74 never matched

75 from open_webui.retrieval.vector.dbs.elasticsearch import ( 

76 ElasticsearchClient, 

77 ) 

78 

79 return ElasticsearchClient() 

80 case VectorType.CHROMA: 80 ↛ 84line 80 didn't jump to line 84 because the pattern on line 80 always matched

81 from open_webui.retrieval.vector.dbs.chroma import ChromaClient 

82 

83 return ChromaClient() 

84 case VectorType.ORACLE23AI: 

85 from open_webui.retrieval.vector.dbs.oracle23ai import Oracle23aiClient 

86 

87 return Oracle23aiClient() 

88 case VectorType.WEAVIATE: 

89 from open_webui.retrieval.vector.dbs.weaviate import WeaviateClient 

90 

91 return WeaviateClient() 

92 case VectorType.VALKEY: 

93 from open_webui.retrieval.vector.dbs.valkey import ValkeyClient 

94 

95 return ValkeyClient() 

96 case _: 

97 raise ValueError(f'Unsupported vector type: {vector_type}') 

98 

99 

100VECTOR_DB_CLIENT = None if USE_SLIM else Vector.get_vector(VECTOR_DB) 

101_vector_client_lock = Lock() 

102 

103 

104def get_vector_db_client() -> VectorDBBase: 

105 """Initialize slim's remote client on first use so chat can start without it.""" 

106 global VECTOR_DB_CLIENT 

107 if VECTOR_DB_CLIENT is not None: 107 ↛ 109line 107 didn't jump to line 109 because the condition on line 107 was always true

108 return VECTOR_DB_CLIENT 

109 with _vector_client_lock: 

110 if VECTOR_DB_CLIENT is None: 

111 from open_webui import config 

112 

113 if VECTOR_DB == VectorType.PGVECTOR and not config.PGVECTOR_DB_URL.startswith('postgres'): 

114 raise HTTPException(503, 'Configure PGVECTOR_DB_URL for remote vector storage.') 

115 try: 

116 VECTOR_DB_CLIENT = Vector.get_vector(VECTOR_DB) 

117 except HTTPException: 

118 raise 

119 except Exception as exc: 

120 raise HTTPException( 

121 503, f'Unable to connect to configured vector database ({VECTOR_DB}): {exc}' 

122 ) from exc 

123 return VECTOR_DB_CLIENT