This cookbook will show you how to integrate Ragie's Python SDK to retrieve relevant information from documents, pass it to an OpenAI model, and generate responses based on the retrieved data. In this example, the system interacts as an HR assistant, retrieving chunks from an employee handbook to answer employee queries.

Connect Your Data Source to Ragie

To begin, you’ll need to connect your data source (e.g., Google Drive, Notion) to Ragie. This allows Ragie to access and index your documents for retrieval.

Note: When the import mode is set to hi_res, images and tables will be extracted from the document. The fast mode only extracts text and can be up to 20 times faster than hi_res.

Getting Started

1. Create a requirements.txt file with the following content to install the necessary packages:

openai
ragie
python-dotenv

2. Install the packages by running:

pip install -r requirements.txt

3. Create a .env file to store your API keys securely.

‍# Instructions on how to get a Ragie API Key: https://secure.ragie.ai/api-keysRAGIE_API_KEY="YOUR_RAGIE_API_KEY"‍

# Instructions on how to get an OpenAI API Key: https://platform.openai.com/api-keysOPENAI_API_KEY="YOUR_OPENAI_API_KEY"

4. Add .env to your .gitignore file if you’re using version control to keep your API keys secure.

File Connection to Ragie and OpenAI

In your main file (index.py), start by importing the required libraries, loading environment variables, and authenticating with Ragie and OpenAI.

import os
from openai import OpenAI
from ragie import Ragie
from dotenv import load_dotenv

load_dotenv() 

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
ragie = Ragie(auth=os.getenv("RAGIE_API_KEY"))

Define a Class to Capture User Queries from the Command Line

class TextInput:
    def __init__(self, prompt="\n\nType your question here: "):
        self.prompt = prompt
        self.user_input = None
    
    def get_input(self):
        try:
            self.user_input = input(self.prompt)
            return self.user_input
        except EOFError:
            print("Error reading input")
            return None

Retrieving Chunks from Ragie

This function retrieves relevant chunks from Ragie based on the user’s query. Each chunk provides specific segments of information from the document to help answer the query accurately.

def retrieve_chunks(query: str):
    retrieval_res = ragie.retrievals.retrieve(request={
        "query": query,
        "rerank": True,
        "top_k": 6,
        "max_chunks_per_document": 0
    })
    if retrieval_res:
        return retrieval_res.scored_chunks
    return []

Processing the Retrieved Chunks

This function processes the retrieved chunks into a continuous string that can be passed to the model.

def process_chunks(chunks):
    return "".join(chunk.text for chunk in chunks)

Run the LLM from OpenAI with the Retrieved Chunks

Define a function to provide context from the retrieved chunks and generate a response using GPT-4 o mini. The context includes content from the employee handbook to create accurate answers.

def generate_response(chunk_text, query):
    system_prompt = f"""You are an internal AI assistant, "Ragie AI", designed to answer questions about the Company Handbook. Here are relevant chunks that you can use to respond to the user:
    
    ==== START RAG_CHUNKS ====
    {chunk_text}
    ==== END RAG_CHUNKS ====

    Answer succinctly and professionally in the tone of an HR manager. If no relevant information is found, politely suggest reaching out to HR.
    """
    
    completion = client.chat.completions.create(
        model="gpt-4-mini",
        stream=True,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": query}
        ]
    )
    print(completion.choices[0].message)

Running the Script

The main function captures user queries, retrieves chunks from Ragie, processes them, and generates a response using the LLM model we connected (GPT-4 o mini). 

if __name__ == "__main__":
    text_handler = TextInput()
    run = True
    while run:
        try:
            query = text_handler.get_input()
            if not query:
                break
            
            # Retrieve and print Ragie chunks
            chunks = retrieve_chunks(query)
            print("Retrieved chunks from Ragie:")
            for chunk in chunks:
                print(chunk.text + "\n")

            # Process chunks and run the model
            processed_chunks = process_chunks(chunks)
            generate_response(processed_chunks, query)
            
        except KeyboardInterrupt:
            run = False
            print("\nOperation cancelled by user.")

Testing the Code

1. Run the Script: Start the program in your terminal by running:

‍python index.py

2. Enter a Query: When prompted, type in a question you’d like the system to answer, such as:

What are the company’s remote work policies?

3. Generating a Response: The system will retrieve relevant content from your indexed documents in Ragie, process it through the OpenAI model, and display a response directly in the terminal.

Usage Example

Suppose a user asks: “What is the policy on remote work?” The system retrieves relevant document chunks from your company’s handbook file in Ragie and passes them to GPT-4 o mini to generate a response.

The output might look like:

Type your question here:

what is the policy on remote work?

At PostHog, we're a fully remote team, which allows us to hire talented individuals from all over the world. We believe that remote work encourages a deeper level of personal thought and writing things down, and it allows for uninterrupted work. Our remote work policy is designed to be flexible, so you can work from anywhere that suits you best. We also provide a budget to help you find a coworking space or cover the costs of coffee for those who prefer not to work at home.

Conclusion

This recipe shows how to use Ragie’s Python SDK with GPT-4 o mini by OpenAI to retrieve relevant content and generate accurate responses from your connected files on Ragie’s Dashboard. Try integrating this setup to add context-rich answers to your AI applications