/ haystack / dataclasses / answer.py
answer.py
  1  # SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
  2  #
  3  # SPDX-License-Identifier: Apache-2.0
  4  
  5  from dataclasses import asdict, dataclass, field
  6  from typing import Any, Optional, Protocol, runtime_checkable
  7  
  8  from haystack.core.serialization import default_from_dict, default_to_dict
  9  from haystack.dataclasses import ChatMessage, Document
 10  from haystack.utils.dataclasses import _warn_on_inplace_mutation
 11  
 12  
 13  @runtime_checkable
 14  @dataclass
 15  class Answer(Protocol):
 16      data: Any
 17      query: str
 18      meta: dict[str, Any]
 19  
 20      def to_dict(self) -> dict[str, Any]:  # noqa: D102
 21          ...
 22  
 23      @classmethod
 24      def from_dict(cls, data: dict[str, Any]) -> "Answer":  # noqa: D102
 25          ...
 26  
 27  
 28  @_warn_on_inplace_mutation
 29  @dataclass
 30  class ExtractedAnswer:
 31      """
 32      Holds an answer extracted by an extractive Reader (query, score, text, and optional document/context).
 33      """
 34  
 35      query: str
 36      score: float
 37      data: str | None = None
 38      document: Document | None = None
 39      context: str | None = None
 40      document_offset: Optional["Span"] = None
 41      context_offset: Optional["Span"] = None
 42      meta: dict[str, Any] = field(default_factory=dict)
 43  
 44      @_warn_on_inplace_mutation
 45      @dataclass
 46      class Span:
 47          start: int
 48          end: int
 49  
 50      def to_dict(self) -> dict[str, Any]:
 51          """
 52          Serialize the object to a dictionary.
 53  
 54          :returns:
 55              Serialized dictionary representation of the object.
 56          """
 57          document = self.document.to_dict(flatten=False) if self.document is not None else None
 58          document_offset = asdict(self.document_offset) if self.document_offset is not None else None
 59          context_offset = asdict(self.context_offset) if self.context_offset is not None else None
 60          return default_to_dict(
 61              self,
 62              data=self.data,
 63              query=self.query,
 64              document=document,
 65              context=self.context,
 66              score=self.score,
 67              document_offset=document_offset,
 68              context_offset=context_offset,
 69              meta=self.meta,
 70          )
 71  
 72      @classmethod
 73      def from_dict(cls, data: dict[str, Any]) -> "ExtractedAnswer":
 74          """
 75          Deserialize the object from a dictionary.
 76  
 77          :param data:
 78              Dictionary representation of the object.
 79          :returns:
 80              Deserialized object.
 81          """
 82          init_params = data.get("init_parameters", {})
 83          if (doc := init_params.get("document")) is not None:
 84              data["init_parameters"]["document"] = Document.from_dict(doc)
 85  
 86          if (offset := init_params.get("document_offset")) is not None:
 87              data["init_parameters"]["document_offset"] = ExtractedAnswer.Span(**offset)
 88  
 89          if (offset := init_params.get("context_offset")) is not None:
 90              data["init_parameters"]["context_offset"] = ExtractedAnswer.Span(**offset)
 91          return default_from_dict(cls, data)
 92  
 93  
 94  @_warn_on_inplace_mutation
 95  @dataclass
 96  class GeneratedAnswer:
 97      """
 98      Holds a generated answer from a Generator (answer text, query, referenced documents, and metadata).
 99      """
100  
101      data: str
102      query: str
103      documents: list[Document]
104      meta: dict[str, Any] = field(default_factory=dict)
105  
106      def to_dict(self) -> dict[str, Any]:
107          """
108          Serialize the object to a dictionary.
109  
110          :returns:
111              Serialized dictionary representation of the object.
112          """
113          documents = [doc.to_dict(flatten=False) for doc in self.documents]
114  
115          # Serialize ChatMessage objects to dicts
116          meta = self.meta
117          all_messages = meta.get("all_messages")
118  
119          # all_messages is either a list of ChatMessage objects or a list of strings
120          if all_messages and isinstance(all_messages[0], ChatMessage):
121              meta = {**meta, "all_messages": [msg.to_dict() for msg in all_messages]}
122  
123          return default_to_dict(self, data=self.data, query=self.query, documents=documents, meta=meta)
124  
125      @classmethod
126      def from_dict(cls, data: dict[str, Any]) -> "GeneratedAnswer":
127          """
128          Deserialize the object from a dictionary.
129  
130          :param data:
131              Dictionary representation of the object.
132  
133          :returns:
134              Deserialized object.
135          """
136          init_params = data.get("init_parameters", {})
137  
138          if (documents := init_params.get("documents")) is not None:
139              init_params["documents"] = [Document.from_dict(d) for d in documents]
140  
141          meta = init_params.get("meta", {})
142          if (all_messages := meta.get("all_messages")) is not None and isinstance(all_messages[0], dict):
143              meta["all_messages"] = [ChatMessage.from_dict(m) for m in all_messages]
144          init_params["meta"] = meta
145  
146          return default_from_dict(cls, data)