Guardrails Output Parsing

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

%pip install llama-index-llms-openai
%pip install llama-index-output-parsers-guardrails

!pip install guardrails-ai

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.core import VectorStoreIndex, SimpleDirectoryReader
from IPython.display import Markdown, display

import os

os.environ["OPENAI_API_KEY"] = "sk-..."

# 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.guardrails import GuardrailsOutputParser

Define custom QA and Refine Prompts

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}
"""

from llama_index.llms.openai import OpenAI

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

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

# attach to an llm object
llm = OpenAI(output_parser=output_parser)

from llama_index.core.prompts.default_prompts import (
    DEFAULT_TEXT_QA_PROMPT_TMPL,
)

# take a look at the new QA template!
fmt_qa_tmpl = output_parser.format(DEFAULT_TEXT_QA_PROMPT_TMPL)
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(
    llm=llm,
)
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"
          }
        ]
      }
    }
  }
}