Basic RAG inference pipeline

Use your RAG components to generate responses to prompts.

Now that we have our index store, we can use it to make queries based on the documents in the index store. We use some utility functions to make this happen but no external libraries are needed beyond an interface to the index store as well as the LLM itself.

If you've been following along with the guide, you should have some documents ingested already and you can pass a query in as a flag to the Python command used to run the pipeline:

python run.py --rag-query "how do I use a custom materializer inside my own zenml

steps? i.e. how do I set it? inside the @step decorator?" --model=gpt4

This inference query itself is not a ZenML pipeline, but rather a function call which uses the outputs and components of our pipeline to generate the response. For a more complex inference setup, there might be even more going on here, but for the purposes of this initial guide we will keep it simple.

Bringing everything together, the code for the inference pipeline is as follows:

def process_input_with_retrieval(

input: str, model: str = OPENAI_MODEL, n_items_retrieved: int = 5

) -> str:

delimiter = "```"

# Step 1: Get documents related to the user input from database

related_docs = get_topn_similar_docs(

get_embeddings(input), get_db_conn(), n=n_items_retrieved

# Step 2: Get completion from OpenAI API

# Set system message to help set appropriate tone and context for model

system_message = f"""

You are a friendly chatbot. \

You can answer questions about ZenML, its features and its use cases. \

You respond in a concise, technically credible tone. \

You ONLY use the context from the ZenML documentation to provide relevant

answers. \

You do not make up answers or provide opinions that you don't have

information to support. \

If you are unsure or don't know, just say so. \

"""

# Prepare messages to pass to model