Embeddings and vector search with pgvector

import { embed } from "./embed.ts";

function dot(a: number[], b: number[]): number {
  let sum = 0;
  for (let i = 0; i < a.length; i++) sum += a[i] * b[i];
  return sum;
}

function norm(a: number[]): number {
  return Math.sqrt(dot(a, a));
}

// Cosine similarity: dot product divided by the product of lengths
export function cosine(a: number[], b: number[]): number {
  return dot(a, b) / (norm(a) * norm(b));
}

const texts = [
  "Нужен интернет-магазин с каталогом товаров",
  "Разработка сайта с корзиной и онлайн-оплатой",
  "Как настроить резервное копирование базы данных",
];
const vectors = await embed(texts, "query");

console.log("dimensions:", vectors[0].length);
console.log("vector length:", norm(vectors[0]).toFixed(4));
console.log("first 5 numbers:", vectors[0].slice(0, 5).map((x) => x.toFixed(4)).join(", "));
for (let i = 0; i < texts.length; i++) {
  for (let j = i + 1; j < texts.length; j++) {
    console.log(`${i + 1}↔${j + 1}: ${cosine(vectors[i], vectors[j]).toFixed(3)}`);
  }
}

$ node src/cosine.ts
dimensions: 384
vector length: 1.0000
first 5 numbers: -0.0070, -0.0174, -0.0829, -0.0610, 0.0510
1↔2: 0.859
1↔3: 0.794
2↔3: 0.801

import { env, pipeline } from "@huggingface/transformers";

// Keep the weight cache in the project folder, not inside node_modules
// (survives reinstalls; add .models to .gitignore)
env.cacheDir = "./.models";

// Multilingual model: 384 dimensions, handles Russian and English.
// The first run downloads the weights from the Hugging Face Hub (~120 MB for q8).
const MODEL = "Xenova/multilingual-e5-small";
const extractor = await pipeline("feature-extraction", MODEL, { dtype: "q8" });

// E5 was trained with prefixes: "query: " for questions, "passage: " for documents.
export type Kind = "query" | "passage";

export async function embed(texts: string[], kind: Kind): Promise<number[][]> {
  const output = await extractor(
    texts.map((t) => `${kind}: ${t}`),
    { pooling: "mean", normalize: true }, // mean over tokens, vector length 1
  );
  return output.tolist() as number[][];
}

export const DIMENSIONS = 384;

import OpenAI from "openai";

// Keys come only from the environment (.env on the server), never from code or the frontend
const openai = new OpenAI(); // reads OPENAI_API_KEY

export async function embedOpenAI(texts: string[]): Promise<number[][]> {
  const res = await openai.embeddings.create({
    model: "text-embedding-3-small",
    input: texts,
    dimensions: 384, // shorten the vector to fit a vector(384) column
  });
  return res.data.map((d) => d.embedding);
}

// Voyage AI is the provider covered in Anthropic's embeddings guide.
// The key is created in MongoDB Atlas (AI Model APIs section).
const VOYAGE_URL = process.env.VOYAGE_BASE_URL ?? "https://ai.mongodb.com/v1";

export async function embedVoyage(
  texts: string[],
  inputType: "query" | "document",
): Promise<number[][]> {
  const res = await fetch(`${VOYAGE_URL}/embeddings`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${process.env.VOYAGE_API_KEY}`,
    },
    body: JSON.stringify({ input: texts, model: "voyage-4", input_type: inputType }),
  });
  if (!res.ok) throw new Error(`Voyage ${res.status}: ${await res.text()}`);
  const json = (await res.json()) as { data: { embedding: number[]; index: number }[] };
  return json.data.sort((a, b) => a.index - b.index).map((d) => d.embedding);
}

# PostgreSQL 17 + pgvector; the password test is for a throwaway local container only
docker run -d --rm --name kb-pgvector-emb -e POSTGRES_PASSWORD=test -p 55440:5432 pgvector/pgvector:pg17

npm init -y && npm pkg set type=module
npm install @huggingface/transformers@4.3.1 pg@8.23.1 pgvector@0.3.0
npm install -D typescript @types/node@24 @types/pg

# schema, corpus load and the first search
docker exec -i kb-pgvector-emb psql -U postgres -v ON_ERROR_STOP=1 < sql/01-schema.sql
export DATABASE_URL=postgres://postgres:test@localhost:55440/postgres
node src/seed.ts
node src/search.ts "нужно приложение для телефона"

[
  {"slug":"web-development","title":"Разработка сайтов","body":"Корпоративные и продуктовые сайты, каталоги и интернет-магазины: структура, UX/UI-дизайн, frontend, SEO-ready архитектура, подключение аналитики и запуск на домене заказчика."},
  {"slug":"web-app-development","title":"Веб-приложения","body":"Личные кабинеты, внутренние порталы и панели управления: роли и права доступа, бизнес-логика, база данных, API и интеграции с внешними сервисами."},
  {"slug":"mobile-development","title":"Мобильные приложения","body":"Приложения для iOS и Android: сценарии, дизайн под гайдлайны платформ, backend с API, авторизация, push-уведомления и публикация в App Store и Google Play."},
  …
]

CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE services (
  id        bigint GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
  slug      text UNIQUE NOT NULL,
  title     text NOT NULL,
  body      text NOT NULL,
  -- the model fixes the size: multilingual-e5-small returns 384 numbers
  embedding vector(384) NOT NULL,
  -- for hybrid search: lexemes of the Russian text, kept up to date automatically
  tsv       tsvector GENERATED ALWAYS AS (
              to_tsvector('russian', title || ' ' || body)
            ) STORED
);

CREATE INDEX services_tsv_idx ON services USING gin (tsv);

import pg from "pg";
import pgvector from "pgvector/pg";

// Connection string comes from the environment; no password in code
export const pool = new pg.Pool({
  connectionString: process.env.DATABASE_URL,
  // register the vector type for every new pool connection
  onConnect: async (client) => {
    await pgvector.registerTypes(client);
  },
});

import { readFile } from "node:fs/promises";
import pgvector from "pgvector/pg";
import { pool } from "./db.ts";
import { embed } from "./embed.ts";

type Doc = { slug: string; title: string; body: string };
const docs: Doc[] = JSON.parse(await readFile("data/services.json", "utf8"));

// Embed the documents with the passage prefix, in one batch
const vectors = await embed(docs.map((d) => `${d.title}. ${d.body}`), "passage");

for (const [i, d] of docs.entries()) {
  await pool.query(
    `INSERT INTO services (slug, title, body, embedding) VALUES ($1, $2, $3, $4)
     ON CONFLICT (slug) DO UPDATE
       SET title = EXCLUDED.title, body = EXCLUDED.body, embedding = EXCLUDED.embedding`,
    [d.slug, d.title, d.body, pgvector.toSql(vectors[i])],
  );
}
console.log(`documents loaded: ${docs.length}`);
await pool.end();

import pgvector from "pgvector/pg";
import { pool } from "./db.ts";
import { embed } from "./embed.ts";

const question = process.argv[2] ?? "need a mobile app for our customers";
const [q] = await embed([question], "query");

// <=> is cosine distance; similarity = 1 - distance
const { rows } = await pool.query(
  `SELECT slug, title, round((1 - (embedding <=> $1))::numeric, 3) AS similarity
     FROM services
    ORDER BY embedding <=> $1
    LIMIT 3`,
  [pgvector.toSql(q)],
);
console.log(`«${question}»`);
console.table(rows);
await pool.end();

$ node src/seed.ts
documents loaded: 12
$ node src/search.ts "нужно приложение для телефона"
┌─────────┬──────────────────────┬────────────────────────┬────────────┐
│ (index) │ slug                 │ title                  │ similarity │
├─────────┼──────────────────────┼────────────────────────┼────────────┤
│ 0       │ 'mobile-development' │ 'Мобильные приложения' │ '0.888'    │
│ 1       │ 'support'            │ 'Сопровождение'        │ '0.834'    │
│ 2       │ 'optimization'       │ 'Оптимизация продукта' │ '0.833'    │
└─────────┴──────────────────────┴────────────────────────┴────────────┘
$ node src/search.ts "хотим, чтобы заказы из Kaspi сами попадали в учёт"
│ 0       │ 'integrations'     │ 'Интеграции и API' │ '0.829'    │
│ 1       │ 'support'          │ 'Сопровождение'    │ '0.812'    │
│ 2       │ 'saas-development' │ 'SaaS-продукты'    │ '0.801'    │
$ node src/search.ts "who can review our code and security?"
│ 0       │ 'support'             │ 'Сопровождение'        │ '0.779'    │
│ 1       │ 'optimization'        │ 'Оптимизация продукта' │ '0.773'    │
│ 2       │ 'web-app-development' │ 'Веб-приложения'       │ '0.770'    │
$ node src/search.ts "technical audit of an existing app"
│ 0       │ 'audit'              │ 'Технический аудит'    │ '0.821'    │
│ 1       │ 'support'            │ 'Сопровождение'        │ '0.794'    │
│ 2       │ 'mobile-development' │ 'Мобильные приложения' │ '0.790'    │

-- Test bench: 20,000 vectors with the model's dimensionality.
-- Real embeddings form semantic clusters, so we generate vectors
-- around 200 random centers instead of uniformly.
SELECT setseed(0.42);
CREATE FUNCTION rnd_vec(scale float8) RETURNS vector LANGUAGE sql VOLATILE AS
  $$ SELECT array_agg((random() - 0.5) * scale)::vector FROM generate_series(1, 384) $$;

CREATE TABLE centers AS SELECT k, rnd_vec(1) AS c FROM generate_series(0, 199) AS k;
CREATE TABLE bench (id int PRIMARY KEY, embedding vector(384));
INSERT INTO bench
SELECT i, c + rnd_vec(1.4) FROM generate_series(1, 20000) AS i JOIN centers ON k = i % 200;

-- 50 queries and the ground truth: exact search while there is no index
CREATE TABLE queries AS
SELECT q, (c + rnd_vec(1.4))::vector(384) AS v FROM generate_series(1, 50) AS q JOIN centers ON k = q * 3;
CREATE TABLE truth AS
SELECT q.q, b.id
  FROM queries q
  CROSS JOIN LATERAL (SELECT id FROM bench ORDER BY embedding <=> q.v LIMIT 10) b;
ANALYZE bench;

EXPLAIN (COSTS OFF)
SELECT id FROM bench ORDER BY embedding <=> (SELECT v FROM queries WHERE q = 1) LIMIT 10;

\timing on
CREATE INDEX bench_hnsw ON bench USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);
\timing off

EXPLAIN (COSTS OFF)
SELECT id FROM bench ORDER BY embedding <=> (SELECT v FROM queries WHERE q = 1) LIMIT 10;

-- recall@10: share of the exact neighbors the index found
CREATE FUNCTION recall_at_10() RETURNS numeric LANGUAGE sql AS $$
  SELECT round(count(*)::numeric / (SELECT count(*) FROM truth), 3)
    FROM queries q
    CROSS JOIN LATERAL (SELECT id FROM bench ORDER BY embedding <=> q.v LIMIT 10) a
    JOIN truth t ON t.q = q.q AND t.id = a.id
$$;

SET hnsw.ef_search = 10;  SELECT 'ef_search=10'  AS setting, recall_at_10();
SET hnsw.ef_search = 40;  SELECT 'ef_search=40'  AS setting, recall_at_10();
SET hnsw.ef_search = 200; SELECT 'ef_search=200' AS setting, recall_at_10();
SELECT pg_size_pretty(pg_relation_size('bench')) AS table_size,
       pg_size_pretty(pg_relation_size('bench_hnsw')) AS hnsw_size;

DROP INDEX bench_hnsw;
-- lists ≈ rows / 1000 for tables up to 1M rows (pgvector guidance)
\timing on
CREATE INDEX bench_ivf ON bench USING ivfflat (embedding vector_cosine_ops) WITH (lists = 20);
\timing off

SET ivfflat.probes = 1; SELECT 'probes=1' AS setting, recall_at_10();
SET ivfflat.probes = 4; SELECT 'probes=4' AS setting, recall_at_10();
SELECT pg_size_pretty(pg_relation_size('bench_ivf')) AS ivfflat_size;

$ docker exec -i kb-pgvector-emb psql -U postgres < sql/03-index.sql
 Limit
   ->  Sort
         Sort Key: ((bench.embedding <=> (InitPlan 1).col1))
         ->  Seq Scan on bench
CREATE INDEX
Time: 1698.993 ms (00:01.699)
 Limit
   ->  Index Scan using bench_hnsw on bench
         Order By: (embedding <=> (InitPlan 1).col1)
 ef_search=10  |        0.910
 ef_search=40  |        1.000
 ef_search=200 |        1.000
 table_size | hnsw_size
 31 MB      | 39 MB

$ docker exec -i kb-pgvector-emb psql -U postgres < sql/04-ivfflat.sql
CREATE INDEX
Time: 234.909 ms
 probes=1 |        0.916
 probes=4 |        0.962
 ivfflat_size
 31 MB

import pgvector from "pgvector/pg";
import { pool } from "./db.ts";
import { embed } from "./embed.ts";

const question = process.argv[2] ?? "Kaspi";
const [q] = await embed([question], "query");

// Reciprocal Rank Fusion: add up 1 / (k + rank) from both lists.
// k = 60 smooths out the gap between the top positions.
const { rows } = await pool.query(
  `WITH semantic AS (
     SELECT id, rank() OVER (ORDER BY embedding <=> $2) AS r
       FROM services
      ORDER BY embedding <=> $2
      LIMIT 20
   ),
   keyword AS (
     SELECT id, rank() OVER (ORDER BY ts_rank_cd(tsv, query) DESC) AS r
       FROM services, websearch_to_tsquery('russian', $1) AS query
      WHERE tsv @@ query
      ORDER BY ts_rank_cd(tsv, query) DESC
      LIMIT 20
   )
   SELECT s.slug, semantic.r AS semantic_rank, keyword.r AS keyword_rank,
          round((coalesce(1.0 / ($3 + semantic.r), 0)
               + coalesce(1.0 / ($3 + keyword.r), 0))::numeric, 4) AS score
     FROM semantic
     FULL JOIN keyword USING (id)
     JOIN services s USING (id)
    ORDER BY score DESC
    LIMIT 5`,
  [question, pgvector.toSql(q), 60],
);
console.log(`«${question}»`);
console.table(rows);
await pool.end();

$ node src/hybrid.ts "импорт из Excel"
┌─────────┬───────────────────────┬───────────────┬──────────────┬──────────┐
│ (index) │ slug                  │ semantic_rank │ keyword_rank │ score    │
├─────────┼───────────────────────┼───────────────┼──────────────┼──────────┤
│ 0       │ 'crm-development'     │ '4'           │ '1'          │ '0.0320' │
│ 1       │ 'integrations'        │ '1'           │ null         │ '0.0164' │
│ 2       │ 'optimization'        │ '2'           │ null         │ '0.0161' │
│ 3       │ 'saas-development'    │ '3'           │ null         │ '0.0159' │
│ 4       │ 'business-automation' │ '5'           │ null         │ '0.0154' │
└─────────┴───────────────────────┴───────────────┴──────────────┴──────────┘