burr_101_haystack.py

from burr.core import action, State, ApplicationBuilder
from haystack.components.embedders import SentenceTransformersTextEmbedder


# 1. create actions
@action(reads=[], writes=["query_embedding"])
def embed_text(state: State, user_question: str) -> State:
    text_embedder = SentenceTransformersTextEmbedder(
        model="sentence-transformers/all-MiniLM-L6-v2"
    )
    results = text_embedder.run(text=user_question)
    return state.update(query_embedding=results["embedding"])


@action(reads=["query_embedding"], writes=["documents"])
def retrieve_documents(state: State) -> State:
    # ...
    return state.update(documents=...)


@action(reads=["documents"], writes=["question_prompt"])
def build_prompt(state: State, user_question: str) -> State:
    # ...
    return state.update(question_prompt=...)


@action(reads=["question_prompt"], writes=["answer"])
def generate_answer(state: State) -> State:
    # ...
    return state.update(answer=...)


# 2. define application
app = (
    ApplicationBuilder()
    .with_actions(
        embed_text,
        retrieve_documents,
        build_prompt,
        generate_answer
    )
    .with_transitions(
        ("embed_text", "retrieve_documents"),
        ("retrieve_documents", "build_prompt"),
        ("build_prompt", "generate_answer"))
    .with_entrypoint("embed_text")
    .build()
)

# 3. show application
app.visualize(include_state=True)

# 4. run application
user_query = "What is the capital of France?"
app.run(
    halt_after=["generate_answer"],
    inputs={
        "text": user_query,
        "question": user_query,
    }
)