@titan-design/brain 0.3.0 → 0.5.0

package/dist/chunk-ZVXSW52A.js

@@ -0,0 +1,307 @@
+import {
+  addFrontmatterField
+} from "./chunk-4SD4JRLS.js";
+
+// src/services/search.ts
+import { existsSync } from "fs";
+
+// src/services/reranker.ts
+var DEFAULT_MODEL = "Xenova/ms-marco-MiniLM-L-6-v2";
+var cachedPipeline = null;
+async function getPipeline() {
+  if (cachedPipeline) return cachedPipeline;
+  const { pipeline } = await import("@huggingface/transformers");
+  cachedPipeline = await pipeline("text-classification", DEFAULT_MODEL, {
+    dtype: "q8"
+  });
+  return cachedPipeline;
+}
+async function rerank(query, results, topK) {
+  if (results.length <= 1) return results;
+  const classifier = await getPipeline();
+  const pairs = results.map((r) => [query, r.excerpt]);
+  const scores = await classifier(pairs, { top_k: 1 });
+  const scored = results.map((result, i) => ({
+    result,
+    rerankScore: scores[i].score
+  }));
+  scored.sort((a, b) => b.rerankScore - a.rerankScore);
+  const limit = topK ?? results.length;
+  return scored.slice(0, limit).map(({ result, rerankScore }) => ({
+    ...result,
+    score: rerankScore
+  }));
+}
+
+// src/services/search.ts
+var RRF_K = 60;
+var EXCERPT_MAX_LENGTH = 500;
+var OVERFETCH_MULTIPLIER = 3;
+function distanceToCosineSim(distance) {
+  return 1 - distance * distance / 2;
+}
+function truncateExcerpt(content) {
+  if (content.length <= EXCERPT_MAX_LENGTH) return content;
+  return content.slice(0, EXCERPT_MAX_LENGTH);
+}
+async function embedQuery(embedder, query) {
+  const queryText = embedder.model.includes("nomic") ? `search_query: ${query}` : query;
+  const [embedding] = await embedder.embed([queryText]);
+  return new Float32Array(embedding);
+}
+function parseTags(tagsStr) {
+  if (!tagsStr) return [];
+  return tagsStr.split(",").map((t) => t.trim()).filter(Boolean);
+}
+function applyDropoffFilter(results, threshold) {
+  if (results.length <= 1) return results;
+  let maxDrop = 0;
+  let cutIndex = -1;
+  for (let i = 1; i < results.length; i++) {
+    const prev = results[i - 1].score;
+    if (prev <= 0) continue;
+    const drop = (prev - results[i].score) / prev;
+    if (drop > maxDrop) {
+      maxDrop = drop;
+      cutIndex = i;
+    }
+  }
+  if (cutIndex > 0 && maxDrop >= threshold) {
+    return results.slice(0, cutIndex);
+  }
+  return results;
+}
+function fuseByRRF(ftsResults, vectorByNote, weights) {
+  const fusionMap = /* @__PURE__ */ new Map();
+  for (let i = 0; i < ftsResults.length; i++) {
+    const { noteId } = ftsResults[i];
+    fusionMap.set(noteId, {
+      noteId,
+      bm25Rank: i + 1,
+      vectorRank: null,
+      chunkId: null
+    });
+  }
+  let vectorRank = 1;
+  for (const [noteId, vr] of vectorByNote) {
+    const existing = fusionMap.get(noteId);
+    if (existing) {
+      existing.vectorRank = vectorRank;
+      existing.chunkId = vr.chunkId;
+    } else {
+      fusionMap.set(noteId, {
+        noteId,
+        bm25Rank: null,
+        vectorRank,
+        chunkId: vr.chunkId
+      });
+    }
+    vectorRank++;
+  }
+  const scored = [];
+  for (const entry of fusionMap.values()) {
+    let score = 0;
+    if (entry.bm25Rank !== null) {
+      score += weights.bm25 * (1 / (RRF_K + entry.bm25Rank));
+    }
+    if (entry.vectorRank !== null) {
+      score += weights.vector * (1 / (RRF_K + entry.vectorRank));
+    }
+    scored.push({ noteId: entry.noteId, score, chunkId: entry.chunkId });
+  }
+  return scored;
+}
+function fuseByScore(ftsResults, vectorByNote, weights) {
+  const fusionMap = /* @__PURE__ */ new Map();
+  for (const { noteId, rank } of ftsResults) {
+    fusionMap.set(noteId, {
+      noteId,
+      bm25Score: rank,
+      vectorDistance: null,
+      chunkId: null
+    });
+  }
+  for (const [noteId, vr] of vectorByNote) {
+    const existing = fusionMap.get(noteId);
+    if (existing) {
+      existing.vectorDistance = vr.distance;
+      existing.chunkId = vr.chunkId;
+    } else {
+      fusionMap.set(noteId, {
+        noteId,
+        bm25Score: null,
+        vectorDistance: vr.distance,
+        chunkId: vr.chunkId
+      });
+    }
+  }
+  const bm25Scores = [...fusionMap.values()].map((e) => e.bm25Score).filter((s) => s !== null);
+  let bm25Best = 0;
+  let bm25Worst = 0;
+  let bm25Range = 0;
+  if (bm25Scores.length > 0) {
+    bm25Best = Math.min(...bm25Scores);
+    bm25Worst = Math.max(...bm25Scores);
+    bm25Range = bm25Worst - bm25Best;
+  }
+  const scored = [];
+  for (const entry of fusionMap.values()) {
+    let score = 0;
+    if (entry.bm25Score !== null) {
+      const normBm25 = bm25Range === 0 ? 1 : (bm25Worst - entry.bm25Score) / bm25Range;
+      score += weights.bm25 * normBm25;
+    }
+    if (entry.vectorDistance !== null) {
+      const cosineSim = distanceToCosineSim(entry.vectorDistance);
+      score += weights.vector * Math.max(0, cosineSim);
+    }
+    scored.push({ noteId: entry.noteId, score, chunkId: entry.chunkId });
+  }
+  return scored;
+}
+async function search(db, embedder, query, options, fusionWeights = { bm25: 0.3, vector: 0.7 }, moduleRegistry) {
+  if (!query.trim()) return [];
+  const limit = options.limit;
+  const overfetchLimit = limit * OVERFETCH_MULTIPLIER;
+  let allowedNoteIds = db.getFilteredNoteIds({
+    tier: options.tier,
+    category: options.category,
+    confidence: options.confidence,
+    since: options.since,
+    tags: options.tags
+  });
+  if (moduleRegistry) {
+    const excludePm = options.excludePm || options.includePm === false;
+    const idsToExclude = getPrivateModuleNoteIds(db, moduleRegistry, excludePm ? void 0 : "pm");
+    if (idsToExclude.size > 0) {
+      if (allowedNoteIds) {
+        for (const id of idsToExclude) {
+          allowedNoteIds.delete(id);
+        }
+      } else {
+        const allNoteIds = new Set(db.getAllNotes().map((n) => n.id));
+        for (const id of idsToExclude) {
+          allNoteIds.delete(id);
+        }
+        allowedNoteIds = allNoteIds;
+      }
+    }
+  }
+  if (options.filters?.length) {
+    const metaIds = db.getFilteredNoteIdsByMetadata(options.filters, allowedNoteIds ?? void 0);
+    allowedNoteIds = metaIds;
+  }
+  const ftsResults = db.searchFTS(query, overfetchLimit);
+  const filteredFts = allowedNoteIds ? ftsResults.filter((r) => allowedNoteIds.has(r.noteId)) : ftsResults;
+  const queryVec = await embedQuery(embedder, query);
+  const vectorResults = db.searchVector(queryVec, overfetchLimit);
+  const filteredVector = allowedNoteIds ? vectorResults.filter((r) => allowedNoteIds.has(r.noteId)) : vectorResults;
+  const bestVectorByNote = /* @__PURE__ */ new Map();
+  for (const vr of filteredVector) {
+    const existing = bestVectorByNote.get(vr.noteId);
+    if (!existing || vr.distance < existing.distance) {
+      bestVectorByNote.set(vr.noteId, vr);
+    }
+  }
+  const strategy = options.fusionStrategy ?? "score";
+  const scored = strategy === "score" ? fuseByScore(filteredFts, bestVectorByNote, fusionWeights) : fuseByRRF(filteredFts, bestVectorByNote, fusionWeights);
+  scored.sort((a, b) => b.score - a.score);
+  const DEFAULT_MIN_SCORE = 0.25;
+  const effectiveMinScore = options.minScore != null ? options.minScore : strategy === "score" ? DEFAULT_MIN_SCORE : null;
+  const filtered = effectiveMinScore != null ? scored.filter((s) => s.score >= effectiveMinScore) : scored;
+  const afterDropoff = options.dropoff != null ? applyDropoffFilter(filtered, options.dropoff) : filtered;
+  const topResults = afterDropoff.slice(0, limit);
+  const results = buildSearchResults(db, topResults);
+  for (const result of results) {
+    db.recordAccess(result.noteId, "search_hit");
+  }
+  if (options.rerank && results.length > 1) {
+    return rerank(query, results, limit);
+  }
+  return results;
+}
+function computeFacets(db, facetFields, noteIds) {
+  return facetFields.map((field) => ({
+    field,
+    values: db.getFacetCounts(field, noteIds)
+  }));
+}
+function buildSearchResults(db, topResults) {
+  const noteIds = topResults.map((r) => r.noteId);
+  const notesById = db.getNotesByIds(noteIds);
+  const results = [];
+  for (const item of topResults) {
+    const note = notesById.get(item.noteId);
+    if (!note) continue;
+    const excerptContent = item.chunkId ? db.getChunkContent(item.chunkId) : db.getFirstChunkForNote(item.noteId)?.content ?? "";
+    results.push({
+      score: item.score,
+      filePath: note.filePath,
+      noteId: note.id,
+      heading: db.getChunkHeading(item.chunkId, item.noteId),
+      excerpt: truncateExcerpt(excerptContent),
+      tier: note.tier,
+      tags: parseTags(note.tags),
+      confidence: note.confidence
+    });
+  }
+  return results;
+}
+async function searchMemories(db, embedder, query, limit = 10, containerTag) {
+  if (!query.trim()) return [];
+  const queryVec = await embedQuery(embedder, query);
+  const vectorResults = db.searchMemoryVectors(queryVec, limit * 3);
+  if (vectorResults.length === 0) return [];
+  const memoryIds = vectorResults.map((vr) => vr.memoryId);
+  const memoriesById = db.getMemoriesByIds(memoryIds);
+  const results = [];
+  for (const vr of vectorResults) {
+    const memory = memoriesById.get(vr.memoryId);
+    if (!memory) continue;
+    if (!memory.isLatest || memory.isForgotten) continue;
+    if (containerTag && memory.containerTag !== containerTag) continue;
+    const cosineSim = distanceToCosineSim(vr.distance);
+    results.push({
+      score: Math.max(0, cosineSim),
+      memory: memory.memory,
+      memoryId: memory.id,
+      sourceNoteId: memory.sourceNoteId,
+      containerTag: memory.containerTag,
+      createdAt: memory.createdAt
+    });
+  }
+  results.sort((a, b) => b.score - a.score);
+  return results.slice(0, limit);
+}
+var DEFAULT_PROMOTION_THRESHOLD = 10;
+function checkAndPromote(db, noteId, threshold = DEFAULT_PROMOTION_THRESHOLD) {
+  const note = db.getNoteById(noteId);
+  if (!note || note.tier !== "fast") return false;
+  const count = db.getAccessCount(noteId);
+  if (count < threshold) return false;
+  const promoted = { ...note, tier: "slow", reviewInterval: null };
+  db.upsertNote(promoted);
+  if (existsSync(note.filePath)) {
+    addFrontmatterField(note.filePath, "tier", "slow");
+  }
+  return true;
+}
+function getPrivateModuleNoteIds(db, registry, skipModule) {
+  const privateIds = /* @__PURE__ */ new Set();
+  for (const { module: moduleName, filter } of registry.getFilters()) {
+    if (filter.visibility === "private" && moduleName !== skipModule) {
+      const noteIds = db.getModuleNoteIds({ module: moduleName });
+      for (const id of noteIds) {
+        privateIds.add(id);
+      }
+    }
+  }
+  return privateIds;
+}
+
+export {
+  search,
+  computeFacets,
+  searchMemories,
+  checkAndPromote
+};