rag.py

import os
import chromadb
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

client = OpenAI(api_key=os.environ["OPEN_AI_API_KEY"])

# Vector database
chroma = chromadb.PersistentClient(path="./chroma")

collection = chroma.get_or_create_collection(
    name="company_docs"
)


# -------------------------
# INDEXING
# -------------------------

documents = [
    {
        "id": "refund",
        "text": """
        Our refund policy allows customers to request a refund within
        30 days of receiving their order. Shipping fees are not refundable.
        Refunds are normally processed within 5 business days.
        """
    },
    {
        "id": "shipping",
        "text": """
        Standard shipping usually takes 3 to 5 business days.
        Express shipping takes 1 to 2 business days.
        Express shipping is available only for orders above $20.
        """
    }
]


def embed(text):
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )

    return response.data[0].embedding


# Add documents to vector DB
for document in documents:

    collection.upsert(
        ids=[document["id"]],
        documents=[document["text"]],
        embeddings=[embed(document["text"])]
    )


# -------------------------
# RETRIEVAL
# -------------------------

def retrieve(question, k=2):

    query_embedding = embed(question)

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=k
    )

    return results["documents"][0]


# -------------------------
# GENERATION
# -------------------------

def answer(question):

    chunks = retrieve(question)

    context = "\n\n".join(chunks)

    prompt = f"""
You are a helpful customer support assistant.

Answer the question using ONLY the information
provided in the context.

If the answer is not in the context, say:
"I don't know based on the available information."

Context:
{context}

Question:
{question}
"""

    response = client.chat.completions.create(
        model="gpt-5",
        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ]
    )

    return response.choices[0].message.content


# -------------------------
# CHAT
# -------------------------

while True:

    question = input("\nYou: ")

    if question.lower() in ["exit", "quit"]:
        break

    print("\nAI:", answer(question))