RAG Engine Overview

​

Links

• Playground: Try it in the browser, with all the API parameters.
• API documentation: Learn how to use our RAG Engine APIs:
  • Contextual Answers using the RAG Engine
  • Manage your document library

​

Library specifics

​

Features

• Seamless integration between retrieval and generation: RAG Engine automatically integrates with many of our task-specific models, so you can surface search results or provide a grounded answer to a query based on your organizational data – all within a single API call. You can also connect RAG Engine with a foundation model like Jurassic – and use Semantic Search results within a prompt.
• Support for several file formats: The RAG Engine supports several different file types including PDF. The parser understands complex layouts, including tables.
• Built-in data source integration: You can integrate your organization’s data sources, such as Google Drive, Amazon S3, and others, to automatically sync documents with RAG Engine. To enable data source integration, contact us.
• Secure and appropriate access to documents: The RAG Engine adheres to organizational user and group permissions with a comprehensive approach that addresses the intricate needs of document management in modern, data-intensive environments.

​

Current limitations

• OCR functionality is not available at this time. Only PDF documents that contain a text layer are supported.
• Some PDFs may take longer to process, depending on their content.

​

How the RAG engine works

• Adocument library manager: Upload files into the library and parse the content into text and content. The library can read several different formats, including PDFs, and can understand complex structures such as tables.

1. Indexing: When a document is uploaded, the engine parses it into segments by topic. Adjacent sections of the doc that are about the same topic are considered to be the same segment (up until a maximum size). Each segment has an embedding value that expresses the meaning of the segment in numerical form. The engine stores a list of segments and their embedding values.
   • Pro tips: A document is actually segmented and rated in two different ways: using sparse evaluation, which is a simple list of keywords in the segment, and dense evaluation, which is a vector of 1024 values that give a more complex evaluation of what the segment is about (monkeys, pets, care and feeding, etc) as calculated by the embedding engine.
2. Retrieval: The engine extracts the user’s question, calculates an embedding value for it that represents the meaning of the question, then looks through the index for segments that seem to be about the same topic. Matching segments are put into a candidate pool for the next step.
   • Pro tip: There is a similarity threshold between the user’s question and the segment. The retrieval engine looks for segments within this threshold.
3. Question answering: All segments in the candidate pool are combined into a prompt, along with the question, and sent to an LLM to generate an answer. (The prompt is basically, “here is some information and a question; generate an answer to the question using only the information given here; if you don’t know the answer, say ‘I don’t know’“)

​

Using the RAG engine

​

Step 1: Upload your files

from ai21 import AI21Client

client = AI21Client(
    # This is the default and can be omitted
    api_key=os.environ.get("AI21_API_KEY"),
)

def upload_rag_file(path, labels=None, path_meta=None, url=None):
    file_id = client.library.files.create(
        file_path=path,
        labels=labels,
        path=path_meta,
        public_url=url
    )
    print(file_id)

# Note that each file name (ignoring the path) must be unique.
# path and label parameters are used for filtering, as described below.
upload_rag_file("/Users/dafna/uk_employee_policy.pdf",
        path_meta='hr/UK',labels=['hr', 'policy'])
upload_rag_file("/Users/dafna/us_employee_policy.md",
        path_meta='hr/US', labels=['hr', 'policy'])
upload_rag_file("/Users/dafna/it_security.html",
        labels=['security'])

​

Labeling and filtering requests

​

Step 2: Ask a question

client = AI21Client()

def query_library(query, labels=None, path=None):
    response = client.library.answer.create(
        question=query,
        path=path,
        labels=labels
    )
    print(response.answer)

# Filter the question to documents with the case-sensitive label "hr"
query_library("How many days can I work from home in the US?",labels=["hr"])

"In the US, employees can work from home up to two days per week."

xxx

if response.answer_in_context:
    # Do something
else:
    # Do something else

​

Improving your results

​

Working with PDFs

​

Benchmarking your answers