Coverage for haystack/dataclasses/document.py: 96%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import hashlib 

6import json 

7from dataclasses import asdict, dataclass, field, fields 

8from typing import Any 

9 

10from numpy import ndarray 

11 

12from haystack.dataclasses.byte_stream import ByteStream 

13from haystack.dataclasses.sparse_embedding import SparseEmbedding 

14from haystack.utils.dataclasses import _warn_on_inplace_mutation 

15 

16_LEGACY_FIELDS = ["content_type", "id_hash_keys", "dataframe"] 

17 

18 

19class _RemoveLegacyFields(type): 

20 def __call__(cls, *args: Any, **kwargs: Any) -> Any: 

21 """ 

22 Called before Document.__init__, removes the legacy fields. 

23 """ 

24 for field_name in _LEGACY_FIELDS: 

25 kwargs.pop(field_name, None) 

26 

27 return super().__call__(*args, **kwargs) 

28 

29 

30@_warn_on_inplace_mutation 

31@dataclass 

32class Document(metaclass=_RemoveLegacyFields): # noqa: PLW1641 

33 """ 

34 Base data class containing some data to be queried. 

35 

36 Can contain text snippets and file paths to images or audios. Documents can be sorted by score and saved 

37 to/from dictionary and JSON. 

38 

39 :param id: Unique identifier for the document. When not set, it's generated based on the Document fields' values. 

40 :param content: Text of the document, if the document contains text. 

41 :param blob: Binary data associated with the document, if the document has any binary data associated with it. 

42 :param meta: Additional custom metadata for the document. Must be JSON-serializable. 

43 :param score: Score of the document. Used for ranking, usually assigned by retrievers. 

44 :param embedding: dense vector representation of the document. 

45 :param sparse_embedding: sparse vector representation of the document. 

46 """ 

47 

48 id: str = field(default="") 

49 content: str | None = field(default=None) 

50 blob: ByteStream | None = field(default=None) 

51 meta: dict[str, Any] = field(default_factory=dict) 

52 score: float | None = field(default=None) 

53 embedding: list[float] | None = field(default=None) 

54 sparse_embedding: SparseEmbedding | None = field(default=None) 

55 

56 def __post_init__(self) -> None: 

57 """ 

58 Checks content type, converts embedding from 1.x type and generates the ID based on the init parameters. 

59 """ 

60 if self.content is not None and not isinstance(self.content, str): 

61 raise ValueError("The `content` field must be a string or None.") 

62 

63 # Embeddings were stored as NumPy arrays in 1.x, so we convert them to the new type 

64 if isinstance(self.embedding, ndarray): 

65 self.embedding = self.embedding.tolist() 

66 

67 # Generate an id only if not explicitly set 

68 self.id = self.id or self._create_id() 

69 

70 def __repr__(self) -> str: 

71 fields = [] 

72 if self.content is not None: 

73 fields.append( 

74 f"content: '{self.content}'" if len(self.content) < 100 else f"content: '{self.content[:100]}...'" 

75 ) 

76 if self.blob is not None: 

77 fields.append(f"blob: {len(self.blob.data)} bytes") 

78 if len(self.meta) > 0: 

79 fields.append(f"meta: {self.meta}") 

80 if self.score is not None: 

81 fields.append(f"score: {self.score}") 

82 if self.embedding is not None: 

83 fields.append(f"embedding: vector of size {len(self.embedding)}") 

84 if self.sparse_embedding is not None: 

85 fields.append(f"sparse_embedding: vector with {len(self.sparse_embedding.indices)} non-zero elements") 

86 fields_str = ", ".join(fields) 

87 return f"{self.__class__.__name__}(id={self.id}, {fields_str})" 

88 

89 def __eq__(self, other: object) -> bool: 

90 """ 

91 Compares Documents for equality. 

92 

93 Two Documents are considered equals if their dictionary representation is identical. 

94 """ 

95 if type(self) != type(other): 

96 return False 

97 return self.to_dict(flatten=False) == other.to_dict(flatten=False) 

98 

99 @classmethod 

100 def _field_names(cls) -> set[str]: 

101 return set(_LEGACY_FIELDS) | {f.name for f in fields(cls)} 

102 

103 def _create_id(self) -> str: 

104 """ 

105 Creates a hash of the given content that acts as the document's ID. 

106 """ 

107 text = self.content or None 

108 dataframe = None # this allows the ID creation to remain unchanged even if the dataframe field has been removed 

109 blob = self.blob.data if self.blob is not None else None 

110 mime_type = self.blob.mime_type if self.blob is not None else None 

111 # Sort keys so meta order doesn't affect the ID. Keep "{}" for empty meta so existing IDs stay stable. 

112 meta = json.dumps(self.meta, sort_keys=True, default=str) if self.meta else "{}" 

113 embedding = self.embedding if self.embedding is not None else None 

114 sparse_embedding = self.sparse_embedding.to_dict() if self.sparse_embedding is not None else "" 

115 data = f"{text}{dataframe}{blob!r}{mime_type}{meta}{embedding}{sparse_embedding}" 

116 return hashlib.sha256(data.encode("utf-8")).hexdigest() 

117 

118 def to_dict(self, flatten: bool = True) -> dict[str, Any]: 

119 """ 

120 Converts Document into a dictionary. 

121 

122 `blob` field is converted to a JSON-serializable type. 

123 

124 :param flatten: 

125 Whether to flatten the `meta` field. Defaults to `True` to be backward-compatible with Haystack 1.x. 

126 Meta keys that clash with document field names are kept in a nested `meta` dictionary. 

127 """ 

128 data = asdict(self) 

129 

130 # Use `ByteStream` and `SparseEmbedding`'s to_dict methods to convert them to JSON-serializable types. 

131 if self.blob is not None: 

132 data["blob"] = self.blob.to_dict() 

133 if self.sparse_embedding is not None: 

134 data["sparse_embedding"] = self.sparse_embedding.to_dict() 

135 

136 if not flatten: 

137 return data 

138 

139 field_names = self._field_names() 

140 flattened_meta: dict[str, Any] = {} 

141 colliding_meta: dict[str, Any] = {} 

142 for key, value in data.pop("meta").items(): 

143 if key in field_names: 

144 colliding_meta[key] = value # would clash with a document field: keep it nested 

145 else: 

146 flattened_meta[key] = value 

147 

148 data = {**flattened_meta, **data} 

149 if colliding_meta: 

150 data["meta"] = colliding_meta 

151 return data 

152 

153 @classmethod 

154 def from_dict(cls, data: dict[str, Any]) -> "Document": 

155 """ 

156 Creates a new Document object from a dictionary. 

157 

158 The `blob` field is converted to its original type. 

159 """ 

160 field_names = cls._field_names() 

161 field_data: dict[str, Any] = {} 

162 flattened_meta: dict[str, Any] = {} 

163 for key, value in data.items(): 

164 if key == "meta": 

165 continue # merged into meta at the end 

166 if key in field_names: 

167 field_data[key] = value # legacy fields included: _RemoveLegacyFields drops them before __init__ 

168 else: 

169 flattened_meta[key] = value 

170 

171 if blob := field_data.get("blob"): 

172 field_data["blob"] = ByteStream.from_dict(blob) 

173 if sparse_embedding := field_data.get("sparse_embedding"): 

174 field_data["sparse_embedding"] = SparseEmbedding.from_dict(sparse_embedding) 

175 

176 nested_meta = data.get("meta") or {} 

177 return cls(**field_data, meta={**nested_meta, **flattened_meta}) 

178 

179 @property 

180 def content_type(self) -> str: 

181 """ 

182 Returns the type of the content for the document. 

183 

184 This is necessary to keep backward compatibility with 1.x. 

185 """ 

186 if self.content is not None: 

187 return "text" 

188 raise ValueError("Content is not set.")