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 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

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 = 3) -> 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]

def ask(question: str) -> str:
    conn = psycopg2.connect(
        host="localhost", dbname="ragdb", user="raguser", password="ragpass"
    )
    embedding = get_embedding(question)
    chunks = retrieve(conn, embedding)
    conn.close()

    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: {question}
Answer:"""

    response = litellm.completion(
        model="anthropic/claude-3-haiku-20240307",
        messages=[{"role": "user", "content": prompt}],
        temperature=0
    )
    return response.choices[0].message.content.strip()

if __name__ == "__main__":
    questions = [
        "How many vacation days do I get per year?",
        "What is the meal expense limit per person?",
        "What tools does the company use for project tracking?",
    ]
    for q in questions:
        print(f"\nQ: {q}")
        print(f"A: {ask(q)}")

Q: ¿Cuántos días de vacaciones tengo por año?
A: According to TechCorp's Vacation Policy, you receive 15 paid vacation days per year.

Q: ¿Cuál es el límite de gastos en restaurantes?
A: The meal expense limit is $50 per person for meals.

Q: ¿Qué herramientas usa la empresa para gestión de proyectos?
A: TechCorp uses Linear for project tracking.