Coverage for haystack/dataclasses/byte_stream.py: 100%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from dataclasses import dataclass, field 

6from pathlib import Path 

7from typing import Any 

8 

9from haystack.utils.dataclasses import _warn_on_inplace_mutation 

10from haystack.utils.misc import _guess_mime_type 

11 

12 

13@_warn_on_inplace_mutation 

14@dataclass(repr=False) 

15class ByteStream: 

16 """ 

17 Base data class representing a binary object in the Haystack API. 

18 

19 :param data: The binary data stored in Bytestream. 

20 :param meta: Additional metadata to be stored with the ByteStream. 

21 :param mime_type: The mime type of the binary data. 

22 """ 

23 

24 data: bytes 

25 meta: dict[str, Any] = field(default_factory=dict, hash=False) 

26 mime_type: str | None = field(default=None) 

27 

28 def to_file(self, destination_path: Path) -> None: 

29 """ 

30 Write the ByteStream to a file. Note: the metadata will be lost. 

31 

32 :param destination_path: The path to write the ByteStream to. 

33 """ 

34 with open(destination_path, "wb") as fd: 

35 fd.write(self.data) 

36 

37 @classmethod 

38 def from_file_path( 

39 cls, 

40 filepath: Path, 

41 mime_type: str | None = None, 

42 meta: dict[str, Any] | None = None, 

43 guess_mime_type: bool = False, 

44 ) -> "ByteStream": 

45 """ 

46 Create a ByteStream from the contents read from a file. 

47 

48 :param filepath: A valid path to a file. 

49 :param mime_type: The mime type of the file. 

50 :param meta: Additional metadata to be stored with the ByteStream. 

51 :param guess_mime_type: Whether to guess the mime type from the file. 

52 """ 

53 if not mime_type and guess_mime_type: 

54 mime_type = _guess_mime_type(filepath) 

55 with open(filepath, "rb") as fd: 

56 return cls(data=fd.read(), mime_type=mime_type, meta=meta or {}) 

57 

58 @classmethod 

59 def from_string( 

60 cls, text: str, encoding: str = "utf-8", mime_type: str | None = None, meta: dict[str, Any] | None = None 

61 ) -> "ByteStream": 

62 """ 

63 Create a ByteStream encoding a string. 

64 

65 :param text: The string to encode 

66 :param encoding: The encoding used to convert the string into bytes 

67 :param mime_type: The mime type of the file. 

68 :param meta: Additional metadata to be stored with the ByteStream. 

69 """ 

70 return cls(data=text.encode(encoding), mime_type=mime_type, meta=meta or {}) 

71 

72 def to_string(self, encoding: str = "utf-8") -> str: 

73 """ 

74 Convert the ByteStream to a string, metadata will not be included. 

75 

76 :param encoding: The encoding used to convert the bytes to a string. Defaults to "utf-8". 

77 :returns: The string representation of the ByteStream. 

78 :raises UnicodeDecodeError: If the ByteStream data cannot be decoded with the specified encoding. 

79 """ 

80 return self.data.decode(encoding) 

81 

82 def __repr__(self) -> str: 

83 """ 

84 Return a string representation of the ByteStream, truncating the data to 100 bytes. 

85 """ 

86 fields = [] 

87 truncated_data = self.data[:100] + b"..." if len(self.data) > 100 else self.data 

88 fields.append(f"data={truncated_data!r}") 

89 fields.append(f"meta={self.meta!r}") 

90 fields.append(f"mime_type={self.mime_type!r}") 

91 fields_str = ", ".join(fields) 

92 return f"{self.__class__.__name__}({fields_str})" 

93 

94 def to_dict(self) -> dict[str, Any]: 

95 """ 

96 Convert the ByteStream to a dictionary representation. 

97 

98 :returns: A dictionary with keys 'data', 'meta', and 'mime_type'. 

99 """ 

100 # Note: The data is converted to a list of integers for serialization since JSON does not support bytes 

101 # directly. 

102 return {"data": list(self.data), "meta": self.meta, "mime_type": self.mime_type} 

103 

104 def _to_trace_dict(self) -> dict[str, Any]: 

105 """ 

106 Convert the ByteStream to a dictionary representation for tracing. 

107 

108 Binary data is replaced with a placeholder string to avoid sending large payloads to the tracing backend. 

109 

110 :returns: 

111 Serialized version of the object only for tracing purposes. 

112 """ 

113 return {"data": f"Binary data ({len(self.data)} bytes)", "meta": self.meta, "mime_type": self.mime_type} 

114 

115 @classmethod 

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

117 """ 

118 Create a ByteStream from a dictionary representation. 

119 

120 :param data: A dictionary with keys 'data', 'meta', and 'mime_type'. 

121 

122 :returns: A ByteStream instance. 

123 """ 

124 return ByteStream(data=bytes(data["data"]), meta=data.get("meta", {}), mime_type=data.get("mime_type"))