Guardrails Output Parsing

If you’re opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

Download Data

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

Load documents, build the VectorStoreIndex

import logging
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

from llama_index import VectorStoreIndex, SimpleDirectoryReader
from IPython.display import Markdown, display

import openai

openai.api_key = "<YOUR_OPENAI_API_KEY>"

# load documents
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()

index = VectorStoreIndex.from_documents(documents, chunk_size=512)

INFO:llama_index.token_counter.token_counter:> [build_index_from_documents] Total LLM token usage: 0 tokens
> [build_index_from_documents] Total LLM token usage: 0 tokens
INFO:llama_index.token_counter.token_counter:> [build_index_from_documents] Total embedding token usage: 18579 tokens
> [build_index_from_documents] Total embedding token usage: 18579 tokens

Define Query + Guardrails Spec

from llama_index.output_parsers import GuardrailsOutputParser
from llama_index.llm_predictor import StructuredLLMPredictor

llm_predictor = StructuredLLMPredictor()

Define custom QA and Refine Prompts

from llama_index.prompts import PromptTemplate
from llama_index.prompts.default_prompts import (
    DEFAULT_TEXT_QA_PROMPT_TMPL,
    DEFAULT_REFINE_PROMPT_TMPL,
)

Define Guardrails Spec

# You can either define a RailSpec and initialise a Guard object from_rail_string()
# OR define Pydantic classes and initialise a Guard object from_pydantic()
# For more info: https://docs.guardrailsai.com/defining_guards/pydantic/
# Guardrails recommends Pydantic

from pydantic import BaseModel, Field
from typing import List
import guardrails as gd


class Point(BaseModel):
    # In all the fields below, you can define validators as well
    # Left out for brevity
    explanation: str = Field()
    explanation2: str = Field()
    explanation3: str = Field()


class BulletPoints(BaseModel):
    points: List[Point] = Field(
        description="Bullet points regarding events in the author's life."
    )


# Define the prompt
prompt = """
Query string here.

${gr.xml_prefix_prompt}

${output_schema}

${gr.json_suffix_prompt_v2_wo_none}
"""

# Create a guard object
guard = gd.Guard.from_pydantic(output_class=BulletPoints, prompt=prompt)

# Create output parse object
output_parser = GuardrailsOutputParser(guard, llm=llm_predictor.llm)

# NOTE: we use the same output parser for both prompts, though you can choose to use different parsers
# NOTE: here we add formatting instructions to the prompts.

fmt_qa_tmpl = output_parser.format(DEFAULT_TEXT_QA_PROMPT_TMPL)
fmt_refine_tmpl = output_parser.format(DEFAULT_REFINE_PROMPT_TMPL)

qa_prompt = PromptTemplate(fmt_qa_tmpl, output_parser=output_parser)
refine_prompt = PromptTemplate(fmt_refine_tmpl, output_parser=output_parser)

# take a look at the new QA template!
print(fmt_qa_tmpl)

Context information is below.
---------------------
{context_str}
---------------------
Given the context information and not prior knowledge, answer the query.
Query: {query_str}
Answer: 


Given below is XML that describes the information to extract from this document and the tags to extract it into.


<output>
    <list name="points" description="Bullet points regarding events in the author's life.">
        <object>
            <string name="explanation"/>
            <string name="explanation2"/>
            <string name="explanation3"/>
        </object>
    </list>
</output>



ONLY return a valid JSON object (no other text is necessary). The JSON MUST conform to the XML format, including any types and format requests e.g. requests for lists, objects and specific types. Be correct and concise.

Query Index

query_engine = index.as_query_engine(
    text_qa_template=qa_prompt,
    refine_template=refine_prompt,
    llm_predictor=llm_predictor,
)
response = query_engine.query(
    "What are the three items the author did growing up?",
)

INFO:llama_index.token_counter.token_counter:> [query] Total LLM token usage: 754 tokens
> [query] Total LLM token usage: 754 tokens
INFO:llama_index.token_counter.token_counter:> [query] Total embedding token usage: 11 tokens
> [query] Total embedding token usage: 11 tokens

print(response)

{
  "output": {
    "list": {
      "name": "points",
      "description": "Bullet points regarding events in the author's life.",
      "object": {
        "string": [
          {
            "name": "explanation",
            "content": "Writing short stories"
          },
          {
            "name": "explanation2",
            "content": "Programming on the IBM 1401"
          },
          {
            "name": "explanation3",
            "content": "Building a microcomputer"
          }
        ]
      }
    }
  }
}