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

class DocumentLoader(Protocol):
    def supports(self, path: Path) -> bool: ...
    def load(self, path: Path) -> Iterator[Document]: ...


DEFAULT_LOADERS: list[DocumentLoader] = [PdfLoader(), MarkdownLoader(), TextLoader()]


def loader_for(path: Path, loaders: list[DocumentLoader]) -> DocumentLoader:
    for loader in loaders:
        if loader.supports(path):
            return loader
    raise UnsupportedDocument(path)

class PdfLoader:
    """One document per page, so an answer can cite the page it came from."""

    def supports(self, path: Path) -> bool:
        return path.suffix.lower() == ".pdf"

    def load(self, path: Path) -> Iterator[Document]:
        reader = PdfReader(str(path))
        for number, page in enumerate(reader.pages, start=1):
            text = (page.extract_text() or "").strip()
            if not text:
                continue  # Image-only pages carry nothing to embed.
            yield Document(text=text, source=path.name, title=path.stem, page=number)

limit = max_tokens - overlap
pieces: list[str] = []
for paragraph in (p.strip() for p in _PARAGRAPH.split(text) if p.strip()):
    pieces.extend(_fit(paragraph, encoding, limit))

chunks: list[str] = []
current: list[int] = []
for piece in pieces:
    tokens = encoding.encode(piece if not current else " " + piece)
    if current and len(current) + len(tokens) > max_tokens:
        chunks.append(encoding.decode(current).strip())
        current = current[-overlap:] if overlap else []   # carry the overlap
        tokens = encoding.encode(piece)
    current.extend(tokens)

CREATE TABLE IF NOT EXISTS chunks (
    id BIGSERIAL PRIMARY KEY,
    source TEXT NOT NULL,
    title TEXT NOT NULL,
    page INT,
    content TEXT NOT NULL,
    embedding vector(1536) NOT NULL
);

CREATE INDEX IF NOT EXISTS chunks_embedding_idx
    ON chunks USING hnsw (embedding vector_cosine_ops);

SELECT title, source, page, content,
       1 - (embedding <=> %s) AS similarity
  FROM chunks
 ORDER BY embedding <=> %s
 LIMIT %s;

{
  "answer": "According to the TechCorp Vacation Policy (page 1), TechCorp grants **15 paid vacation days per year** to all full-time employees.",
  "sources": [
    {
      "title": "vacation-policy",
      "source": "vacation-policy.pdf",
      "page": 1,
      "content": "TechCorp Vacation Policy\nTechCorp grants 15 paid vacation days per year to all full-time employees...",
      "similarity": 0.559
    },
    {
      "title": "vacation-policy",
      "source": "vacation-policy.pdf",
      "page": 2,
      "content": "Requesting Time Off\nRequests must be submitted at least 2 weeks in advance...",
      "similarity": 0.474
    },
    {
      "title": "expense-policy",
      "source": "expense-policy.pdf",
      "page": 1,
      "content": "TechCorp Expense Policy\nTechCorp reimburses pre-approved business expenses only...",
      "similarity": 0.249
    },
    {
      "title": "expense-policy",
      "source": "expense-policy.pdf",
      "page": 2,
      "content": "Submitting Expenses\nEquipment purchases above 500 dollars require pre-approval...",
      "similarity": 0.237
    }
  ]
}