Coverage for open_webui/retrieval/vector/main.py: 77%

1from abc import ABC, abstractmethod 

2from typing import Any, Dict, List, Optional, Union 

3 

4from pydantic import BaseModel 

5 

6 

7class VectorItem(BaseModel): 

8 id: str 

9 text: str 

10 vector: List[float | int] 

11 metadata: Any 

12 

13 

14class GetResult(BaseModel): 

15 ids: Optional[List[List[str]]] 

16 documents: Optional[List[List[str]]] 

17 metadatas: Optional[List[List[Any]]] 

18 

19 

20class SearchResult(GetResult): 

21 distances: Optional[List[List[float | int]]] 

22 

23 

24class VectorDBBase(ABC): 

25 """ 

26 Abstract base class for all vector database backends. 

27 

28 Implementations of this class provide methods for collection management, 

29 vector insertion, deletion, similarity search, and metadata filtering. 

30 

31 Any custom vector database integration must inherit from this class and 

32 implement all abstract methods. 

33 """ 

34 

35 @abstractmethod 

36 def has_collection(self, collection_name: str) -> bool: 

37 """Check if the collection exists in the vector DB.""" 

38 pass 

39 

40 @abstractmethod 

41 def delete_collection(self, collection_name: str) -> None: 

42 """Delete a collection from the vector DB.""" 

43 pass 

44 

45 @abstractmethod 

46 def insert(self, collection_name: str, items: List[VectorItem]) -> None: 

47 """Insert a list of vector items into a collection.""" 

48 pass 

49 

50 @abstractmethod 

51 def upsert(self, collection_name: str, items: List[VectorItem]) -> None: 

52 """Insert or update vector items in a collection.""" 

53 pass 

54 

55 @abstractmethod 

56 def search( 

57 self, 

58 collection_name: str, 

59 vectors: List[List[Union[float, int]]], 

60 filter: Optional[Dict] = None, 

61 limit: int = 10, 

62 ) -> Optional[SearchResult]: 

63 """Search for similar vectors in a collection.""" 

64 pass 

65 

66 def hybrid_search( 

67 self, 

68 collection_name: str, 

69 query: str, 

70 vectors: List[List[Union[float, int]]], 

71 filter: Optional[Dict] = None, 

72 limit: int = 10, 

73 hybrid_bm25_weight: float = 0.5, 

74 ) -> Optional[SearchResult]: 

75 """Search using a backend-native hybrid keyword/vector implementation when available.""" 

76 return None 

77 

78 @abstractmethod 

79 def query(self, collection_name: str, filter: Dict, limit: Optional[int] = None) -> Optional[GetResult]: 

80 """Query vectors from a collection using metadata filter.""" 

81 pass 

82 

83 @abstractmethod 

84 def get(self, collection_name: str) -> Optional[GetResult]: 

85 """Retrieve all vectors from a collection.""" 

86 pass 

87 

88 @abstractmethod 

89 def delete( 

90 self, 

91 collection_name: str, 

92 ids: Optional[List[str]] = None, 

93 filter: Optional[Dict] = None, 

94 ) -> None: 

95 """Delete vectors by ID or filter from a collection.""" 

96 pass 

97 

98 @abstractmethod 

99 def reset(self) -> None: 

100 """Reset the vector database by removing all collections or those matching a condition.""" 

101 pass