Эмбеддинги и векторный поиск на 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));
}

// Косинусная близость: скалярное произведение, делённое на произведение длин
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("размерность:", vectors[0].length);
console.log("длина вектора:", norm(vectors[0]).toFixed(4));
console.log("первые 5 чисел:", 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
размерность: 384
длина вектора: 1.0000
первые 5 чисел: -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";

// Кеш весов — в папке проекта, а не внутри node_modules
// (переживёт переустановку пакетов; добавьте .models в .gitignore)
env.cacheDir = "./.models";

// Многоязычная модель: 384 измерения, понимает русский и английский.
// Первый запуск скачивает веса с Hugging Face Hub (~120 МБ для q8).
const MODEL = "Xenova/multilingual-e5-small";
const extractor = await pipeline("feature-extraction", MODEL, { dtype: "q8" });

// E5 обучали с префиксами: "query: " для вопроса, "passage: " для документа.
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 }, // среднее по токенам, длина вектора 1
  );
  return output.tolist() as number[][];
}

export const DIMENSIONS = 384;

import OpenAI from "openai";

// Ключи — только из окружения (.env на сервере), никогда в коде и не во фронтенде
const openai = new OpenAI(); // читает 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, // укоротить вектор под колонку vector(384)
  });
  return res.data.map((d) => d.embedding);
}

// Voyage AI — провайдер из руководства Anthropic по эмбеддингам.
// Ключ создаётся в MongoDB Atlas (раздел AI Model APIs).
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; пароль test — только для локального одноразового контейнера
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

# схема, загрузка корпуса и первый поиск
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,
  -- размерность фиксирована моделью: multilingual-e5-small даёт 384 числа
  embedding vector(384) NOT NULL,
  -- для гибридного поиска: лексемы русского текста, обновляются сами
  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";

// Строка подключения — из окружения, пароль в коде не храним
export const pool = new pg.Pool({
  connectionString: process.env.DATABASE_URL,
  // регистрируем тип vector для каждого нового соединения пула
  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"));

// Документы кодируем с префиксом passage — одним батчем
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(`загружено документов: ${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] ?? "нужно приложение для телефона";
const [q] = await embed([question], "query");

// <=> — косинусное расстояние; близость = 1 - расстояние
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
загружено документов: 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'    │

-- Стенд: 20 000 векторов той же размерности, что у модели.
-- Настоящие эмбеддинги образуют смысловые «облака», поэтому векторы
-- генерируем вокруг 200 случайных центров, а не равномерно.
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 запросов и «правильные ответы» — точный поиск, пока индекса нет
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: какая доля точных соседей нашлась через индекс
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 для таблиц до миллиона строк (рекомендация pgvector)
\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: складываем 1 / (k + место) из двух списков.
// k = 60 сглаживает разницу между первыми местами.
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' │
└─────────┴───────────────────────┴───────────────┴──────────────┴──────────┘