Coverage for open_webui/retrieval/vector/utils.py: 27%

1import datetime as dt 

2from typing import Any 

3 

4from open_webui.env import RAG_METADATA_MAX_VALUE_CHARS 

5from open_webui.retrieval.vector.main import SearchResult 

6from open_webui.utils.misc import sanitize_text_for_db 

7 

8KEYS_TO_EXCLUDE = [ 

9 'content', 

10 'pages', 

11 'tables', 

12 'paragraphs', 

13 'sections', 

14 'figures', 

15 'documents', 

16 'keyValuePairs', 

17 'styles', 

18 'languages', 

19] 

20 

21 

22def filter_metadata(metadata: dict[str, any]) -> dict[str, any]: 

23 # Removes large/redundant fields from metadata dict. 

24 result = {} 

25 for key, value in metadata.items(): 

26 if key in KEYS_TO_EXCLUDE: 26 ↛ 27line 26 didn't jump to line 27 because the condition on line 26 was never true

27 continue 

28 if RAG_METADATA_MAX_VALUE_CHARS is not None and isinstance(value, (list, dict)): 28 ↛ 29line 28 didn't jump to line 29 because the condition on line 28 was never true

29 try: 

30 if len(str(value)) > RAG_METADATA_MAX_VALUE_CHARS: 

31 continue 

32 except (MemoryError, RecursionError, ValueError): 

33 continue 

34 result[key] = value 

35 return result 

36 

37 

38def process_metadata( 

39 metadata: dict[str, any], 

40) -> dict[str, any]: 

41 # Removes large fields, converts non-serializable types (datetime, list, dict) to strings, 

42 # and sanitizes strings for database storage (strips null bytes and invalid surrogates). 

43 result = {} 

44 for key, value in metadata.items(): 

45 # Skip large fields 

46 if key in KEYS_TO_EXCLUDE: 46 ↛ 47line 46 didn't jump to line 47 because the condition on line 46 was never true

47 continue 

48 if value is None: 48 ↛ 49line 48 didn't jump to line 49 because the condition on line 48 was never true

49 continue 

50 if RAG_METADATA_MAX_VALUE_CHARS is not None and isinstance(value, (list, dict)): 50 ↛ 51line 50 didn't jump to line 51 because the condition on line 50 was never true

51 try: 

52 if len(str(value)) > RAG_METADATA_MAX_VALUE_CHARS: 

53 continue 

54 except (MemoryError, RecursionError, ValueError): 

55 continue 

56 # Convert non-serializable fields to strings 

57 if isinstance(value, (dt.datetime, list, dict)): 57 ↛ 58line 57 didn't jump to line 58 because the condition on line 57 was never true

58 result[key] = sanitize_text_for_db(str(value)) 

59 else: 

60 result[key] = sanitize_text_for_db(value) 

61 return result 

62 

63 

64def iter_filter_conditions(filter: dict[str, Any] | None): 

65 for key, value in (filter or {}).items(): 

66 if isinstance(value, dict): 

67 if set(value) != {'$in'}: 

68 raise ValueError(f"Unsupported metadata filter for '{key}': {value}") 

69 yield key, '$in', list(value['$in']) 

70 else: 

71 yield key, '$eq', value 

72 

73 

74def normalize_filter(filter: dict[str, Any] | None) -> dict[str, Any]: 

75 return {key: {'$in': value} if op == '$in' else value for key, op, value in iter_filter_conditions(filter)} 

76 

77 

78def metadata_matches_filter(metadata: dict[str, Any], filter: dict[str, Any] | None) -> bool: 

79 if not isinstance(metadata, dict): 

80 return False 

81 for key, op, value in iter_filter_conditions(filter): 

82 actual = metadata.get(key) 

83 if op == '$in': 

84 if actual not in value: 

85 return False 

86 elif actual != value: 

87 return False 

88 return True 

89 

90 

91def merge_hybrid_search_results( 

92 vector_result: SearchResult | None, 

93 fts_results: list[dict[str, Any]], 

94 num_queries: int, 

95 limit: int, 

96 hybrid_bm25_weight: float, 

97) -> SearchResult: 

98 rank_constant = 60.0 

99 bm25_weight = min(max(hybrid_bm25_weight, 0.0), 1.0) 

100 vector_weight = 1.0 - bm25_weight 

101 

102 ids = [[] for _ in range(num_queries)] 

103 distances = [[] for _ in range(num_queries)] 

104 documents = [[] for _ in range(num_queries)] 

105 metadatas = [[] for _ in range(num_queries)] 

106 

107 for qid in range(num_queries): 

108 candidates: dict[str, dict[str, Any]] = {} 

109 

110 if vector_result and vector_result.ids and qid < len(vector_result.ids): 

111 for rank, item_id in enumerate(vector_result.ids[qid] or [], start=1): 

112 score = vector_weight / (rank_constant + rank) if vector_weight > 0 else 0 

113 if score <= 0: 

114 continue 

115 

116 candidate = candidates.setdefault( 

117 item_id, 

118 { 

119 'score': 0.0, 

120 'document': vector_result.documents[qid][rank - 1], 

121 'metadata': vector_result.metadatas[qid][rank - 1], 

122 }, 

123 ) 

124 candidate['score'] += score 

125 

126 for rank, row in enumerate(fts_results, start=1): 

127 score = bm25_weight / (rank_constant + rank) if bm25_weight > 0 else 0 

128 if score <= 0: 

129 continue 

130 

131 item_id = row['id'] 

132 candidate = candidates.setdefault( 

133 item_id, 

134 { 

135 'score': 0.0, 

136 'document': row['text'], 

137 'metadata': row['vmetadata'], 

138 }, 

139 ) 

140 candidate['score'] += score 

141 

142 ranked = sorted(candidates.items(), key=lambda item: item[1]['score'], reverse=True)[:limit] 

143 ids[qid] = [item_id for item_id, _ in ranked] 

144 distances[qid] = [candidate['score'] for _, candidate in ranked] 

145 documents[qid] = [candidate['document'] for _, candidate in ranked] 

146 metadatas[qid] = [candidate['metadata'] for _, candidate in ranked] 

147 

148 return SearchResult(ids=ids, distances=distances, documents=documents, metadatas=metadatas)