Build an AI Search Engine with Hybrid Retrieval

Query -> [expand with Claude] -> parallel search:
  |-- PostgreSQL tsvector + GIN   (keyword)
  `-- pgvector HNSW, cosine       (semantic)
        -> RRF merge -> facets -> results

npm install @prisma/client@7 @prisma/adapter-pg@7 pg dotenv cohere-ai @anthropic-ai/sdk
npm install -D prisma@7 typescript @types/node @types/pg

// prisma/schema.prisma
generator client {
  provider = "prisma-client"
  output   = "../src/generated/prisma"
}

datasource db {
  provider = "postgresql"
}

model Document {
  id        String                      @id @default(cuid())
  title     String
  content   String
  category  String
  tags      String[]
  createdAt DateTime                    @default(now())
  embedding Unsupported("vector(1024)")?

  @@index([category])
  @@index([createdAt])
}

// prisma.config.ts
import "dotenv/config";
import { defineConfig, env } from "prisma/config";

export default defineConfig({
  schema: "prisma/schema.prisma",
  datasource: { url: env("DATABASE_URL") },
});

CREATE EXTENSION IF NOT EXISTS vector;

ALTER TABLE "Document" ADD COLUMN search_vector tsvector
  GENERATED ALWAYS AS (to_tsvector('english', title || ' ' || content)) STORED;

CREATE INDEX documents_search_idx ON "Document" USING GIN (search_vector);

-- HNSW needs no training step and can be built on an empty table (vector: up to 2,000 dimensions)
CREATE INDEX documents_embedding_idx ON "Document"
  USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64);

// src/lib.ts
import "dotenv/config";
import { PrismaPg } from "@prisma/adapter-pg";
import { CohereClientV2 } from "cohere-ai";
import Anthropic from "@anthropic-ai/sdk";
import { Prisma, PrismaClient } from "./generated/prisma/client";

export const prisma = new PrismaClient({
  adapter: new PrismaPg({ connectionString: process.env.DATABASE_URL }),
});
export const cohere = new CohereClientV2({ token: process.env.COHERE_API_KEY });
const anthropic = new Anthropic(); // reads ANTHROPIC_API_KEY

async function embed(texts: string[], inputType: "search_document" | "search_query") {
  const res = await cohere.embed({
    texts, // max 96 per call
    model: "embed-english-v3.0", // 1024 dimensions
    inputType,
    embeddingTypes: ["float"],
  });
  return res.embeddings.float ?? [];
}

export async function embedDocuments(documentIds: string[]) {
  const docs = await prisma.document.findMany({
    where: { id: { in: documentIds } },
    select: { id: true, title: true, content: true },
  });
  for (let i = 0; i < docs.length; i += 96) {
    const batch = docs.slice(i, i + 96);
    const vectors = await embed(
      batch.map((d) => `${d.title}\n${d.content.slice(0, 2000)}`),
      "search_document",
    );
    for (let j = 0; j < batch.length; j++) {
      const vec = `[${vectors[j].join(",")}]`;
      await prisma.$executeRaw`
        UPDATE "Document" SET embedding = ${vec}::vector WHERE id = ${batch[j].id}`;
    }
  }
}

export async function expandQuery(query: string): Promise<string[]> {
  const message = await anthropic.messages.create({
    model: "claude-haiku-5-5",
    max_tokens: 1024,
    messages: [{
      role: "user",
      content: `Write 3 alternative search queries for: "${query}"\nReturn only the queries, one per line, no numbering.`,
    }],
  });
  const block = message.content.find((b) => b.type === "text");
  const lines = block?.type === "text" ? block.text.trim().split("\n").filter(Boolean) : [];
  return [query, ...lines.slice(0, 3)];
}

const K = 60; // RRF constant
type Filters = { category?: string; tags?: string[]; dateFrom?: Date };

function whereClause(f: Filters = {}) {
  const parts = [Prisma.sql`TRUE`];
  if (f.category) parts.push(Prisma.sql`category = ${f.category}`);
  if (f.tags?.length) parts.push(Prisma.sql`tags @> ${f.tags}::text[]`);
  if (f.dateFrom) parts.push(Prisma.sql`"createdAt" >= ${f.dateFrom}`);
  return Prisma.join(parts, " AND ");
}

export async function hybridSearch(query: string, filters?: Filters) {
  const queries = await expandQuery(query);
  const where = whereClause(filters);
  const [queryVec] = await embed([query], "search_query");
  const vec = `[${queryVec.join(",")}]`;
  // Keyword side: original query plus the expansions, OR-ed together
  const tsQuery = queries.join(" or ");

  const [keyword, semantic] = await Promise.all([
    prisma.$queryRaw<{ id: string }[]>`
      SELECT id FROM "Document"
      WHERE search_vector @@ websearch_to_tsquery('english', ${tsQuery}) AND ${where}
      ORDER BY ts_rank_cd(search_vector, websearch_to_tsquery('english', ${tsQuery})) DESC
      LIMIT 100`,
    prisma.$transaction(async (tx) => {
      await tx.$executeRaw`SET LOCAL hnsw.iterative_scan = strict_order`;
      return tx.$queryRaw<{ id: string }[]>`
        SELECT id FROM "Document"
        WHERE embedding IS NOT NULL AND ${where}
        ORDER BY embedding <=> ${vec}::vector
        LIMIT 100`;
    }),
  ]);

  const scores = new Map<string, number>();
  for (const list of [keyword, semantic]) {
    list.forEach((doc, rank) =>
      scores.set(doc.id, (scores.get(doc.id) ?? 0) + 1 / (K + rank + 1)));
  }
  const ranked = [...scores.entries()].sort((a, b) => b[1] - a[1]).slice(0, 20).map(([id]) => id);

  const rows = await prisma.document.findMany({
    where: { id: { in: ranked } },
    select: { id: true, title: true, category: true, tags: true, createdAt: true },
  });
  return ranked.map((id) => rows.find((r) => r.id === id)!); // keep RRF order
}

// src/route.ts
import { hybridSearch, prisma } from "./lib";

// app/api/search/route.ts (Next.js App Router)
export async function GET(req: Request) {
  const { searchParams } = new URL(req.url);
  const from = searchParams.get("from");
  const results = await hybridSearch(searchParams.get("q") ?? "", {
    category: searchParams.get("category") ?? undefined,
    tags: searchParams.getAll("tag"),
    dateFrom: from ? new Date(from) : undefined,
  });
  const facets = await prisma.document.groupBy({
    by: ["category"],
    _count: { id: true },
    where: { id: { in: results.map((r) => r.id) } },
  });
  return Response.json({ results, facets });
}