SPB Git forge

spb/polyllm

Public
15commits 1branches 0releases
2.2 MBsize
maindefault branch
13 days agolast push
TypeScript 97.4% SQL 1% JavaScript 0.9% CSS 0.6%
18.9 KB · 445 lines typescript
Raw Blame History
1import "server-only";2import { and, desc, eq, ilike, inArray, or, sql, type SQL } from "drizzle-orm";3import { getDb, conversations, messages, folders, projects, prompts, promptPresets, modelPresets, sharedConversations } from "@/db";4import { listRegistryModels } from "@/lib/ai/registry";5import { PROVIDER_IDS, type ProviderId } from "@/lib/ai/core/types";6import { parseQuery, makeSnippet, escapeLike, prefixTsQuery, type ParsedQuery, type FilterToken } from "./query";78/* ------------------------------------------------------------------------------------------------9 * Public shapes (re-exported for the client through src/lib/client/types.ts)10 * ---------------------------------------------------------------------------------------------- */11export const SEARCH_GROUPS = ["conversations", "messages", "models", "prompts", "projects", "presets"] as const;12export type SearchGroup = (typeof SEARCH_GROUPS)[number];1314export interface SearchConversationHit {15  id: string;16  title: string;17  modelKey: string | null;18  provider: string | null;19  projectId: string | null;20  folderId: string | null;21  pinned: boolean;22  archived: boolean;23  shared: boolean;24  messageCount: number;25  updatedAt: string;26  lastMessageAt: string | null;27}2829export interface SearchMessageHit {30  id: string;31  conversationId: string;32  title: string;33  role: string;34  modelKey: string | null;35  createdAt: string;36  /** Plain-text excerpt centred on the first matching term (highlight client-side with `query.terms`). */37  snippet: string;38}3940export interface SearchModelHit {41  key: string;42  displayName: string;43  provider: string;44  status: string;45}4647export interface SearchPromptHit {48  id: string;49  name: string;50  description: string | null;51  /** `library` = prompts table (workspace upgrade), `legacy` = prompt_presets. */52  kind: "library" | "legacy";53}5455export interface SearchProjectHit {56  id: string;57  name: string;58  description: string | null;59  icon: string | null;60  color: string | null;61}6263export interface SearchPresetHit {64  id: string;65  name: string;66  description: string | null;67  modelKey: string;68}6970export interface SearchResponse {71  query: {72    raw: string;73    text: string;74    terms: string[];75    tokens: FilterToken[];76    hasFilters: boolean;77  };78  conversations: SearchConversationHit[];79  messages: SearchMessageHit[];80  models: SearchModelHit[];81  prompts: SearchPromptHit[];82  projects: SearchProjectHit[];83  presets: SearchPresetHit[];84  nextCursor: string | null;85  /** Which engine produced the message hits (diagnostics / QA). */86  engine: "fts" | "ilike" | "none";87  tookMs: number;88}8990export interface SearchOptions {91  limit?: number;92  cursor?: string | null;93  groups?: SearchGroup[];94  now?: Date;95}9697/* ------------------------------------------------------------------------------------------------98 * Cursor (opaque, base64url JSON): { c?: updatedAt ISO | "done", m?: createdAt ISO | "done" }99 * ---------------------------------------------------------------------------------------------- */100interface Cursor {101  c?: string;102  m?: string;103}104105export function encodeCursor(c: Cursor): string {106  return Buffer.from(JSON.stringify(c), "utf8").toString("base64url");107}108109export function decodeCursor(raw: string | null | undefined): Cursor | null {110  if (!raw) return null;111  try {112    const v = JSON.parse(Buffer.from(raw, "base64url").toString("utf8")) as Cursor;113    if (typeof v !== "object" || v === null) return null;114    return { c: typeof v.c === "string" ? v.c : undefined, m: typeof v.m === "string" ? v.m : undefined };115  } catch {116    return null;117  }118}119120function cursorDate(v: string | undefined): Date | null {121  if (!v || v === "done") return null;122  const d = new Date(v);123  return Number.isNaN(d.getTime()) ? null : d;124}125126/* ------------------------------------------------------------------------------------------------127 * Filter resolution128 * ---------------------------------------------------------------------------------------------- */129const PROVIDER_ALIASES: Record<string, ProviderId> = { google: "gemini", grok: "xai", claude: "anthropic", gpt: "openai", moonshot: "kimi", "open router": "openrouter", local: "custom", ollama: "custom" };130131function resolveProviders(values: string[]): ProviderId[] {132  const out = new Set<ProviderId>();133  for (const raw of values) {134    const v = raw.toLowerCase();135    if (PROVIDER_ALIASES[v]) out.add(PROVIDER_ALIASES[v]);136    for (const p of PROVIDER_IDS) if (p.includes(v) || v.includes(p)) out.add(p);137  }138  return Array.from(out);139}140141async function resolveProjectIds(userId: string, values: string[]): Promise<string[]> {142  if (!values.length) return [];143  const conds = values.flatMap((v) => [eq(projects.id, v), ilike(projects.name, `%${escapeLike(v)}%`)]);144  const rows = await getDb()145    .select({ id: projects.id })146    .from(projects)147    .where(and(eq(projects.userId, userId), or(...conds)))148    .limit(50);149  return rows.map((r) => r.id);150}151152async function resolveFolderIds(userId: string, values: string[]): Promise<string[]> {153  if (!values.length) return [];154  const conds = values.flatMap((v) => [eq(folders.id, v), ilike(folders.name, `%${escapeLike(v)}%`)]);155  const rows = await getDb()156    .select({ id: folders.id })157    .from(folders)158    .where(and(eq(folders.userId, userId), or(...conds)))159    .limit(50);160  return rows.map((r) => r.id);161}162163async function resolveModelKeys(values: string[]): Promise<string[]> {164  if (!values.length) return [];165  const all = await listRegistryModels({ includeDeprecated: true, includeHidden: true });166  const keys = new Set<string>();167  for (const v of values) {168    for (const m of all) if (m.displayName.toLowerCase().includes(v) || m.key.toLowerCase().includes(v) || m.id.toLowerCase().includes(v)) keys.add(m.key);169  }170  return Array.from(keys);171}172173interface Resolved {174  providers: ProviderId[] | null; // null = no filter175  projectIds: string[] | null;176  folderIds: string[] | null;177  modelKeys: string[]; // extra keys resolved from display names178  impossible: boolean; // a filter matched nothing → empty result179}180181async function resolveFilters(userId: string, q: ParsedQuery): Promise<Resolved> {182  const f = q.filters;183  const [projectIds, folderIds, modelKeys] = await Promise.all([f.project.length ? resolveProjectIds(userId, f.project) : Promise.resolve(null), f.folder.length ? resolveFolderIds(userId, f.folder) : Promise.resolve(null), resolveModelKeys(f.model)]);184  const providers = f.provider.length ? resolveProviders(f.provider) : null;185  const impossible = (projectIds !== null && projectIds.length === 0) || (folderIds !== null && folderIds.length === 0) || (providers !== null && providers.length === 0);186  return { providers, projectIds, folderIds, modelKeys, impossible };187}188189/* ------------------------------------------------------------------------------------------------190 * SQL fragments191 * ---------------------------------------------------------------------------------------------- */192const sharedExists = sql<boolean>`exists (select 1 from ${sharedConversations} s where s.conversation_id = ${conversations.id} and s.revoked_at is null)`;193const convActivity = sql`coalesce(${conversations.lastMessageAt}, ${conversations.updatedAt})`;194195// Must match the expressions indexed by drizzle/0003_search_index.sql exactly.196const msgTsv = sql`to_tsvector('simple', coalesce(${messages.content}, ''))`;197const titleTsv = sql`to_tsvector('simple', coalesce(${conversations.title}, ''))`;198199function modelCond(col: SQL, values: string[], extraKeys: string[]): SQL | null {200  if (!values.length) return null;201  const parts: SQL[] = values.map((v) => sql`${col} ILIKE ${`%${escapeLike(v)}%`}`);202  if (extraKeys.length) parts.push(sql`${col} in (${sql.join(extraKeys.map((k) => sql`${k}`), sql`, `)})`);203  return or(...parts)!;204}205206/** Conversation-level conditions shared by both groups (the message query joins conversations). */207function conversationScope(userId: string, q: ParsedQuery, r: Resolved): SQL[] {208  const f = q.filters;209  const conds: SQL[] = [eq(conversations.userId, userId)];210  if (f.is.includes("pinned")) conds.push(eq(conversations.pinned, true));211  if (f.is.includes("archived")) conds.push(eq(conversations.archived, true));212  if (f.is.includes("shared")) conds.push(sharedExists);213  if (r.providers) conds.push(inArray(conversations.provider, r.providers));214  if (r.projectIds) conds.push(inArray(conversations.projectId, r.projectIds));215  if (r.folderIds) conds.push(inArray(conversations.folderId, r.folderIds));216  return conds;217}218219function ilikeAll(col: SQL | typeof messages.content | typeof conversations.title, q: ParsedQuery): SQL | null {220  const needles = [...q.words, ...q.phrases];221  if (!needles.length) return null;222  return and(...needles.map((n) => sql`${col} ILIKE ${`%${escapeLike(n)}%`}`))!;223}224225/* ------------------------------------------------------------------------------------------------226 * Groups227 * ---------------------------------------------------------------------------------------------- */228async function searchConversations(userId: string, q: ParsedQuery, r: Resolved, limit: number, after: Date | null) {229  const f = q.filters;230  const conds = conversationScope(userId, q, r);231  const mc = modelCond(sql`${conversations.modelKey}`, f.model, r.modelKeys);232  if (mc) conds.push(mc);233  if (f.after) conds.push(sql`${convActivity} >= ${f.after}`);234  if (f.before) conds.push(sql`${convActivity} < ${f.before}`);235  if (q.text) {236    const like = ilikeAll(conversations.title, q)!;237    conds.push(or(like, sql`${titleTsv} @@ websearch_to_tsquery('simple', ${q.text})`)!);238  }239  if (after) conds.push(sql`${conversations.updatedAt} < ${after}`);240  const rows = await getDb()241    .select({242      id: conversations.id,243      title: conversations.title,244      modelKey: conversations.modelKey,245      provider: conversations.provider,246      projectId: conversations.projectId,247      folderId: conversations.folderId,248      pinned: conversations.pinned,249      archived: conversations.archived,250      messageCount: conversations.messageCount,251      updatedAt: conversations.updatedAt,252      lastMessageAt: conversations.lastMessageAt,253      shared: sharedExists,254    })255    .from(conversations)256    .where(and(...conds))257    .orderBy(desc(conversations.updatedAt))258    .limit(limit + 1);259  const hasMore = rows.length > limit;260  const page: SearchConversationHit[] = rows.slice(0, limit).map((c) => ({261    id: c.id,262    title: c.title,263    modelKey: c.modelKey,264    provider: c.provider,265    projectId: c.projectId,266    folderId: c.folderId,267    pinned: c.pinned,268    archived: c.archived,269    shared: Boolean(c.shared),270    messageCount: c.messageCount,271    updatedAt: c.updatedAt.toISOString(),272    lastMessageAt: c.lastMessageAt?.toISOString() ?? null,273  }));274  return { page, hasMore };275}276277async function searchMessages(userId: string, q: ParsedQuery, r: Resolved, limit: number, after: Date | null): Promise<{ page: SearchMessageHit[]; hasMore: boolean; engine: SearchResponse["engine"] }> {278  const f = q.filters;279  const base: SQL[] = [eq(messages.userId, userId), eq(messages.active, true), sql`${messages.status} <> 'streaming'`, ...conversationScope(userId, q, r)];280  if (f.role) base.push(eq(messages.role, f.role));281  const mc = modelCond(sql`coalesce(${messages.modelKey}, ${conversations.modelKey})`, f.model, r.modelKeys);282  if (mc) base.push(mc);283  if (f.after) base.push(sql`${messages.createdAt} >= ${f.after}`);284  if (f.before) base.push(sql`${messages.createdAt} < ${f.before}`);285  if (after) base.push(sql`${messages.createdAt} < ${after}`);286287  const run = async (textCond: SQL | null) => {288    const conds = textCond ? [...base, textCond] : base;289    return getDb()290      .select({ id: messages.id, conversationId: messages.conversationId, role: messages.role, modelKey: messages.modelKey, content: messages.content, createdAt: messages.createdAt, title: conversations.title })291      .from(messages)292      .innerJoin(conversations, eq(messages.conversationId, conversations.id))293      .where(and(...conds))294      .orderBy(desc(messages.createdAt))295      .limit(limit + 1);296  };297298  let engine: SearchResponse["engine"] = "none";299  let rows: Awaited<ReturnType<typeof run>> = [];300  if (q.text) {301    const prefix = prefixTsQuery(q.words);302    const fts = prefix && !q.phrases.length ? or(sql`${msgTsv} @@ websearch_to_tsquery('simple', ${q.text})`, sql`${msgTsv} @@ to_tsquery('simple', ${prefix})`)! : sql`${msgTsv} @@ websearch_to_tsquery('simple', ${q.text})`;303    rows = await run(fts);304    engine = "fts";305    if (rows.length === 0 && !after) {306      // Partial words inside tokens (e.g. "config" in "reconfigure") — bounded per-user ILIKE fallback.307      rows = await run(ilikeAll(messages.content, q));308      engine = rows.length ? "ilike" : "fts";309    }310  } else if (f.role || f.model.length || f.after || f.before) {311    rows = await run(null);312    engine = "none";313  } else {314    return { page: [], hasMore: false, engine };315  }316  const hasMore = rows.length > limit;317  const page: SearchMessageHit[] = rows.slice(0, limit).map((m) => ({318    id: m.id,319    conversationId: m.conversationId,320    title: m.title,321    role: m.role,322    modelKey: m.modelKey,323    createdAt: m.createdAt.toISOString(),324    snippet: makeSnippet(m.content, q.terms),325  }));326  return { page, hasMore, engine };327}328329async function searchModels(q: ParsedQuery, r: Resolved, limit = 8): Promise<SearchModelHit[]> {330  if (!q.words.length && !q.phrases.length) return [];331  const all = await listRegistryModels();332  const needles = [...q.words, ...q.phrases.map((p) => p.toLowerCase())];333  return all334    .filter((m) => !r.providers || r.providers.includes(m.provider))335    .filter((m) => {336      const hay = `${m.displayName} ${m.key} ${m.provider} ${m.family ?? ""}`.toLowerCase();337      return needles.every((n) => hay.includes(n));338    })339    .slice(0, limit)340    .map((m) => ({ key: m.key, displayName: m.displayName, provider: m.provider, status: m.status }));341}342343async function searchPrompts(userId: string, q: ParsedQuery, limit = 6): Promise<SearchPromptHit[]> {344  if (!q.text) return [];345  const db = getDb();346  const needles = [...q.words, ...q.phrases];347  const like = (col: typeof prompts.name | typeof prompts.description | typeof prompts.content) => and(...needles.map((n) => ilike(col, `%${escapeLike(n)}%`)))!;348  const [lib, legacy] = await Promise.all([349    db350      .select({ id: prompts.id, name: prompts.name, description: prompts.description })351      .from(prompts)352      .where(and(eq(prompts.userId, userId), or(like(prompts.name), like(prompts.description), like(prompts.content))))353      .orderBy(desc(prompts.favorite), desc(prompts.updatedAt))354      .limit(limit),355    db356      .select({ id: promptPresets.id, name: promptPresets.name, description: promptPresets.description })357      .from(promptPresets)358      .where(and(eq(promptPresets.userId, userId), or(and(...needles.map((n) => ilike(promptPresets.name, `%${escapeLike(n)}%`))), and(...needles.map((n) => ilike(promptPresets.systemPrompt, `%${escapeLike(n)}%`))))))359      .limit(limit),360  ]);361  return [...lib.map((p) => ({ ...p, kind: "library" as const })), ...legacy.map((p) => ({ ...p, kind: "legacy" as const }))].slice(0, limit + 2);362}363364async function searchProjects(userId: string, q: ParsedQuery, limit = 6): Promise<SearchProjectHit[]> {365  if (!q.text) return [];366  const needles = [...q.words, ...q.phrases];367  const rows = await getDb()368    .select({ id: projects.id, name: projects.name, description: projects.description, icon: projects.icon, color: projects.color })369    .from(projects)370    .where(and(eq(projects.userId, userId), eq(projects.archived, false), or(and(...needles.map((n) => ilike(projects.name, `%${escapeLike(n)}%`))), and(...needles.map((n) => ilike(projects.description, `%${escapeLike(n)}%`))))))371    .orderBy(desc(projects.updatedAt))372    .limit(limit);373  return rows;374}375376async function searchPresets(userId: string, q: ParsedQuery, limit = 6): Promise<SearchPresetHit[]> {377  if (!q.text) return [];378  const needles = [...q.words, ...q.phrases];379  return getDb()380    .select({ id: modelPresets.id, name: modelPresets.name, description: modelPresets.description, modelKey: modelPresets.modelKey })381    .from(modelPresets)382    .where(and(eq(modelPresets.userId, userId), or(and(...needles.map((n) => ilike(modelPresets.name, `%${escapeLike(n)}%`))), and(...needles.map((n) => ilike(modelPresets.description, `%${escapeLike(n)}%`))))))383    .limit(limit);384}385386/* ------------------------------------------------------------------------------------------------387 * Entry point388 * ---------------------------------------------------------------------------------------------- */389export const SEARCH_MIN_CHARS = 2;390391export function isSearchable(q: ParsedQuery): boolean {392  return q.text.replace(/"/g, "").trim().length >= SEARCH_MIN_CHARS || q.hasFilters;393}394395export function emptyResponse(q: ParsedQuery, tookMs = 0): SearchResponse {396  return { query: { raw: q.raw, text: q.text, terms: q.terms, tokens: q.tokens, hasFilters: q.hasFilters }, conversations: [], messages: [], models: [], prompts: [], projects: [], presets: [], nextCursor: null, engine: "none", tookMs };397}398399export async function searchAll(userId: string, raw: string, opts: SearchOptions = {}): Promise<SearchResponse> {400  const started = Date.now();401  const q = parseQuery(raw, { now: opts.now });402  if (!isSearchable(q)) return emptyResponse(q);403  const limit = Math.min(Math.max(opts.limit ?? 12, 1), 50);404  const groups = new Set<SearchGroup>(opts.groups?.length ? opts.groups : SEARCH_GROUPS);405  const cursor = decodeCursor(opts.cursor);406  const firstPage = !cursor;407408  const r = await resolveFilters(userId, q);409  if (r.impossible) return emptyResponse(q, Date.now() - started);410411  const wantConv = groups.has("conversations") && cursor?.c !== "done" && !q.filters.role;412  const wantMsg = groups.has("messages") && cursor?.m !== "done";413414  const [conv, msg, models, promptHits, projectHits, presetHits] = await Promise.all([415    wantConv ? searchConversations(userId, q, r, limit, cursorDate(cursor?.c)) : Promise.resolve({ page: [] as SearchConversationHit[], hasMore: false }),416    wantMsg ? searchMessages(userId, q, r, limit, cursorDate(cursor?.m)) : Promise.resolve({ page: [] as SearchMessageHit[], hasMore: false, engine: "none" as const }),417    firstPage && groups.has("models") ? searchModels(q, r) : Promise.resolve([]),418    firstPage && groups.has("prompts") ? searchPrompts(userId, q) : Promise.resolve([]),419    firstPage && groups.has("projects") ? searchProjects(userId, q) : Promise.resolve([]),420    firstPage && groups.has("presets") ? searchPresets(userId, q) : Promise.resolve([]),421  ]);422423  const nextCursor =424    conv.hasMore || msg.hasMore425      ? encodeCursor({426          c: conv.hasMore ? conv.page[conv.page.length - 1].updatedAt : "done",427          m: msg.hasMore ? msg.page[msg.page.length - 1].createdAt : "done",428        })429      : null;430431  return {432    query: { raw: q.raw, text: q.text, terms: q.terms, tokens: q.tokens, hasFilters: q.hasFilters },433    conversations: conv.page,434    messages: msg.page,435    models,436    prompts: promptHits,437    projects: projectHits,438    presets: presetHits,439    nextCursor,440    engine: msg.engine,441    tookMs: Date.now() - started,442  };443}444445