Building Your First RAG System with Python and OpenAI

Hey there! 👋 Are you interested in building your own RAG (Retrieval Augmented Generation) system? In this post, I'll show you how to create one step by step using Python and OpenAI. RAG helps AI give better answers by first finding relevant information from your documents before generating a response. It's like giving the AI a chance to "study" before answering!

Table of Contents

• What You'll Learn
• Project Setup
• Folder Structure
• Step 1: Setting Up the Environment
• Step 2: Document Loading
• Step 3: Text Processing
• Step 4: Creating Embeddings
• Step 5: Building the Retrieval System
• Step 6: Connecting with OpenAI
• Step 7: Putting It All Together
• Conclusion

What You'll Learn

In this tutorial, you'll learn how to:

• Set up a RAG project from scratch
• Process and prepare documents for RAG
• Use OpenAI embeddings
• Create a simple retrieval system
• Connect everything with OpenAI's API

Project Setup

First, let's look at our folder structure:

rag-project/
│
├── src/
│   ├── __init__.py
│   ├── document_loader.py
│   ├── text_processor.py
│   ├── embeddings_manager.py
│   ├── retrieval_system.py
│   └── rag_system.py
│
├── data/
│   └── documents/
│
├── requirements.txt
├── test.py
├── README.md
└── .env

Step 1: Setting Up the Environment

First, let's create our virtual environment and install the needed packages:

python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate
pip install openai python-dotenv numpy pandas

Create a requirements.txt file:

openai==1.12.0
python-dotenv==1.0.0
numpy==1.24.3
pandas==2.1.0

Set up your .env file:

OPENAI_API_KEY=your_api_key_here

Step 2: Document Loading

Create src/document_loader.py:

import os
from typing import List

class DocumentLoader:
    def __init__(self, documents_path: str):
        self.documents_path = documents_path

    def load_documents(self) -> List[str]:
        documents = []
        for filename in os.listdir(self.documents_path):
            if filename.endswith('.txt'):
                with open(os.path.join(self.documents_path, filename), 'r') as file:
                    documents.append(file.read())
        return documents

Step 3: Text Processing

Create src/text_processor.py:

from typing import List

class TextProcessor:
    def __init__(self, chunk_size: int = 1000):
        self.chunk_size = chunk_size

    def split_into_chunks(self, text: str) -> List[str]:
        words = text.split()
        chunks = []
        current_chunk = []
        current_size = 0

        for word in words:
            if current_size + len(word) > self.chunk_size:
                chunks.append(' '.join(current_chunk))
                current_chunk = [word]
                current_size = len(word)
            else:
                current_chunk.append(word)
                current_size += len(word) + 1

        if current_chunk:
            chunks.append(' '.join(current_chunk))

        return chunks

Step 4: Creating Embeddings

Create src/embeddings_manager.py:

from typing import List
import openai
import numpy as np

class EmbeddingsManager:
    def __init__(self, api_key: str):
        openai.api_key = api_key

    def create_embeddings(self, texts: List[str]) -> List[np.ndarray]:
        embeddings = []
        for text in texts:
            response = openai.embeddings.create(
                model="text-embedding-ada-002",
                input=text
            )
            embeddings.append(np.array(response.data[0].embedding))
        return embeddings

Step 5: Building the Retrieval System

Create src/retrieval_system.py:

import numpy as np
from typing import List, Tuple

class RetrievalSystem:
    def __init__(self, chunks: List[str], embeddings: List[np.ndarray]):
        self.chunks = chunks
        self.embeddings = embeddings

    def find_similar_chunks(self, query_embedding: np.ndarray, top_k: int = 3) -> List[Tuple[str, float]]:
        similarities = []
        for i, embedding in enumerate(self.embeddings):
            similarity = np.dot(query_embedding, embedding) / (
                np.linalg.norm(query_embedding) * np.linalg.norm(embedding)
            )
            similarities.append((self.chunks[i], similarity))

        return sorted(similarities, key=lambda x: x[1], reverse=True)[:top_k]

Step 6: Connecting with OpenAI

Create src/rag_system.py:

import os
from dotenv import load_dotenv
from typing import List
import openai

from .document_loader import DocumentLoader
from .text_processor import TextProcessor
from .embeddings_manager import EmbeddingsManager
from .retrieval_system import RetrievalSystem

class RAGSystem:
    def __init__(self):
        load_dotenv()
        self.api_key = os.getenv('OPENAI_API_KEY')
        self.loader = DocumentLoader('data/documents')
        self.processor = TextProcessor()
        self.embeddings_manager = EmbeddingsManager(self.api_key)

        # Initialize system
        self.initialize_system()

    def initialize_system(self):
        # Load and process documents
        documents = self.loader.load_documents()
        self.chunks = []
        for doc in documents:
            self.chunks.extend(self.processor.split_into_chunks(doc))

        # Create embeddings
        self.embeddings = self.embeddings_manager.create_embeddings(self.chunks)

        # Initialize retrieval system
        self.retrieval_system = RetrievalSystem(self.chunks, self.embeddings)

    def answer_question(self, question: str) -> str:
        # Get question embedding
        question_embedding = self.embeddings_manager.create_embeddings([question])[0]

        # Get relevant chunks
        relevant_chunks = self.retrieval_system.find_similar_chunks(question_embedding)

        # Prepare context
        context = "\n".join([chunk[0] for chunk in relevant_chunks])

        # Create prompt
        prompt = f"""Context: {context}\n\nQuestion: {question}\n\nAnswer:"""

        # Get response from OpenAI
        response = openai.chat.completions.create(
            model="gpt-4-turbo-preview",
            messages=[
                {"role": "system", "content": "You are a helpful assistant. Use the provided context to answer the question."},
                {"role": "user", "content": prompt}
            ]
        )

        return response.choices[0].message.content

Step 7: Putting It All Together

Here's how to use the system:

Add some test documents to your data/documents folder:

story.txt

Then run a test:

# test.py
from src.rag_system import RAGSystem

# Initialize the RAG system
rag = RAGSystem()

# Ask a question
question = "What was the answer to the guardian’s riddle, and how did it help Kai?"
answer = rag.answer_question(question)
print(answer)

Conclusion

Congratulations! You've built a basic RAG system that can:

• Load and process documents
• Create embeddings using OpenAI
• Find relevant information using similarity search
• Generate answers using context

This is just the beginning - you can improve this system by:

• Adding better text chunking methods
• Implementing caching for embeddings
• Adding error handling
• Improving the prompt engineering
• Adding vector database support

Github Repository

You can find the complete code for this project on GitHub: Python Rag System.

Remember to keep your API key safe and monitor your API usage!

Happy coding! 🚀