Vector Search via Embedding

Database entries can be too many and too large to fit into a prompt. The AI tools need a more compact form, so the app does two things:

The @operation.tool({ hidden: true }) decorator keeps the “get all” procedures out of the derived AI tools.
Each user and task gets a vector embedding when it is created or updated. An OpenAI embeddings model creates it, and pgvector stores it in the embedding column of each model, declared as Unsupported("vector(1536)") in the Prisma schema. See the Database article.

The EmbeddingService class creates the embeddings and runs vector search with the pgvector extension in Postgres. Its two main methods:

generateEntityEmbedding – creates an embedding from an entity’s string fields and writes the vector to its database row.
vectorSearch – searches the stored embeddings for a query string and returns the most similar entries.

Both methods work with any entity type, user and task alike.

src/modules/embedding/embedding-service.ts


import { openai } from '@ai-sdk/openai';
import { Prisma } from '@prisma/client';
import type { EntityType } from '@schemas/index';
import type { TaskType } from '@schemas/models/Task.schema';
import type { UserType } from '@schemas/models/User.schema';
import { embed } from 'ai';
import { capitalize, omit } from 'lodash';
import { BASE_KEYS } from '@/constants';
import DatabaseService from '../database/database-service';
 
export default class EmbeddingService {
  static async generateEmbedding(value: string): Promise<number[]> {
    const { embedding } = await embed({
      model: openai.embeddingModel('text-embedding-3-small'),
      value,
    });
 
    return embedding;
  }
 
  static generateEntityEmbedding = async (
    entityType: EntityType,
    entityId: UserType['id'] | TaskType['id'],
  ) => {
    const entity = await DatabaseService.prisma[
      entityType as 'user'
    ].findUnique({
      where: { id: entityId },
    });
    const capitalizedEntityType = capitalize(entityType);
    if (!entity) throw new Error(`${capitalizedEntityType} not found`);
 
    const embedding = await this.generateEmbedding(
      Object.values(omit(entity, BASE_KEYS))
        .filter((v) => typeof v === 'string')
        .join(' ')
        .trim()
        .toLowerCase(),
    );
 
    await DatabaseService.prisma.$executeRawUnsafe(
      `
    UPDATE "${capitalizedEntityType}" 
    SET embedding = $1::vector
    WHERE id = $2
    `,
      `[${embedding.join(',')}]`,
      entityId,
    );
 
    return embedding;
  };
 
  static async vectorSearch<T>(
    entityType: EntityType,
    query: string,
    limit: number = 10,
    similarityThreshold: number = 0.4,
  ) {
    const queryEmbedding = await EmbeddingService.generateEmbedding(
      query.trim().toLowerCase(),
    );
    const capitalizedEntityType = capitalize(entityType);
 
    // find similar vectors and return entity IDs
    const vectorResults = await DatabaseService.prisma.$queryRaw<
      { id: string; similarity: number }[]
    >`
    SELECT
      id,
      1 - (embedding <=> ${`[${queryEmbedding.join(',')}]`}::vector) as similarity
    FROM ${Prisma.raw(`"${capitalizedEntityType}"`)}
    WHERE embedding IS NOT NULL
      AND 1 - (embedding <=> ${`[${queryEmbedding.join(',')}]`}::vector) > ${similarityThreshold}
    ORDER BY embedding <=> ${`[${queryEmbedding.join(',')}]`}::vector
    LIMIT ${limit}
  `;
 
    return DatabaseService.prisma[entityType as 'user'].findMany({
      where: {
        id: {
          in: vectorResults.map((r) => r.id as string),
        },
      },
    }) as Promise<T[]>;
  }
}

The code above is fetched from GitHub repository.