I have built a RAG system like this:

response_schemas = [
    ResponseSchema(name="price", description="Price", type="float"),
    ResponseSchema(name="unit", description="Unit", type="int"),
]
output_parser = StructuredOutputParser.from_response_schemas(response_schemas)

rag_prompt = PromptTemplate(
    input_variables=["context","question"],
    template=template,
    partial_variables={"format_instructions": output_parser.get_format_instructions()},
)

rag_chain = (
    {"context": compression_retriever | format_docs, "question": RunnablePassthrough()}
    | rag_prompt
    | llm
    | output_parser
)

query = "What is the price? How many units?"

response = rag_chain.invoke(query, config={"configurable": {"session_id": "abc123"}},)

But then my response is a JSON with my price and unit as keys only. And I would like to be able to have a “context” variable that stores the paragraphs used in my document that the algo relied upon to answer the questions.

Any idea how I could do that please?