from openai import OpenAI
import chromadb
import numpy as np
from typing import Optional

client = OpenAI()

# Initialize vector store
chroma_client = chromadb.Client()
collection = chroma_client.create_collection("knowledge_base")

def add_documents(documents: list[dict]):
    """Add documents to the vector store."""
    texts = [doc["text"] for doc in documents]

    # Embed all documents in one API call (more efficient)
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )

    embeddings = [item.embedding for item in response.data]
    ids = [doc["id"] for doc in documents]
    metadatas = [doc.get("metadata", {}) for doc in documents]

    collection.add(
        embeddings=embeddings,
        documents=texts,
        metadatas=metadatas,
        ids=ids
    )

def retrieve(query: str, top_k: int = 3) -> list[str]:
    """Find most relevant documents for a query."""
    query_embedding = client.embeddings.create(
        model="text-embedding-3-small",
        input=query
    ).data[0].embedding

    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k
    )
    return results["documents"][0]

def rag_answer(question: str, top_k: int = 3) -> str:
    """RAG: retrieve context, then generate answer."""
    # Step 1: retrieve relevant documents
    context_docs = retrieve(question, top_k=top_k)
    context = "\n\n".join([f"[{i+1}] {doc}" for i, doc in enumerate(context_docs)])

    # Step 2: generate answer with context
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {
                "role": "system",
                "content": """Answer the question based ONLY on the provided context documents.
If the context doesn't contain the answer, say: "I don't have that information."
Always cite which document number you used: [1], [2], etc."""
            },
            {
                "role": "user",
                "content": f"Context:\n{context}\n\nQuestion: {question}"
            }
        ],
        temperature=0
    )
    return response.choices[0].message.content

# Usage
add_documents([
    {"id": "1", "text": "Our refund policy allows returns within 30 days.", "metadata": {"source": "policy"}},
    {"id": "2", "text": "Shipping takes 5-7 business days standard.", "metadata": {"source": "shipping"}},
    {"id": "3", "text": "Products have 1-year warranty.", "metadata": {"source": "warranty"}},
])

answer = rag_answer("What's the return window?")
print(answer)
RAG — Retrieval Augmented Generation

RAG — Retrieval Augmented Generation

RAG — Retrieval Augmented Generation
