RAG: cómo darle memoria a un LLM con tus propios documentos

services:
  postgres:
    image: pgvector/pgvector:pg16
    environment:
      POSTGRES_DB: ragdb
      POSTGRES_USER: raguser
      POSTGRES_PASSWORD: ragpass
    ports:
      - "5432:5432"
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata:

docker compose up -d

pip install fastapi uvicorn litellm psycopg2-binary

export OPENAI_API_KEY=sk-...
export ANTHROPIC_API_KEY=sk-ant-...

import psycopg2
from psycopg2.extras import execute_values
import litellm

# Fictional TechCorp internal documents
DOCUMENTS = [
    {
        "title": "Vacation Policy",
        "content": """TechCorp grants 15 paid vacation days per year to all full-time employees.
Days accrue at 1.25 days per month. Unused days roll over up to 30 days maximum.
Requests must be submitted at least 2 weeks in advance.
Requests during December require 4 weeks notice.
Vacation pay is calculated based on the employee base salary."""
    },
    {
        "title": "Onboarding Guide",
        "content": """Welcome to TechCorp! Your first week:
Day 1: IT setup, equipment pickup, and access provisioning.
Day 2: HR orientation and benefits enrollment.
Day 3-4: Team introduction and codebase walkthrough.
Day 5: First task assignment and buddy system pairing.
Tools: Slack for communication, Linear for project tracking,
GitHub for version control, Notion for documentation.
Contact IT helpdesk at it@techcorp.com for technical issues."""
    },
    {
        "title": "Expense Policy",
        "content": """TechCorp reimburses pre-approved business expenses.
Meals: up to $50 per person, $150 per team event.
Travel: economy class flights, hotels up to $200 per night.
Equipment: pre-approval required for purchases over $500.
Coworking: up to $30 per day when working remotely.
Submit expense reports within 30 days via Expensify.
Attach all receipts. Expenses over $100 require manager approval.
Reimbursements processed within 2 business weeks."""
    }
]

def chunk_text(text: str, chunk_size: int = 250) -> list[str]:
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    chunks, current, length = [], [], 0
    for para in paragraphs:
        if length + len(para) > chunk_size and current:
            chunks.append(" ".join(current))
            current, length = [para], len(para)
        else:
            current.append(para)
            length += len(para)
    if current:
        chunks.append(" ".join(current))
    return chunks

def get_embedding(text: str) -> list[float]:
    # Claude (Anthropic) has no native embeddings endpoint.
    # LiteLLM lets us use OpenAI embeddings here and Claude for generation.
    response = litellm.embedding(model="text-embedding-3-small", input=text)
    return response.data[0].embedding

def setup_db(conn) -> None:
    with conn.cursor() as cur:
        cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
        cur.execute("""
            CREATE TABLE IF NOT EXISTS documents (
                id SERIAL PRIMARY KEY,
                title TEXT NOT NULL,
                content TEXT NOT NULL,
                embedding vector(1536)
            );
        """)
        cur.execute("TRUNCATE documents;")
    conn.commit()

def main() -> None:
    conn = psycopg2.connect(
        host="localhost", dbname="ragdb", user="raguser", password="ragpass"
    )
    print("Setting up database...")
    setup_db(conn)

    rows = []
    for doc in DOCUMENTS:
        chunks = chunk_text(doc["content"])
        print(f"Processing '{doc['title']}': {len(chunks)} chunks")
        for chunk in chunks:
            rows.append((doc["title"], chunk, get_embedding(chunk)))

    with conn.cursor() as cur:
        execute_values(
            cur,
            "INSERT INTO documents (title, content, embedding) VALUES %s",
            rows,
            template="(%s, %s, %s::vector)"
        )
    conn.commit()
    conn.close()
    print(f"Done. Inserted {len(rows)} chunks.")

if __name__ == "__main__":
    main()

python ingest.py

Setting up database...
Processing 'Vacation Policy': 2 chunks
Processing 'Onboarding Guide': 3 chunks
Processing 'Expense Policy': 3 chunks
Done. Inserted 8 chunks.

import psycopg2
import litellm
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

app = FastAPI(title="TechCorp RAG API")

DB_CONFIG = {
    "host": "localhost",
    "dbname": "ragdb",
    "user": "raguser",
    "password": "ragpass",
}


class AskRequest(BaseModel):
    question: str
    top_k: int = 3


class SourceChunk(BaseModel):
    title: str
    content: str
    similarity: float


class AskResponse(BaseModel):
    answer: str
    sources: list[SourceChunk]


def get_embedding(text: str) -> list[float]:
    response = litellm.embedding(model="text-embedding-3-small", input=text)
    return response.data[0].embedding


def retrieve(conn, embedding: list[float], top_k: int) -> list[dict]:
    with conn.cursor() as cur:
        # <=> is pgvector's cosine distance operator (lower = more similar)
        # 1 - distance = cosine similarity
        cur.execute("""
            SELECT title, content, 1 - (embedding <=> %s::vector) AS similarity
            FROM documents
            ORDER BY embedding <=> %s::vector
            LIMIT %s;
        """, (embedding, embedding, top_k))
        rows = cur.fetchall()
    return [{"title": r[0], "content": r[1], "similarity": round(r[2], 3)} for r in rows]


@app.get("/health")
def health() -> dict:
    return {"status": "ok"}


@app.post("/ask", response_model=AskResponse)
def ask(req: AskRequest) -> AskResponse:
    try:
        conn = psycopg2.connect(**DB_CONFIG)
        embedding = get_embedding(req.question)
        chunks = retrieve(conn, embedding, req.top_k)
        conn.close()
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))

    context = "\n\n".join(f"[{c['title']}]\n{c['content']}" for c in chunks)
    prompt = f"""You are an internal assistant for TechCorp.
Answer ONLY using the context below. If the information is not there, say so clearly.

Context:
{context}

Question: {req.question}
Answer:"""

    response = litellm.completion(
        model="anthropic/claude-3-haiku-20240307",
        messages=[{"role": "user", "content": prompt}],
        temperature=0,
    )
    return AskResponse(
        answer=response.choices[0].message.content.strip(),
        sources=[SourceChunk(**c) for c in chunks],
    )

uvicorn query:app --reload

curl -X POST http://localhost:8000/ask \
  -H "Content-Type: application/json" \
  -d '{"question": "How many vacation days do I get per year?"}'

{
  "answer": "According to TechCorp's Vacation Policy, you receive 15 paid vacation days per year.",
  "sources": [
    {
      "title": "Vacation Policy",
      "content": "TechCorp grants 15 paid vacation days per year to all full-time employees. Days accrue at 1.25 days per month...",
      "similarity": 0.847
    },
    {
      "title": "Vacation Policy",
      "content": "Requests must be submitted at least 2 weeks in advance...",
      "similarity": 0.612
    }
  ]
}