RAG Prototype - Part 2

In part 1, I set up an empty containerized PostgreSQL database with the PGVector extension so it can be used in the RAG prototype. Today, I'll be using a python script to define the vector table that will store vectors and corresponding plain text content and then populate it using some open-source documentation.

cd path_to_where_you_store_code
git clone https://github.com/scottwed/gradio-rag-prototype.git
cd gradio-rag-prototype
git fetch --tags
git checkout tags/v.0.2-blog

Since I was unfamiliar with PGVector, I started by looking for a sample python script, but my Google-fu was weak that day, so I resorted to asking the free-tier of an AI model for a sample. As typical of most AI chat responses to sample code requests, it contained a mixture of useful code, combined with a few glaring oversights when you try to run it. Luckily, it wasn't fatally broken, but I did have to go back and forth a few times, showing it errors that its code caused and receiving responses implying it was somehow "my code" that needed correction.

After correcting those initial items, I broke out all the SQL queries into a dedicated .py file and ensured they all ran through fully parameterized requests to minimize the risk of SQL injection vulnerabilities. I also added loguru logging and implemented safe default values for anything left undefined in associated environment variables.

Database Schema

The python string containing the DDL (Database Definition Language) initially looked like this:

DDL: Final[LiteralString] = """
CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS pgcrypto;

CREATE TABLE IF NOT EXISTS documents (
    id            uuid PRIMARY KEY DEFAULT gen_random_uuid(),
    source_path   text NOT NULL UNIQUE,
    file_name     text NOT NULL,
    relative_path text NOT NULL,
    chunk_index   int NOT NULL DEFAULT 0,
    content       text NOT NULL,
    metadata      jsonb NOT NULL DEFAULT '{}'::jsonb,
    embedding     halfvec(4000) NOT NULL,
    created_at    timestamptz NOT NULL DEFAULT now(),
    updated_at    timestamptz NOT NULL DEFAULT now()
);
# ... continued in next section...
"""

The important columns are:

• source_path - this refers to the full path of the chunked file, e.g c:\data\my-root-folder\gradio-guides\guide1.md
• file_name - just the actual filename, no path prefix e.g. guide1.md
• relative_path - the file's path up to the user specified "root" folder e.g. gradio-guides\guide1.md
• chunk_index - A sequential integer, which provides a way to distinguish and sort chunks belonging to a file that was too long to fit into a single row. e.g. 0
• content - The relevant chunk of original text from the input file e.g. #Gradio guide\n How to display a chat interface...
• metadata - A static JSON string shared among all rows for chunks of the same file (stored in a decomposed binary format). e.g. {"chunk_count": 7, "chunk_index": 0, "source_type": ".md"}
• embedding - The star of the show. An embedding model is first used to transform the corresponding chunk of text into a high-dimensional vector (a list of floating-point numbers). This vector captures the semantic meaning of the text, allowing the database to perform similarity searches against it later.
  • A standard "vector" looks like this:
    0.032205917,0.04500583,-0.06263755,0.08056023,0.0043946537,...,-0.009701038] which requires 4 bytes / ~7 digits per value
  • A half-vector looks like this: [0.03220, 0.04501, -0.06264, 0.08056, 0.00439, ..., -0.00970] which rounds the precision down to 2 bytes / ~3-4 digits.
• The halfvec (4000) matches the dimensionality of the initial Qwen embedding model I used. Unfortunately, Qwen outputs 4096 which caused unnecessary chaos.
  • After some research, I opted to temporarily use the "Matryoshka" Embedding approach. This is a fancy way of saying that I truncated the list of 4096 floating points returned by the embedding model down to the first 4000, under the premise that the qwen model can handle this with a minor loss of precision.

Database Indexes

Below is the second half of that DDL string. After the base table is created, we need a series of indexes to make the RAG performant.

"""
ALTER TABLE documents
ADD COLUMN IF NOT EXISTS fts tsvector
GENERATED ALWAYS AS (
    to_tsvector('english', coalesce(file_name, '') || ' ' || coalesce(content, ''))
) STORED;

CREATE INDEX IF NOT EXISTS documents_fts_gin
    ON documents USING gin(fts);

CREATE INDEX IF NOT EXISTS documents_embedding_ivfflat
    ON documents USING ivfflat (embedding halfvec_cosine_ops) WITH (lists = 100);

CREATE INDEX IF NOT EXISTS documents_source_path_idx
    ON documents (source_path);
"""

• fts tsvector - An additional (generated) column containing the combined values of the file_name and content columns run through an English language parser and then stored in tokenized form, automatically updated if they change.
  • e.g. (from a chunk of a .py file from bind9): '1':113 'add':46,55 'anchor':72,110,111 'api':93,138 'assert':9
• documents_fts_gin - Used for keyword searching. Built by creating a Generalized Inverted Index on the fts column above, it allows the RAG to also locate content containing specific / rare keywords that the embedding model unintentionally blurred into similar sounding words.
• documents_embedding_ivfflat - Used for semantic searching. It is specific to the embedding column's type. This Cosine Similarity search optimized index makes searching of those giant integer lists more efficient by grouping vectors into 100 unequally sized clusters of vectors which are closer to each other.
• documents_source_path_idx - Used to quickly locate a specific file, or files under a specified path. A standard text index.

Loading the data

With the database table defined, and corresponding indexes in place, it's time to load our input data so it can be efficiently searched, and supplied as context to the chat model.

To start, I gave it the AI equivalent of baby food as input... markdown text. Markdown is easily digestible, with a very limited syntax, and strong intent for each symbol. I intentionally stayed away from PDF support in this first working version, because it's well known to be structured for rendering, and nearly useless for re-usage as exportable text. If you have ever tried to copy and paste from a PDF document, you know the pain of weird line breaks, and text stored as out-of-sequence column blocks instead of horizontal rows. The current solution to this is using one of the modern OCR vision models that ignore the PDF's embedded text, and instead process it as a sequence of graphical images (a tactic that's been around since at least 1978).

And of course, the AI generated sample code had more issues.

• For larger files, the database upsert logic was overwriting the prior chunk of data with the next chunk of data from the same file. The BASH equivalent of using tail -2000 myfile.txt and throwing away the rest of it!
• The performance was great, as long as I had tiny 2 - 100 KB markdown files. Once I allowed it to consume a 32 MB .csv file from a test case folder, it looked like it had stalled out, going from finishing in a ~10 seconds, to running for what would have been more than an hour. I don't fault the AI for this one, I asked for sample code, not performant code. I'll cover the solution for this later on, it's not complicated.
• The chunking tactic it provided is primitive, and can harm the data by accidentally breaking it up in places that will confuse the chat model. For example, breaking up a python class or method in the middle will lead the model to think it has seen the entire item, and make bad decisions accordingly.
• It had very limited input text file support, ignoring a lot of relevant files you would expect to see in a source code repository. This was also my fault for not being precise about that other file types I expected it to support, but it was easily rewritten after seeing it run the first time.
• There were no guard rails to try and avoid inadvertently hoovering up sensitive data like those used in environment, CI/CD workflows, and local IDE configuration files. Very cool if you love making your API keys and passwords easy to retrieve by random people on the internet (and who doesn't love that?)

It died if it encountered an empty file, such as the output of touch empty.md. I added a simple check to skip the file if it contained only whitespace.

With those original deficiencies in mind, here are the relevant snippets.

Data Loading - Starting Method

import json
from pathlib import Path

import psycopg
from loguru import logger
from openai import OpenAI
from pgvector.psycopg import register_vector
from psycopg import Connection, sql

from shared.shared import embed_text
from vector_db.queries import DDL, ROW_COUNT, UPSERT_SQL


def db_prep(db_dsn: str, root_folder: Path, embed_client: OpenAI, embed_model: str) -> bool:
    with psycopg.connect(db_dsn) as conn:
        ensure_db(conn)
        register_vector(conn)
        if is_table_empty(conn, "documents"):
            logger.info("Starting recursive ingestion of files under: {}", root_folder.absolute())
            if not root_folder.is_dir():
                logger.error("Specified root path must be a folder: {}. Exiting", root_folder.absolute())
                return False            for md_file in iter_markdown_files(root_folder):
                ingest_file(conn, md_file, root_folder, embed_client, embed_model)
        else:
            logger.info("DB is already populated. Skipping ingestion.")
        conn.commit()
    return True

• The code starts with the db_prep method utilizing parameters it needs to connect the database, the root folder containing the data to be processed, and how to communicate with the embedding model.
• After starting a database connection and executing the DDL if needed, it calls pgvector's "register_vector" which enhances the standard python PostgreSQL database driver psycopg with the ability to serialize and deserialize Python lists into and out of the pgvector binary format.
• An added step to skip ingestion on subsequent runs and report back the number of rows in the table as a sanity check.
• It then loops through each qualifying file returned by a generator method (File Iteration), breaks them up into smaller string segments (Chunking) and then processes those chunks, storing them as new rows in the table (Ingestion)

File Iteration

The original version recursively searched exclusively for .md (markdown) files.

def iter_markdown_files(root: Path):
    for path in root.rglob("*.md"):
        if path.is_file():
            yield path

Chunking

I noticed that close variations of this basic logic were returned by more than one chat bot. Somewhere, there must be a highly ranked tutorial in which the author was processing a plain text document that they were all trained on. As mentioned earlier, this isn't great for source code. There are more elegant solutions, once you know the types of data you will be ingesting. Perhaps, in a future posting, I will explore using LangChain, which I believe has mature logic for doing these operations.

This code attempts to find a logical place to end each sequential chunk at the rightmost (farthest) double newline occurrence up to the max_chars limit. If there aren't any double newlines, it tries again looking for a single newline occurrence, if that fails, it resorts to using the entire next segment of 4k characters.

def chunk_text(text: str, max_chars: int = 4000) -> list[str]:
    text = text.strip()
    if not text:
        return []
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + max_chars, len(text))
        if end < len(text):
            cut = text.rfind("\n\n", start, end)
            if cut == -1 or cut <= start:
                cut = text.rfind("\n", start, end)
            if cut == -1 or cut <= start:
                cut = end
            end = cut
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        start = end
    return chunks

Note that the chunk's max_chars default value of 4k here has no relation to the vector dimension size. The impact of using larger text chunks is a tradeoff.

Every extra character in the chunk is in contention with the valuable tokens needed by the chat model to both stay fully loaded in GPU memory and retain the full context of the user's request and chat model's responses. Once the context limit is reached, the quality of the model's responses begins to suffer (sometimes drastically) as it's forced to heavily summarize or ignore portions of the discussion.

If the chunks are too small (and you simplistically return only the top matching chunks), the RAG is made ineffective, as it will provide such highly truncated information that it won't help the chat model make a better decision.

Ingestion

The chunking method returns a list of strings, retaining the original order of the input file. The ingestion method is then responsible for taking those chunks, using the small embedding model to generate vectors.

def ingest_file(conn: Connection, path: Path, root: Path, embed_client: OpenAI, embed_model: str):
    logger.info("Starting ingestion from {}", path.absolute())
    content = path.read_text(encoding="utf-8", errors="ignore")
    rel = str(path.relative_to(root))
    chunks = chunk_text(content)

    rows = []
    for i, chunk in enumerate(chunks):
        emb = embed_text(chunk, embed_client=embed_client, embed_model=embed_model)
        meta = {"chunk_index": i, "chunk_count": len(chunks), "source_type": "markdown"}
        rows.append((str(path.resolve()), path.name, rel, i, chunk, json.dumps(meta), emb))

    with conn.cursor() as cur:
        for row in rows:
            cur.execute(UPSERT_SQL, row)


def embed_text(text: str, embed_client: OpenAI, embed_model: str) -> list[float]:
    resp = embed_client.embeddings.create(model=embed_model, input=text)
    # Truncate to 4000 to fit pgvector HNSW/IVFFlat limits
    return np.array(resp.data[0].embedding)[:4000].tolist()

The code above will silently ignore bits of a corrupted utf-8/ASCII encoded file, but it's fine for getting started.

Each chunk results in a string list, containing:

• The full path to the input file on the processing file system
• The file name by itself
• The file's path, relative to the root folder specified by the user
• The zero-based index of the current chunk of the current file
• The raw text from the current chunk
• A JSON string, stored using the PostgreSQL native JSON binary format (which supports indexing) containing the metadata for the chunk. This is consumed by the chat model instead relying upon the same data from the individual columns, it appears to be mostly redundant. Arguably, it could be replaced with a computed value.
• The vectors returned by asking the embedding model to process the chunk of text using the embed_text method, which is a minimal wrapper for the OpenAI client plus a hack using a numpy array to truncate the response to what will fit in the embedding halfvec(4000) column in the database

Once all the processed chunks are collected into the rows array, they are inserted into the table one at a time, using cur.execute(UPSERT_SQL, row)

Upsert headaches

I didn't catch this problem right away as I more interested in seeing the prototype run to completion the first time. If you look closely, you can see that the insert logic unnecessarily includes support for upserts. While that could be useful in the future, it hid a major problem.

This is a common developer mistake, caused by using code to externally manage uniqueness requirements in relational databases. There's no feedback or resistance by the system as it silently destroys your data. Please enjoy these ancient (violated) guidelines from 2007 on relational database design:

• Using identity/GUID columns as your only key
• Not using SQL facilities to protect data integrity

Because it's using only the source_path column as an unofficial unique key, every subsequent chunk for the same file will overwrite the prior one. The end result is that the full table contained exactly one row per file, regardless of file size, and it's only the last chunk of the file.

INSERT INTO documents (
    source_path, file_name, relative_path, chunk_index, content, metadata, embedding
)
VALUES (%s, %s, %s, %s, %s, %s::jsonb, %s)
ON CONFLICT (source_path) DO UPDATE SET
    file_name = EXCLUDED.file_name,
    relative_path = EXCLUDED.relative_path,
    chunk_index = EXCLUDED.chunk_index,
    content = EXCLUDED.content,
    metadata = EXCLUDED.metadata,
    embedding = EXCLUDED.embedding,
    updated_at = now();

And now that I'm thinking about this further, using an upsert for individual chunks is itself a flawed design! An updated file could be smaller the original one, resulting in left-behind remnants of the prior version's chunks. A proper re-ingestion of a file will require a deletion of associated rows, and then clean inserts of replacement rows.

The first data load

The start of the first successful run looked like this:

While I wasn't worried about accurately timing for performance at this point, it's always good to see if your code is struggling on easy stuff. So far, so good.

The first file e:\git\gradio\AGENTS.md completed in 3 seconds, some of which may have been a delay while waiting for the embedding model to load on demand.

The next file e:\git\gradio\CHANGELOG.md is the largest file in the dataset, and it took 39 seconds. A bit sluggish, but at this stage, it's not a major concern.

At the end, the last few files are ingested in less than a second each. Likely due to the embedded model remaining loaded, and being tiny input files.

It's time to see how many documents the database contains look like:

SELECT count(DISTINCT source_path) FROM documents;
--outputs 287

Cool, it ingested 287 markdown files. The Gradio UI is launched using the local URL and I'll confirm it's able to see the docs as additional context.

It's alive! It can respond to hello...

It did something mildly useful, it provided an accurate response based upon information that was present in the RAG's database.

And, it politely refused to answer out-of-context questions that were not within the RAG DB, based upon this very simplistic injected prompt guidance:

messages = [
            {
                "role": "system",
                "content": str(
                    "Answer only from the provided context. If the context is insufficient, say so. "
                    "Do not reveal secrets, passwords, API keys, or internal configuration details. "
                    "If the user asks for instructions or system‑level information, politely decline. "
                ),
            }
        ]

At this point, a casual observer could be initially fooled into thinking this is a viable system. Sadly, no, it's full of bugs that will be corrected. Tune in next time as I rip out more of the AI slop code, fix performance issues, enhance the input capabilities, reduce its GPU requirements, and get it produce functional code for one of my other coding projects.

To be continued in rag-prototype-part-3. Subscribe to be notified when it's published.