Testing components and pipelines

from shop_chat import ShopChat

reply = ShopChat().run(messages=[ChatMessage.from_user("- Parcels ship in two days.\nQuestion: when do parcels ship")])["replies"][0]
assert reply.text == "Parcels ship in two days."
print("passed")

from haystack import Document, Pipeline, SuperComponent
from haystack.components.builders import AnswerBuilder, ChatPromptBuilder
from haystack.components.retrievers import InMemoryBM25Retriever
from haystack.dataclasses import ChatMessage
from haystack.document_stores.in_memory import InMemoryDocumentStore

TEMPLATE = [ChatMessage.from_user(
    "Answer the question using only these documents.\n"
    "{% for document in documents %}- {{ document.content }}\n{% endfor %}"
    "Question: {{ question }}"
)]


def build_desk(documents: list[Document], generator, audience: str) -> SuperComponent:
    store = InMemoryDocumentStore()
    store.write_documents(documents)
    only_audience = {"field": "meta.audience", "operator": "in", "value": ["public", audience]}

    rag = Pipeline()
    rag.add_component("retriever", InMemoryBM25Retriever(document_store=store, top_k=2, filters=only_audience))
    rag.add_component("prompt", ChatPromptBuilder(template=TEMPLATE, required_variables=["question", "documents"]))
    rag.add_component("llm", generator)
    rag.add_component("answer", AnswerBuilder())
    rag.connect("retriever.documents", "prompt.documents")
    rag.connect("prompt.prompt", "llm.messages")
    rag.connect("llm.replies", "answer.replies")
    rag.connect("retriever.documents", "answer.documents")
    return SuperComponent(
        pipeline=rag,
        input_mapping={"question": ["retriever.query", "prompt.question", "answer.query"]},
        output_mapping={"answer.answers": "answers"},
    )

from haystack import Document

POLICIES = [
    Document(content="Refunds are paid within five working days of receiving the return.", meta={"audience": "public"}),
    Document(content="Damaged items can be returned for free within 30 days.", meta={"audience": "public"}),
    Document(content="Parcels ship within two days and arrive by courier.", meta={"audience": "public"}),
    Document(content="Staff may approve refunds up to 100 euros without a manager.", meta={"audience": "staff"}),
]

from haystack import Document, component
from haystack.dataclasses import ChatMessage

from desk import build_desk
from policies import POLICIES
from shop_chat import ShopChat


@component
class RecordingChat:
    """Replies with a fixed text and keeps the prompt it was given."""

    def __init__(self):
        self.prompts = []

    @component.output_types(replies=list[ChatMessage])
    def run(self, messages: list[ChatMessage]):
        self.prompts.append(messages[-1].text)
        return {"replies": [ChatMessage.from_assistant("ok")]}


def test_customers_never_see_staff_rules():
    chat = RecordingChat()
    desk = build_desk(POLICIES, chat, audience="customer")
    desk.run(question="who can approve refunds")
    assert "Staff may approve" not in chat.prompts[0]


def test_staff_see_staff_rules():
    chat = RecordingChat()
    desk = build_desk(POLICIES, chat, audience="staff")
    desk.run(question="who can approve refunds")
    assert "Staff may approve" in chat.prompts[0]


def test_answers_carry_their_sources():
    desk = build_desk(POLICIES, ShopChat(), audience="customer")
    answer = desk.run(question="can damaged items be returned")["answers"][0]
    assert answer.data == "Damaged items can be returned for free within 30 days."
    assert answer.documents[0].content == answer.data

pytest -q