deepset-ai--haystack/test/dataclasses/test_answer.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import warnings
from copy import deepcopy

import pytest

from haystack.dataclasses import ChatMessage, Document
from haystack.dataclasses.answer import Answer, ExtractedAnswer, GeneratedAnswer


class TestExtractedAnswer:
    def test_init(self):
        answer = ExtractedAnswer(
            data="42",
            query="What is the answer?",
            document=Document(content="I thought a lot about this. The answer is 42."),
            context="The answer is 42.",
            score=1.0,
            document_offset=ExtractedAnswer.Span(42, 44),
            context_offset=ExtractedAnswer.Span(14, 16),
            meta={"meta_key": "meta_value"},
        )
        assert answer.data == "42"
        assert answer.query == "What is the answer?"
        assert answer.document == Document(content="I thought a lot about this. The answer is 42.")
        assert answer.context == "The answer is 42."
        assert answer.score == 1.0
        assert answer.document_offset == ExtractedAnswer.Span(42, 44)
        assert answer.context_offset == ExtractedAnswer.Span(14, 16)
        assert answer.meta == {"meta_key": "meta_value"}

    def test_protocol(self):
        answer = ExtractedAnswer(
            data="42",
            query="What is the answer?",
            document=Document(content="I thought a lot about this. The answer is 42."),
            context="The answer is 42.",
            score=1.0,
            document_offset=ExtractedAnswer.Span(42, 44),
            context_offset=ExtractedAnswer.Span(14, 16),
            meta={"meta_key": "meta_value"},
        )
        assert isinstance(answer, Answer)

    def test_to_dict(self):
        document = Document(content="I thought a lot about this. The answer is 42.")
        answer = ExtractedAnswer(
            data="42",
            query="What is the answer?",
            document=document,
            context="The answer is 42.",
            score=1.0,
            document_offset=ExtractedAnswer.Span(42, 44),
            context_offset=ExtractedAnswer.Span(14, 16),
            meta={"meta_key": "meta_value"},
        )
        assert answer.to_dict() == {
            "type": "haystack.dataclasses.answer.ExtractedAnswer",
            "init_parameters": {
                "data": "42",
                "query": "What is the answer?",
                "document": document.to_dict(flatten=False),
                "context": "The answer is 42.",
                "score": 1.0,
                "document_offset": {"start": 42, "end": 44},
                "context_offset": {"start": 14, "end": 16},
                "meta": {"meta_key": "meta_value"},
            },
        }

    def test_from_dict(self):
        answer = ExtractedAnswer.from_dict(
            {
                "type": "haystack.dataclasses.answer.ExtractedAnswer",
                "init_parameters": {
                    "data": "42",
                    "query": "What is the answer?",
                    "document": {
                        "id": "8f800a524b139484fc719ecc35f971a080de87618319bc4836b784d69baca57f",
                        "content": "I thought a lot about this. The answer is 42.",
                    },
                    "context": "The answer is 42.",
                    "score": 1.0,
                    "document_offset": {"start": 42, "end": 44},
                    "context_offset": {"start": 14, "end": 16},
                    "meta": {"meta_key": "meta_value"},
                },
            }
        )
        assert answer.data == "42"
        assert answer.query == "What is the answer?"
        assert answer.document == Document(
            id="8f800a524b139484fc719ecc35f971a080de87618319bc4836b784d69baca57f",
            content="I thought a lot about this. The answer is 42.",
        )
        assert answer.context == "The answer is 42."
        assert answer.score == 1.0
        assert answer.document_offset == ExtractedAnswer.Span(42, 44)
        assert answer.context_offset == ExtractedAnswer.Span(14, 16)
        assert answer.meta == {"meta_key": "meta_value"}

    def test_from_dict_does_not_mutate_input(self):
        data = {
            "type": "haystack.dataclasses.answer.ExtractedAnswer",
            "init_parameters": {
                "data": "42",
                "query": "What is the answer?",
                "document": {
                    "id": "8f800a524b139484fc719ecc35f971a080de87618319bc4836b784d69baca57f",
                    "content": "I thought a lot about this. The answer is 42.",
                },
                "context": "The answer is 42.",
                "score": 1.0,
                "document_offset": {"start": 42, "end": 44},
                "context_offset": {"start": 14, "end": 16},
                "meta": {"meta_key": "meta_value"},
            },
        }
        snapshot = deepcopy(data)
        first = ExtractedAnswer.from_dict(data)
        # from_dict must not mutate its input dictionary
        assert data == snapshot
        # deserializing the same dictionary again must still work and be equal
        assert ExtractedAnswer.from_dict(data) == first

    def test_no_warning_on_init(self):
        with warnings.catch_warnings():
            warnings.simplefilter("error", Warning)
            ExtractedAnswer(query="q", score=1.0)

    def test_warn_on_inplace_mutation(self):
        answer = ExtractedAnswer(query="q", score=1.0)
        with pytest.warns(Warning, match="dataclasses.replace"):
            answer.query = "new"

    def test_span_no_warning_on_init(self):
        with warnings.catch_warnings():
            warnings.simplefilter("error", Warning)
            ExtractedAnswer.Span(start=0, end=5)

    def test_span_warn_on_inplace_mutation(self):
        span = ExtractedAnswer.Span(start=0, end=5)
        with pytest.warns(Warning, match="dataclasses.replace"):
            span.start = 1


class TestGeneratedAnswer:
    def test_init(self):
        answer = GeneratedAnswer(
            data="42",
            query="What is the answer?",
            documents=[
                Document(id="1", content="The answer is 42."),
                Document(id="2", content="I believe the answer is 42."),
                Document(id="3", content="42 is definitely the answer."),
            ],
            meta={"meta_key": "meta_value"},
        )
        assert answer.data == "42"
        assert answer.query == "What is the answer?"
        assert answer.documents == [
            Document(id="1", content="The answer is 42."),
            Document(id="2", content="I believe the answer is 42."),
            Document(id="3", content="42 is definitely the answer."),
        ]
        assert answer.meta == {"meta_key": "meta_value"}

    def test_protocol(self):
        answer = GeneratedAnswer(
            data="42",
            query="What is the answer?",
            documents=[
                Document(id="1", content="The answer is 42."),
                Document(id="2", content="I believe the answer is 42."),
                Document(id="3", content="42 is definitely the answer."),
            ],
            meta={"meta_key": "meta_value"},
        )
        assert isinstance(answer, Answer)

    def test_to_dict(self):
        answer = GeneratedAnswer(data="42", query="What is the answer?", documents=[])
        assert answer.to_dict() == {
            "type": "haystack.dataclasses.answer.GeneratedAnswer",
            "init_parameters": {"data": "42", "query": "What is the answer?", "documents": [], "meta": {}},
        }

    def test_to_dict_with_meta(self):
        answer = GeneratedAnswer(
            data="42",
            query="What is the answer?",
            documents=[],
            meta={"meta_key": "meta_value", "all_messages": ["What is the answer?"]},
        )
        assert answer.to_dict() == {
            "type": "haystack.dataclasses.answer.GeneratedAnswer",
            "init_parameters": {
                "data": "42",
                "query": "What is the answer?",
                "documents": [],
                "meta": {"meta_key": "meta_value", "all_messages": ["What is the answer?"]},
            },
        }

    def test_to_dict_with_chat_message_in_meta(self):
        documents = [
            Document(id="1", content="The answer is 42."),
            Document(id="2", content="I believe the answer is 42."),
            Document(id="3", content="42 is definitely the answer."),
        ]
        answer = GeneratedAnswer(
            data="42",
            query="What is the answer?",
            documents=documents,
            meta={"meta_key": "meta_value", "all_messages": [ChatMessage.from_user("What is the answer?")]},
        )
        assert answer.to_dict() == {
            "type": "haystack.dataclasses.answer.GeneratedAnswer",
            "init_parameters": {
                "data": "42",
                "query": "What is the answer?",
                "documents": [d.to_dict(flatten=False) for d in documents],
                "meta": {
                    "meta_key": "meta_value",
                    "all_messages": [ChatMessage.from_user("What is the answer?").to_dict()],
                },
            },
        }

    def test_from_dict(self):
        answer = GeneratedAnswer.from_dict(
            {
                "type": "haystack.dataclasses.answer.GeneratedAnswer",
                "init_parameters": {"data": "42", "query": "What is the answer?", "documents": [], "meta": {}},
            }
        )
        assert answer.data == "42"
        assert answer.query == "What is the answer?"
        assert answer.documents == []
        assert answer.meta == {}

    def test_from_dict_with_meta(self):
        answer = GeneratedAnswer.from_dict(
            {
                "type": "haystack.dataclasses.answer.GeneratedAnswer",
                "init_parameters": {
                    "data": "42",
                    "query": "What is the answer?",
                    "documents": [],
                    "meta": {"meta_key": "meta_value", "all_messages": ["What is the answer?"]},
                },
            }
        )
        assert answer.data == "42"
        assert answer.query == "What is the answer?"
        assert answer.documents == []
        assert answer.meta["meta_key"] == "meta_value"
        assert answer.meta["all_messages"] == ["What is the answer?"]

    def test_from_dict_with_chat_message_in_meta(self):
        answer = GeneratedAnswer.from_dict(
            {
                "type": "haystack.dataclasses.answer.GeneratedAnswer",
                "init_parameters": {
                    "data": "42",
                    "query": "What is the answer?",
                    "documents": [
                        {"id": "1", "content": "The answer is 42."},
                        {"id": "2", "content": "I believe the answer is 42."},
                        {"id": "3", "content": "42 is definitely the answer."},
                    ],
                    "meta": {
                        "meta_key": "meta_value",
                        "all_messages": [ChatMessage.from_user("What is the answer?").to_dict()],
                    },
                },
            }
        )
        assert answer.data == "42"
        assert answer.query == "What is the answer?"
        assert answer.documents == [
            Document(id="1", content="The answer is 42."),
            Document(id="2", content="I believe the answer is 42."),
            Document(id="3", content="42 is definitely the answer."),
        ]
        assert answer.meta["meta_key"] == "meta_value"
        assert answer.meta["all_messages"] == [ChatMessage.from_user("What is the answer?")]

    def test_from_dict_does_not_mutate_input(self):
        data = {
            "type": "haystack.dataclasses.answer.GeneratedAnswer",
            "init_parameters": {
                "data": "42",
                "query": "What is the answer?",
                "documents": [
                    {"id": "1", "content": "The answer is 42."},
                    {"id": "2", "content": "I believe the answer is 42."},
                ],
                "meta": {
                    "meta_key": "meta_value",
                    "all_messages": [ChatMessage.from_user("What is the answer?").to_dict()],
                },
            },
        }
        snapshot = deepcopy(data)
        first = GeneratedAnswer.from_dict(data)
        # from_dict must not mutate its input dictionary
        assert data == snapshot
        # deserializing the same dictionary again must still work and be equal
        assert GeneratedAnswer.from_dict(data) == first

    def test_from_dict_with_empty_all_messages(self):
        # An empty `all_messages` list must not crash deserialization: `is not None`
        # let `[]` through and then indexed `all_messages[0]`, raising IndexError.
        answer = GeneratedAnswer.from_dict(
            {
                "type": "haystack.dataclasses.answer.GeneratedAnswer",
                "init_parameters": {
                    "data": "42",
                    "query": "What is the answer?",
                    "documents": [],
                    "meta": {"all_messages": []},
                },
            }
        )
        assert answer.meta["all_messages"] == []

    def test_to_dict_from_dict_round_trip_with_empty_all_messages(self):
        answer = GeneratedAnswer(data="42", query="What is the answer?", documents=[], meta={"all_messages": []})
        assert GeneratedAnswer.from_dict(answer.to_dict()) == answer

    def test_no_warning_on_init(self):
        with warnings.catch_warnings():
            warnings.simplefilter("error", Warning)
            GeneratedAnswer(data="42", query="q", documents=[])

    def test_warn_on_inplace_mutation(self):
        answer = GeneratedAnswer(data="42", query="q", documents=[])
        with pytest.warns(Warning, match="dataclasses.replace"):
            answer.data = "new"