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

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from dataclasses import asdict, dataclass, field 

6from typing import Any, Optional, Protocol, runtime_checkable 

7 

8from haystack.dataclasses import ChatMessage, Document 

9from haystack.utils.dataclasses import _warn_on_inplace_mutation 

10 

11 

12@runtime_checkable 

13@dataclass 

14class Answer(Protocol): 

15 data: Any 

16 query: str 

17 meta: dict[str, Any] 

18 

19 def to_dict(self) -> dict[str, Any]: # noqa: D102 

20 ... 

21 

22 @classmethod 

23 def from_dict(cls, data: dict[str, Any]) -> "Answer": # noqa: D102 

24 ... 

25 

26 

27@_warn_on_inplace_mutation 

28@dataclass 

29class ExtractedAnswer: 

30 """ 

31 Holds an answer extracted by an extractive Reader (query, score, text, and optional document/context). 

32 """ 

33 

34 query: str 

35 score: float 

36 data: str | None = None 

37 document: Document | None = None 

38 context: str | None = None 

39 document_offset: Optional["Span"] = None 

40 context_offset: Optional["Span"] = None 

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

42 

43 @_warn_on_inplace_mutation 

44 @dataclass 

45 class Span: 

46 start: int 

47 end: int 

48 

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

50 """ 

51 Serialize the object to a dictionary. 

52 

53 :returns: 

54 Serialized dictionary representation of the object. 

55 """ 

56 return { 

57 "data": self.data, 

58 "query": self.query, 

59 "document": self.document.to_dict(flatten=False) if self.document is not None else None, 

60 "context": self.context, 

61 "score": self.score, 

62 "document_offset": asdict(self.document_offset) if self.document_offset is not None else None, 

63 "context_offset": asdict(self.context_offset) if self.context_offset is not None else None, 

64 "meta": self.meta, 

65 } 

66 

67 @classmethod 

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

69 """ 

70 Deserialize the object from a dictionary. 

71 

72 :param data: 

73 Dictionary representation of the object. 

74 :returns: 

75 Deserialized object. 

76 """ 

77 # Backward compatibility: the old format wrapped the fields in an `init_parameters` envelope. 

78 if "init_parameters" in data: 

79 data = data["init_parameters"] 

80 

81 document = data.get("document") 

82 if document is not None: 

83 document = Document.from_dict(document) 

84 

85 document_offset = data.get("document_offset") 

86 if document_offset is not None: 

87 document_offset = ExtractedAnswer.Span(**document_offset) 

88 

89 context_offset = data.get("context_offset") 

90 if context_offset is not None: 

91 context_offset = ExtractedAnswer.Span(**context_offset) 

92 

93 return cls( 

94 data=data.get("data"), 

95 query=data["query"], 

96 score=data["score"], 

97 document=document, 

98 context=data.get("context"), 

99 document_offset=document_offset, 

100 context_offset=context_offset, 

101 meta=data.get("meta", {}), 

102 ) 

103 

104 

105@_warn_on_inplace_mutation 

106@dataclass 

107class GeneratedAnswer: 

108 """ 

109 Holds a generated answer from a Generator (answer text, query, referenced documents, and metadata). 

110 """ 

111 

112 data: str 

113 query: str 

114 documents: list[Document] 

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

116 

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

118 """ 

119 Serialize the object to a dictionary. 

120 

121 :returns: 

122 Serialized dictionary representation of the object. 

123 """ 

124 # all_messages is either a list of ChatMessage objects or a list of strings 

125 meta = self.meta 

126 all_messages = meta.get("all_messages") 

127 if all_messages and isinstance(all_messages[0], ChatMessage): 

128 meta = {**meta, "all_messages": [msg.to_dict() for msg in all_messages]} 

129 

130 return { 

131 "data": self.data, 

132 "query": self.query, 

133 "documents": [doc.to_dict(flatten=False) for doc in self.documents], 

134 "meta": meta, 

135 } 

136 

137 @classmethod 

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

139 """ 

140 Deserialize the object from a dictionary. 

141 

142 :param data: 

143 Dictionary representation of the object. 

144 

145 :returns: 

146 Deserialized object. 

147 """ 

148 # Backward compatibility: the old format wrapped the fields in an `init_parameters` envelope. 

149 if "init_parameters" in data: 

150 data = data["init_parameters"] 

151 

152 documents = [Document.from_dict(d) for d in data.get("documents", [])] 

153 

154 # Copy `meta` before converting `all_messages` so the caller's input dict is left untouched. 

155 meta = dict(data.get("meta", {})) 

156 if (all_messages := meta.get("all_messages")) and isinstance(all_messages[0], dict): 

157 meta["all_messages"] = [ChatMessage.from_dict(m) for m in all_messages] 

158 

159 return cls(data=data["data"], query=data["query"], documents=documents, meta=meta)