Embeddings and pgvector RAG
Ground AI answers in your own data: embeddings stored in Postgres via pgvector.
The chat surface you built with tools can answer “what’s the total on invoice #INV-203?”: a tool fetches that one row and hands it back. But “what’s our refund policy?” or “summarize the themes across my four thousand support tickets” has no single row to fetch. The answer lives in a body of text the model never trained on, your internal handbook or a customer’s ticket archive, and that text is far too big to paste into every prompt.
The move is to pull only the relevant pieces of that text and feed them to the model as context for this one question.
That is retrieval-augmented generation (RAG) : not a new stack, just the same streamText route, Postgres, and session.orgId discipline running a query that enriches the system prompt before the model runs.
The first thing to learn is when to reach for it, because the common failure is building the whole apparatus for text that would have fit in one system prompt.
When to use RAG
Section titled “When to use RAG”Two conditions must both hold before RAG pays for itself.
First, the corpus is internal, so the model can’t have trained on it: your company handbook, a customer’s private knowledge base, a user’s uploaded documents, last week’s call transcripts. Ask any model which HTTP status code means “not found” and it answers correctly on its own, so retrieval over public, well-trodden information is wasted work.
Second, the corpus is either too large for the context window, like a three-hundred-page handbook you’d otherwise resend on every message, or small but changing too fast to freeze, like a pricing table that updates weekly.
Below that threshold, do the boring thing: paste the text into the system prompt and build nothing else. A two-page returns policy that rarely changes is a few hundred tokens of prompt; a vector store for it adds a table, an indexing job, and a query path to solve what the prompt already solved. The threshold is a cost line: stuffing the same N tokens into every request bills you N tokens per call forever, while retrieval pays the embedding cost once at indexing, then sends a bounded handful of passages per question no matter how large the corpus grows. A long-context model can swallow a mid-size corpus whole, but per-call cost and recall that degrades inside long prompts still favor retrieval at scale.
The decision tool below runs these conditions one question at a time.
The model already knows this. No retrieval, no tool, no corpus: a plain free-text answer is correct and cheapest.
Below the threshold, so a vector store is over-engineering. Drop the text into the system prompt and ship. Revisit only when the corpus outgrows the prompt or starts changing under you.
Exact lookup, not semantic search. A tool like getInvoiceById from earlier in this chapter fetches one row deterministically by its key. Embeddings find text that means something similar, the wrong instrument for “give me record X.”
A broad, fuzzy question over a large or fast-changing internal corpus. The case retrieval was built for, and the rest of this lesson builds it.
RAG earns its place only at that last leaf, not for a small corpus (paste it) or an exact-key lookup (use a tool).
What an embedding is
Section titled “What an embedding is”To find the relevant passages, you first need a way to measure meaning.
An embedding is a fixed-length array of floating-point numbers, a vector, produced by an embedding model . Unlike a chat model, which writes text back, an embedding model hands back numbers, the same numbers every time for the same input.
One property makes those numbers useful: text with similar meaning maps to nearby vectors. “Invoice,” “bill,” and “receipt” land close together; “office cat” lands far away. The model learned this during training, so “nearby” tracks meaning rather than spelling: “unpaid invoices” sits near “outstanding bills” even though they share no words.
The standard measure of “nearby” is cosine similarity , where higher means more similar in meaning.
Postgres stores the inverse, cosine distance (1 − similarity), so closer vectors have a smaller distance, and your query orders by distance ascending to get the most similar passages first.
Semantic search is three steps: embed every passage in your corpus once, embed the user’s question when it arrives, then find the corpus vectors nearest to the question’s vector. Those nearest passages are your relevant context.
The vector’s fixed length is its dimensions , and that count sizes your storage.
OpenAI’s text-embedding-3-small, the model this lesson uses, produces 1536-dimensional vectors, a number that has to match the database column exactly.
Embedding with embed and embedMany
Section titled “Embedding with embed and embedMany”The AI SDK exposes embeddings through two functions imported from 'ai', the same operation at one scale and at many.
embed({ model, value }) takes one string and returns one vector plus a usage count.
Use it at query time to embed the single question the user just asked.
embedMany({ model, values }) takes an array of strings and returns an array of vectors in the same order, plus aggregate usage.
Use it at index time, handing it the whole corpus in one call; when the array exceeds the provider’s per-call limit, it splits the work into batches for you. That batching is separate from chunking the documents, a decision you make yourself and one we return to shortly.
Both functions take a model handle from the same registry as your chat models, where the embedding model gets its own named export.
export const embeddingModel = 'openai/text-embedding-3-small';This is the same 'provider/model' AI Gateway string your chat handles use: no provider package imported, no factory called.
We pick text-embedding-3-small over the text-embedding-3-large from the Swappable models & the Gateway lesson because its 1536 dimensions stay under pgvector’s HNSW index limit of 2000; 3-large’s 3072 would overflow it.
That line hides a trap. Chat handles swap freely: point the export at a different provider tomorrow and yesterday’s conversations still make sense, but embedding handles don’t. A different embedding model produces vectors in a different space, incomparable to the old ones, so swapping forces you to re-embed the entire corpus. Treat it as a far stickier choice than the chat model.
import { embed } from 'ai';import { embeddingModel } from '@/lib/llm/models';
const { embedding } = await embed({ model: embeddingModel, value: 'When do I get a refund?',});One string in, one vector out. The query-time call: embed the user’s question to compare it against the stored corpus. embedding is a single number[].
import { embedMany } from 'ai';import { embeddingModel } from '@/lib/llm/models';
const { embeddings } = await embedMany({ model: embeddingModel, values: chunks,});Many strings in, many vectors out, in input order. The index-time call: chunks is an array of passages, and embeddings[i] is the vector for chunks[i]. The SDK auto-batches under the provider’s per-call limit.
The batch call belongs to an indexing job, not a request handler. Run it on document upload or as a one-time backfill, in a plain async function or a script. Embedding ten thousand passages takes time and money, so you pay that once, offline, and the live chat never touches it.
Storing vectors: pgvector and the Drizzle vector column
Section titled “Storing vectors: pgvector and the Drizzle vector column”You have vectors; where do they live? For a web app, the Postgres you already run.
pgvector is the extension that makes this work, and using it is a pure operations win: one fewer service to run, credential to rotate, and thing to break. Reach for a dedicated vector database such as Pinecone, Upstash Vector, or Qdrant only when the corpus outgrows pgvector, roughly tens of millions of vectors, or when you need a managed service. For a typical feature, that day rarely comes.
The schema is a single table, each column earning its place.
export const documentChunks = pgTable( 'document_chunks', { id: uuid().primaryKey().$defaultFn(() => uuidv7()), orgId: uuid().notNull(), documentId: uuid() .notNull() .references(() => documents.id, { onDelete: 'cascade' }), content: text().notNull(), embedding: vector({ dimensions: 1536 }).notNull(), embeddingModel: text().notNull(), }, (t) => [ index('idx_document_chunks_embedding').using( 'hnsw', t.embedding.op('vector_cosine_ops'), ), ],);A normal Drizzle pgTable. The UUIDv7 primary key is the app’s usual convention, and snake_case SQL names like org_id come from the client’s casing: 'snake_case'.
export const documentChunks = pgTable( 'document_chunks', { id: uuid().primaryKey().$defaultFn(() => uuidv7()), orgId: uuid().notNull(), documentId: uuid() .notNull() .references(() => documents.id, { onDelete: 'cascade' }), content: text().notNull(), embedding: vector({ dimensions: 1536 }).notNull(), embeddingModel: text().notNull(), }, (t) => [ index('idx_document_chunks_embedding').using( 'hnsw', t.embedding.op('vector_cosine_ops'), ), ],);The tenancy column: every chunk belongs to one org. The retrieval query filters on it, and getting that filter right is the most important thing in this lesson, so it gets its own section below.
export const documentChunks = pgTable( 'document_chunks', { id: uuid().primaryKey().$defaultFn(() => uuidv7()), orgId: uuid().notNull(), documentId: uuid() .notNull() .references(() => documents.id, { onDelete: 'cascade' }), content: text().notNull(), embedding: vector({ dimensions: 1536 }).notNull(), embeddingModel: text().notNull(), }, (t) => [ index('idx_document_chunks_embedding').using( 'hnsw', t.embedding.op('vector_cosine_ops'), ), ],);The raw passage text. The vector finds the row, but this is the text the query returns and the model reads.
export const documentChunks = pgTable( 'document_chunks', { id: uuid().primaryKey().$defaultFn(() => uuidv7()), orgId: uuid().notNull(), documentId: uuid() .notNull() .references(() => documents.id, { onDelete: 'cascade' }), content: text().notNull(), embedding: vector({ dimensions: 1536 }).notNull(), embeddingModel: text().notNull(), }, (t) => [ index('idx_document_chunks_embedding').using( 'hnsw', t.embedding.op('vector_cosine_ops'), ), ],);The pgvector column. vector({ dimensions: 1536 }) matches text-embedding-3-small’s output exactly; if the two numbers disagree, the insert fails.
export const documentChunks = pgTable( 'document_chunks', { id: uuid().primaryKey().$defaultFn(() => uuidv7()), orgId: uuid().notNull(), documentId: uuid() .notNull() .references(() => documents.id, { onDelete: 'cascade' }), content: text().notNull(), embedding: vector({ dimensions: 1536 }).notNull(), embeddingModel: text().notNull(), }, (t) => [ index('idx_document_chunks_embedding').using( 'hnsw', t.embedding.op('vector_cosine_ops'), ), ],);Records which model produced this vector. Redundant with one model in use, but it makes a future re-index possible by letting you find exactly the rows that need new vectors.
export const documentChunks = pgTable( 'document_chunks', { id: uuid().primaryKey().$defaultFn(() => uuidv7()), orgId: uuid().notNull(), documentId: uuid() .notNull() .references(() => documents.id, { onDelete: 'cascade' }), content: text().notNull(), embedding: vector({ dimensions: 1536 }).notNull(), embeddingModel: text().notNull(), }, (t) => [ index('idx_document_chunks_embedding').using( 'hnsw', t.embedding.op('vector_cosine_ops'), ), ],);The HNSW index, which makes similarity search fast once the table is large. The defaults fit this course’s scale, so we name it and move on.
Two columns carry the migration’s whole risk.
The embedding column’s 1536 must match the model’s output dimension, and the HNSW index must use the vector_cosine_ops operator class for cosine distance: a wrong dimension breaks the insert, a wrong operator class fails the migration. Everything else is ordinary Drizzle.
The two-phase pipeline: index, then query
Section titled “The two-phase pipeline: index, then query”RAG runs in two phases at completely different times.
The index phase is offline and batched: take a document, split it into passages, embed all of them with embedMany, and insert one row per passage into documentChunks, each carrying its content, embedding, orgId, and embeddingModel.
It runs on upload or as a backfill, always before any user asks anything.
The query phase is online and per-request: embed the question, run a similarity query for the nearest passages filtered to the user’s org, stitch their content into a context string, drop that into the system prompt, and call streamText.
This runs every time a user sends a message.
The two phases meet at one place, the documentChunks table: the index phase writes it, the query phase reads it.
Index phase, offline. A document comes in and a chunker splits it into coherent passages, a few sentences to a paragraph each.
embedMany turns every passage into a vector in one batched call.
One row per passage is inserted into documentChunks, carrying its content, embedding, orgId, and embeddingModel. Written once, offline.
Query phase, now per request. The user’s question arrives and embed turns it into a single query vector.
A similarity query reads documentChunks and returns the top-K nearest passages, scoped to this org. This table is the seam between the phases: written by the index phase, read here.
The retrieved passages enrich the system prompt, streamText runs, and the answer is grounded in your corpus.
Chunking: a decision, not a default
Section titled “Chunking: a decision, not a default”The SDK ships no chunker, deliberately: how you split text is a problem-domain decision, not a setting.
Prose splits on paragraphs, code on fixed token windows, transcripts on utterances or speaker turns.
Reach for a library like @langchain/text-splitters, or hand-roll a splitter for your shape of text.
Two guardrails pull in opposite directions. Chunk too large, a whole page each, and a retrieved passage is mostly irrelevant text around the one sentence that mattered, so the answer drifts. Chunk too small, a single sentence each, and the passage loses the context it needs to stand on its own. The sweet spot is a coherent passage, a few sentences to a paragraph, with a little overlap between neighbors so an idea split across a boundary survives in both. Aim for “a human could answer the question from this passage alone,” then adjust.
The index-phase code
Section titled “The index-phase code”A plain async function you’d call on document upload, stripped to its three steps.
export const indexDocument = async (doc: Document) => { const chunks = chunkDocument(doc.content);
const { embeddings } = await embedMany({ model: embeddingModel, values: chunks, });
await db.insert(documentChunks).values( chunks.map((content, i) => ({ orgId: doc.orgId, documentId: doc.id, content, embedding: embeddings[i], embeddingModel: 'openai/text-embedding-3-small', })), );};Real ingestion adds deduplication, batching across many documents, and error handling around each step. These reuse the background-work patterns you already have, so we don’t rebuild them here. Deduplication is a trap worth its own mention later.
The query-phase code
Section titled “The query-phase code”The read side: embed the question, then find the nearest passages, scoped to the org.
export const findRelevantChunks = async (question: string, orgId: string) => { const { embedding } = await embed({ model: embeddingModel, value: question, });
const similarity = sql<number>`1 - (${cosineDistance(documentChunks.embedding, embedding)})`;
return db .select({ content: documentChunks.content, similarity }) .from(documentChunks) .where(eq(documentChunks.orgId, orgId)) .orderBy(desc(similarity)) .limit(5);};Embed the question into a query vector with embed. This is the one string we’re searching with.
export const findRelevantChunks = async (question: string, orgId: string) => { const { embedding } = await embed({ model: embeddingModel, value: question, });
const similarity = sql<number>`1 - (${cosineDistance(documentChunks.embedding, embedding)})`;
return db .select({ content: documentChunks.content, similarity }) .from(documentChunks) .where(eq(documentChunks.orgId, orgId)) .orderBy(desc(similarity)) .limit(5);};cosineDistance measures distance between the stored column and the query vector, and 1 - distance turns it into similarity. Recall the inversion: distance is lower-is-closer, similarity is higher-is-closer. We compute similarity so the ordering reads naturally.
export const findRelevantChunks = async (question: string, orgId: string) => { const { embedding } = await embed({ model: embeddingModel, value: question, });
const similarity = sql<number>`1 - (${cosineDistance(documentChunks.embedding, embedding)})`;
return db .select({ content: documentChunks.content, similarity }) .from(documentChunks) .where(eq(documentChunks.orgId, orgId)) .orderBy(desc(similarity)) .limit(5);};The tenancy filter. eq(documentChunks.orgId, orgId) restricts the search to this org’s chunks and nothing else. This one line is the difference between a working feature and a cross-tenant leak, so it gets its own section below.
export const findRelevantChunks = async (question: string, orgId: string) => { const { embedding } = await embed({ model: embeddingModel, value: question, });
const similarity = sql<number>`1 - (${cosineDistance(documentChunks.embedding, embedding)})`;
return db .select({ content: documentChunks.content, similarity }) .from(documentChunks) .where(eq(documentChunks.orgId, orgId)) .orderBy(desc(similarity)) .limit(5);};Order by similarity descending and take the top 5. K is small on purpose: a large K drags irrelevant passages into the prompt and reintroduces the cost problem retrieval exists to solve.
Pre-retrieval in the chat route handler
Section titled “Pre-retrieval in the chat route handler”Now wire retrieval into the chat handler you already know, in a few lines that run before streamText: embed the user’s latest question, run the org-scoped similarity query, and fold the retrieved passages into the system prompt.
export const POST = authedRoute('member', chatSchema, async ({ messages }, { session }) => { const lastMessage = messages.at(-1); const question = lastMessage?.parts .filter((part) => part.type === 'text') .map((part) => part.text) .join(' ');
const chunks = await findRelevantChunks(question ?? '', session.orgId); const relevantContext = chunks.map((chunk) => chunk.content).join('\n\n');
const result = streamText({ model: smartModel, system: `You answer questions about the company handbook.Answer only from the context below; if it is not there, say you don't know.
Relevant context:${relevantContext}`, messages: convertToModelMessages(messages), stopWhen: stepCountIs(5), maxOutputTokens: 1024, });
return result.toUIMessageStreamResponse();});Same wrapper and signature as the chapter’s handler: the validated body arrives as { messages }, and session carries the session.orgId that scopes retrieval.
export const POST = authedRoute('member', chatSchema, async ({ messages }, { session }) => { const lastMessage = messages.at(-1); const question = lastMessage?.parts .filter((part) => part.type === 'text') .map((part) => part.text) .join(' ');
const chunks = await findRelevantChunks(question ?? '', session.orgId); const relevantContext = chunks.map((chunk) => chunk.content).join('\n\n');
const result = streamText({ model: smartModel, system: `You answer questions about the company handbook.Answer only from the context below; if it is not there, say you don't know.
Relevant context:${relevantContext}`, messages: convertToModelMessages(messages), stopWhen: stepCountIs(5), maxOutputTokens: 1024, });
return result.toUIMessageStreamResponse();});Pull the latest question and fetch its nearest chunks, org-scoped. Running this before the model call is what makes it pre-retrieval.
export const POST = authedRoute('member', chatSchema, async ({ messages }, { session }) => { const lastMessage = messages.at(-1); const question = lastMessage?.parts .filter((part) => part.type === 'text') .map((part) => part.text) .join(' ');
const chunks = await findRelevantChunks(question ?? '', session.orgId); const relevantContext = chunks.map((chunk) => chunk.content).join('\n\n');
const result = streamText({ model: smartModel, system: `You answer questions about the company handbook.Answer only from the context below; if it is not there, say you don't know.
Relevant context:${relevantContext}`, messages: convertToModelMessages(messages), stopWhen: stepCountIs(5), maxOutputTokens: 1024, });
return result.toUIMessageStreamResponse();});Retrieved passages go into the system prompt, the trusted controller, never into messages. The retrieval was authorized server-side under session.orgId, so this text is trusted; the raw user turn is not.
export const POST = authedRoute('member', chatSchema, async ({ messages }, { session }) => { const lastMessage = messages.at(-1); const question = lastMessage?.parts .filter((part) => part.type === 'text') .map((part) => part.text) .join(' ');
const chunks = await findRelevantChunks(question ?? '', session.orgId); const relevantContext = chunks.map((chunk) => chunk.content).join('\n\n');
const result = streamText({ model: smartModel, system: `You answer questions about the company handbook.Answer only from the context below; if it is not there, say you don't know.
Relevant context:${relevantContext}`, messages: convertToModelMessages(messages), stopWhen: stepCountIs(5), maxOutputTokens: 1024, });
return result.toUIMessageStreamResponse();});Unchanged from the chapter’s handler: the step and output caps are your cost guardrails with or without retrieval.
export const POST = authedRoute('member', chatSchema, async ({ messages }, { session }) => { const lastMessage = messages.at(-1); const question = lastMessage?.parts .filter((part) => part.type === 'text') .map((part) => part.text) .join(' ');
const chunks = await findRelevantChunks(question ?? '', session.orgId); const relevantContext = chunks.map((chunk) => chunk.content).join('\n\n');
const result = streamText({ model: smartModel, system: `You answer questions about the company handbook.Answer only from the context below; if it is not there, say you don't know.
Relevant context:${relevantContext}`, messages: convertToModelMessages(messages), stopWhen: stepCountIs(5), maxOutputTokens: 1024, });
return result.toUIMessageStreamResponse();});The same response the client already speaks; retrieval is invisible to it, so nothing on the client changes.
This is the pre-retrieval pattern: retrieve on every turn, before the model runs. It’s the simplest shape and the right default when nearly every message needs the corpus.
The retrieved context rides in the system prompt, not messages, because the retrieval was authorized server-side under session.orgId: it belongs on the trusted side, enriching your instructions rather than handing control to whatever sits in the corpus.
Retrieval grounds the answer but does not validate it. RAG cuts the odds the model contradicts your corpus, but a wrong handbook still yields a confidently wrong answer, so treat retrieved text as the model’s source, not an oracle.
Pre-retrieval vs retrieval as a tool
Section titled “Pre-retrieval vs retrieval as a tool”Pre-retrieval always fires, but on many surfaces most questions never touch the corpus: a general assistant fields “what’s the weather” far more often than “what does the handbook say about X,” and embedding and querying on every turn is wasted work.
A second architecture fixes that: make retrieval a tool.
Define a searchKnowledgeBase tool whose execute runs the embed-and-query you just wrote, and let the model decide when to call it through the same agentic loop as every other tool.
The two split cleanly:
- Pre-retrieval is the default when every turn likely needs the corpus, as in a docs Q&A bot or a “chat with this handbook” surface. One query per turn, deterministic, least machinery.
- Retrieval as a tool fits mixed surfaces, where some questions need the corpus and many don’t. It skips the embedding and the query on turns that don’t need them, folding retrieval into the loop alongside the model’s other tools.
As a tool, the orgId filter lives inside execute, alongside every other tool’s org-scope, and the result feeds back through the same stopWhen loop.
const chunks = await findRelevantChunks(question, session.orgId);const relevantContext = chunks.map((chunk) => chunk.content).join('\n\n');
const result = streamText({ model: smartModel, system: `Answer from the context below.\n\nRelevant context:\n${relevantContext}`, messages: convertToModelMessages(messages), stopWhen: stepCountIs(5),});Use when every turn needs the corpus. Retrieval runs unconditionally before the model, the condensed handler from above.
const searchKnowledgeBase = tool({ description: 'Search the company handbook for relevant passages.', inputSchema: z.object({ query: z.string() }), execute: async ({ query }) => findRelevantChunks(query, session.orgId),});
const result = streamText({ model: smartModel, messages: convertToModelMessages(messages), tools: { searchKnowledgeBase }, stopWhen: stepCountIs(5),});Use on mixed surfaces, where many turns don’t need the corpus. The model calls the tool only when a question needs it; the orgId filter lives inside execute, with session.orgId in scope from the handler.
Don’t run both on the same surface: pre-injecting context and offering a searchKnowledgeBase tool invites double-retrieval, or leaves the model unsure whether it already has what it needs.
One retrieval strategy per surface.
Test the decision against these scenarios.
Your support chat needs to answer questions from a two-page returns policy that’s barely been touched in a year. Which approach earns its weight?
documentChunks and retrieve on every turn before the model runs.getPolicy tool that fetches the row by id.searchKnowledgeBase tool and let the model decide when to consult it.A user pastes the string INV-203 and wants that one invoice’s line items back. Which approach fits?
You’re building a “chat with our 300-page engineering handbook” bot where essentially every message is a question about the handbook. Which approach fits?
lookupSection tool that fetches a page by its heading.A general support assistant fields all kinds of requests — weather, summaries, small talk — and only now and then needs to quote the handbook. Which approach fits?
Authorizing retrieval: the multi-tenant rule
Section titled “Authorizing retrieval: the multi-tenant rule”Every documentChunks row carries an orgId, and every retrieval query must filter by session.orgId. No exceptions.
You already org-scope every query, but here an unscoped one is worse than an ordinary leak. A normal leak hands another org’s rows to your code; this one hands them to the model, which quotes one tenant’s private handbook, pricing, or tickets back as fluent prose inside another tenant’s chat, read as their own answer. It is the worst shape a cross-tenant bug can take.
The filter travels with the query: the route handler for pre-retrieval, or inside execute for retrieval-as-a-tool, where every tool puts its org-scope.
Keeping the corpus fresh
Section titled “Keeping the corpus fresh”A corpus is not a one-time upload. Three operational realities follow the feature for life, each handled with tools you already have.
Embedding models change.
Move from text-embedding-3-small to a newer model and you must re-embed the entire corpus: the new vectors live in a different space, incomparable to the old ones and possibly a different dimension.
The embeddingModel column on each row makes this survivable: you query the rows still tagged with the old model and re-embed them in batches, rolling the corpus forward without a big-bang outage.
Documents change.
When a source document is edited, re-chunk and re-embed its chunks; when it’s deleted, the foreign-key cascade on documentId removes them for free.
Duplicates poison retrieval. A duplicated chunk fills several top-K slots with the same text, spending your small K budget and biasing the answer toward it. Deduplicate at insert time.
A re-index is just a background job, and you already have the toolkit to run one durably.
External resources
Section titled “External resources”The canonical references for the exact APIs this lesson used, plus a tool to make the abstract part concrete.
The reference for embed, embedMany, and the embedding model handle.
The pgvector column helper, the HNSW index, and the cosineDistance query shape.
Rotate and zoom a real embedding space in 3D, and watch nearest neighbours cluster by meaning, live.
The extension itself: the vector column, the cosine operator class, and the HNSW index tuning knobs.