buster/formatter/base.py

from dataclasses import dataclass
from typing import Iterable, NamedTuple


# Should be from the `documents` module.
class Source(NamedTuple):
    title: str
    url: str
    question_similarity: float
    source: str = ""
    # TODO Add answer similarity.
    # answer_similarity: float


# Should be from the `nlp` module.
@dataclass(slots=True)
class Response:
    text: str
    error: bool = False
    error_msg: str | None = None


@dataclass
class ResponseFormatter:
    response_footnote: str
    source_template: str = "{source.name} (relevance: {source.question_similarity:2.1f})"
    error_msg_template: str = """Something went wrong:\n{response.error_msg}"""
    error_fallback_template: str = "Something went very wrong."
    sourced_answer_template: str = (
        """{response.text}\n\n"""
        """📝 Here are the sources I used to answer your question:\n"""
        """{sources}\n\n"""
        """{footnote}"""
    )
    unsourced_answer_template: str = "{response.text}\n\n{footnote}"

    def source_item(self, source: Source) -> str:
        """Format a single source item."""
        return self.source_template.format(source=source)

    def sources_list(self, sources: Iterable[Source]) -> str | None:
        """Format sources into a list."""
        items = [self.source_item(source) for source in sources]
        if not items:
            return None  # No list needed.

        return "\n".join(f"{ind}. {item}" for ind, item in enumerate(items, 1))

    def error(self, response: Response) -> str:
        """Format an error message."""
        if response.error_msg:
            return self.error_msg_template.format(response=response)
        return self.error_fallback_template.format(response=response)

    def answer(self, response: Response, sources: Iterable[Source]) -> str:
        """Format an answer and its sources."""
        sources_list = self.sources_list(sources)
        if sources_list:
            return self.sourced_answer_template.format(
                response=response, sources=sources_list, footnote=self.response_footnote
            )

        return self.unsourced_answer_template.format(response=response, footnote=self.response_footnote)

    def __call__(self, response: Response, sources: Iterable[Source]) -> str:
        """Format an answer and its sources, or an error message."""
        if response.error:
            return self.error(response)
        return self.answer(response, sources)