Step 5: Ingest — Chunk, Embed, Insert

# Option A — use anything you have
cp ~/some-doc.txt data/
 
# Option B — grab the Wikipedia article on RAG (pure illustration)
curl -s "https://en.wikipedia.org/api/rest_v1/page/plain/Retrieval-augmented_generation" \
  > data/rag-wikipedia.txt

# src/chunker.py
"""Split text into overlapping chunks of roughly N characters."""
 
import re
from typing import Iterable
 
# Roughly: 400 tokens ≈ 1600 characters for English prose.
TARGET_CHARS = 1600
OVERLAP_CHARS = 200
 
 
def _split_paragraphs(text: str) -> list[str]:
    return [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
 
 
def chunk_text(text: str) -> Iterable[str]:
    paragraphs = _split_paragraphs(text)
 
    buf = ""
    for p in paragraphs:
        if len(buf) + len(p) + 2 <= TARGET_CHARS:
            buf = f"{buf}\n\n{p}".strip()
            continue
 
        if buf:
            yield buf
 
        # Start the new buffer with the tail of the previous one for overlap.
        tail = buf[-OVERLAP_CHARS:] if buf else ""
        buf = f"{tail}\n\n{p}".strip() if tail else p
 
    if buf:
        yield buf

# src/embeddings.py
"""Wrap Vertex AI text-embedding-005."""
 
from functools import lru_cache
from vertexai.language_models import TextEmbeddingInput, TextEmbeddingModel
import vertexai
 
from .config import load_config
 
EMBED_MODEL = "text-embedding-005"
EMBED_DIM = 768
 
 
@lru_cache(maxsize=1)
def _model() -> TextEmbeddingModel:
    cfg = load_config()
    vertexai.init(project=cfg.project, location=cfg.region)
    return TextEmbeddingModel.from_pretrained(EMBED_MODEL)
 
 
def embed(texts: list[str], task: str = "RETRIEVAL_DOCUMENT") -> list[list[float]]:
    """Embed a batch of texts. Use task='RETRIEVAL_QUERY' for question embedding."""
    inputs = [TextEmbeddingInput(text=t, task_type=task) for t in texts]
    response = _model().get_embeddings(inputs)
    return [r.values for r in response]

# src/ingest.py
"""Walk data/, chunk, embed, insert."""
 
from pathlib import Path
 
from .chunker import chunk_text
from .db import get_conn
from .embeddings import embed
 
DATA_DIR = Path("data")
BATCH_SIZE = 50  # how many chunks to embed per Vertex AI call
 
 
def gather_chunks() -> list[tuple[str, int, str]]:
    """Return a flat list of (source, chunk_index, content)."""
    out: list[tuple[str, int, str]] = []
    for path in sorted(DATA_DIR.glob("*.txt")):
        text = path.read_text(encoding="utf-8")
        for i, content in enumerate(chunk_text(text)):
            out.append((path.name, i, content))
    return out
 
 
def insert_chunks(rows: list[tuple[str, int, str, list[float]]]) -> None:
    """Insert (source, chunk_index, content, embedding) rows."""
    with get_conn() as conn:
        cur = conn.cursor()
        cur.executemany(
            """
            INSERT INTO chunks (source, chunk_index, content, embedding)
            VALUES (%s, %s, %s, %s)
            """,
            # pgvector accepts the standard string format '[1.0,2.0,...]'
            [(s, i, c, _vec_to_pg(v)) for s, i, c, v in rows],
        )
        conn.commit()
 
 
def _vec_to_pg(vec: list[float]) -> str:
    return "[" + ",".join(f"{x:.7f}" for x in vec) + "]"
 
 
def main() -> None:
    chunks = gather_chunks()
    print(f"Found {len(chunks)} chunks across {len(set(c[0] for c in chunks))} files.")
 
    rows: list[tuple[str, int, str, list[float]]] = []
    for start in range(0, len(chunks), BATCH_SIZE):
        batch = chunks[start : start + BATCH_SIZE]
        texts = [c[2] for c in batch]
        vectors = embed(texts, task="RETRIEVAL_DOCUMENT")
        rows.extend((s, i, c, v) for (s, i, c), v in zip(batch, vectors))
        print(f"Embedded {min(start + BATCH_SIZE, len(chunks))}/{len(chunks)}")
 
    insert_chunks(rows)
    print(f"Inserted {len(rows)} chunks into Postgres.")
 
 
if __name__ == "__main__":
    main()

uv run python -m src.ingest

Found 17 chunks across 2 files.
Embedded 17/17
Inserted 17 chunks into Postgres.

gcloud sql connect rag-db --user=$USER_EMAIL --database=rag --quiet

-- Count rows
SELECT COUNT(*), COUNT(DISTINCT source) FROM chunks;
 
-- Peek at one
SELECT source, chunk_index, LEFT(content, 80) AS preview
FROM chunks
LIMIT 3;
 
-- Spot-check the embedding dim
SELECT vector_dims(embedding) FROM chunks LIMIT 1;