rag implementation

16 notes, read out of this brain and free to use. Each one was extracted from a source and is re-checked against its exam.

generateText return sources property

The sources property in generateText return value contains an array of Source objects from all steps. Each Source can have sourceType 'url' with properties: id (string), url (string), title (string, optional), and providerMetadata (SharedV2ProviderMetadata, optional).

RAG definition and purpose

RAG stands for retrieval augmented generation. RAG is the process of providing a Large Language Model with specific information relevant to the prompt. While LLMs are powerful, the information they can reason on is restricted to the data they were trained on. RAG solves this problem by fetching information relevant to the prompt and then passing that to the model as context.

Embeddings definition

Embeddings are a way to represent words, phrases, or images as vectors in a high-dimensional space. In this space, similar words are close to each other, and the distance between words can be used to measure their similarity. The process of calculating the similarity between two vectors is called cosine similarity, where a value of 1 would indicate high similarity and a value of -1 would indicate high opposition.

Chunking for embeddings

Chunking refers to the process of breaking down source material into smaller pieces. The larger the input to an embedding, the lower quality the embedding will be. A simple and common approach to chunking is separating written content by sentences. Once source material is appropriately chunked, each chunk can be embedded and then stored along with its embedding in a database.

Vector database support for RAG

Embeddings can be stored in any database that supports vectors. Postgres with the pgvector plugin is a recommended approach for storing embeddings.

RAG project stack

The recommended stack for building a RAG application includes: Next.js 14 (App Router), AI SDK, Vercel AI Gateway, Drizzle ORM, Postgres with pgvector, shadcn-ui and TailwindCSS for styling.

Embeddings table schema with pgvector

A table for storing embeddings should have the following columns: id (varchar, primaryKey with default nanoid), resourceId (varchar, foreign key to resources table with cascade delete), content (text, not null), and embedding (vector with dimensions 1536, not null). Include an HNSW or IVFFlat index on the embedding column for better performance using cosine distance operations.

Generate chunks function

A basic chunking function can split content by periods: const generateChunks = (input: string): string[] => { return input.trim().split('.').filter(i => i !== ''); }

Generate embeddings using AI SDK

Use the embedMany function from the AI SDK to generate embeddings for multiple chunks at once. The function takes a model (e.g., 'openai/text-embedding-ada-002') and an array of values to embed, returning embeddings that can be mapped with their corresponding content chunks for storage.

Complete embedMany usage example

import { embedMany } from 'ai'; const embeddingModel = 'openai/text-embedding-ada-002'; export const generateEmbeddings = async (value: string): Promise<Array<{ embedding: number[]; content: string }>> => { const chunks = generateChunks(value); const { embeddings } = await embedMany({ model: embeddingModel, values: chunks, }); return embeddings.map((e, i) => ({ content: chunks[i], embedding: e })); };

Single embedding generation for queries

Use the embed function from the AI SDK to generate a single embedding from a query string. This is used when searching for relevant content by comparing the query embedding against stored embeddings.

Complete single embed usage example

import { embed } from 'ai'; export const generateEmbedding = async (value: string): Promise<number[]> => { const input = value.replaceAll('\\n', ' '); const { embedding } = await embed({ model: 'openai/text-embedding-ada-002', value: input, }); return embedding; };

Vector similarity search with cosine distance

To find relevant content, embed the user query and search the database for similar items using cosine distance. Filter results by a similarity threshold (e.g., > 0.5) and order by similarity descending, then limit to a reasonable number of results (e.g., 4).

findRelevantContent implementation using Drizzle

export const findRelevantContent = async (userQuery: string) => { const userQueryEmbedded = await generateEmbedding(userQuery); const similarity = sql<number>`1 - (${cosineDistance(embeddings.embedding, userQueryEmbedded)})`; const similarGuides = await db .select({ name: embeddings.content, similarity }) .from(embeddings) .where(gt(similarity, 0.5)) .orderBy(t => desc(t.similarity)) .limit(4); return similarGuides; };

Creating resources with embeddings in Server Action

'use server'; import { createResource } from '@/lib/actions/resources'; import { generateEmbeddings } from '@/lib/ai/embedding'; import { embeddings as embeddingsTable } from '@/lib/db/schema/embeddings'; import { resources } from '@/lib/db/schema/resources'; import { db } from '@/lib/db'; export const createResource = async (input: NewResourceParams) => { const { content } = insertResourceSchema.parse(input); const [resource] = await db.insert(resources).values({ content }).returning(); const embeddings = await generateEmbeddings(content); await db.insert(embeddingsTable).values( embeddings.map(embedding => ({ resourceId: resource.id, ...embedding, })), ); return 'Resource successfully created and embedded.'; };

Drizzle ORM imports for embeddings query

To implement vector similarity search with Drizzle ORM, import: cosineDistance, desc, gt, sql from 'drizzle-orm' and embeddings from the schema.

Give your agent this brain