Simple RAG Workflow

# Complete minimal RAG system using only Python and the OpenAI API
# Install: pip install openai numpy

import os
import math
import json
from openai import OpenAI

class SimpleRAG:
    """
    A minimal RAG (Retrieval-Augmented Generation) system.
    
    This is a teaching implementation. For production, use a proper
    vector database like Pinecone, Weaviate, or ChromaDB.
    """

    def __init__(self):
        api_key = os.environ.get("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("Set OPENAI_API_KEY environment variable first.")
        self.client = OpenAI(api_key=api_key)
        self.documents = []  # List of {text, embedding} dicts

    def add_document(self, text):
        """Add a document chunk to the knowledge base."""
        # Create embedding for this text
        response = self.client.embeddings.create(
            model="text-embedding-3-small",
            input=text
        )
        embedding = response.data[0].embedding

        self.documents.append({
            "text": text,
            "embedding": embedding
        })
        print(f"Added document: '{text[:60]}...'")

    def search(self, query, top_k=3):
        """Find the most relevant documents for a query."""
        # Get query embedding
        response = self.client.embeddings.create(
            model="text-embedding-3-small",
            input=query
        )
        query_embedding = response.data[0].embedding

        # Calculate similarity with all documents
        similarities = []
        for doc in self.documents:
            sim = self._cosine_similarity(query_embedding, doc["embedding"])
            similarities.append((sim, doc["text"]))

        # Sort by similarity and return top_k
        similarities.sort(reverse=True)
        return [text for _, text in similarities[:top_k]]

    def ask(self, question):
        """Answer a question using retrieved documents."""
        # Retrieve relevant chunks
        relevant_chunks = self.search(question, top_k=3)

        # Build augmented prompt
        context = "\n\n".join([f"- {chunk}" for chunk in relevant_chunks])

        messages = [
            {
                "role": "system",
                "content": "Answer based only on the provided context. If the answer is not in the context, say so."
            },
            {
                "role": "user",
                "content": f"Context:\n{context}\n\nQuestion: {question}"
            }
        ]

        # Generate answer
        response = self.client.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            max_tokens=300
        )

        return response.choices[0].message.content

    def _cosine_similarity(self, vec_a, vec_b):
        dot = sum(a * b for a, b in zip(vec_a, vec_b))
        mag_a = math.sqrt(sum(a ** 2 for a in vec_a))
        mag_b = math.sqrt(sum(b ** 2 for b in vec_b))
        return dot / (mag_a * mag_b) if mag_a and mag_b else 0

# Example usage
rag = SimpleRAG()

# Add company policy documents
rag.add_document("Our refund policy allows returns within 30 days of purchase with original receipt.")
rag.add_document("Customer support is available Monday to Friday, 9 AM to 6 PM EST.")
rag.add_document("Standard shipping takes 5-7 business days. Express shipping takes 2-3 days.")
rag.add_document("Products are covered by a 1-year manufacturer warranty against defects.")

# Ask questions
questions = [
    "How long do I have to return something?",
    "When can I contact support?",
]

for q in questions:
    print(f"\nQ: {q}")
    answer = rag.ask(q)
    print(f"A: {answer}")