agentikit 0.0.9 → 0.0.12

package/src/stash-search.ts

@@ -1,12 +1,10 @@
 import fs from "node:fs"
 import path from "node:path"
-import { type AgentikitAssetType, hasErrnoCode, resolveStashDir } from "./common"
+import { type AgentikitAssetType } from "./common"
 import { ASSET_TYPES, TYPE_DIRS, deriveCanonicalAssetName } from "./asset-spec"
-import { loadSearchIndex, buildSearchText, type IndexedEntry } from "./indexer"
-import { TfIdfAdapter, type ScoredEntry } from "./similarity"
-import { buildToolInfo } from "./tool-runner"
+import { buildSearchText } from "./indexer"
 import { walkStash } from "./walker"
-import { makeOpenRef } from "./stash-ref"
+import { makeAssetRef } from "./stash-ref"
 import type {
   AgentikitSearchType,
   LocalSearchHit,
@@ -18,6 +16,21 @@ import type {
 } from "./stash-types"
 import { loadConfig } from "./config"
 import { searchRegistry } from "./registry-search"
+import {
+  openDatabase,
+  closeDatabase,
+  getDbPath,
+  getMeta,
+  searchFts,
+  searchVec,
+  getAllEntries,
+  getEntryCount,
+  getEntryById,
+  isVecAvailable,
+  type DbSearchResult,
+} from "./db"
+import { tryGetHandler } from "./asset-type-handler"
+import { type StashSource, resolveStashSources, findSourceForPath } from "./stash-source"
 
 type IndexedAsset = {
   type: AgentikitAssetType
@@ -27,29 +40,6 @@ type IndexedAsset = {
 
 const DEFAULT_LIMIT = 20
 
-const DEFAULT_USAGE_GUIDE_BY_TYPE: Record<AgentikitAssetType, string[]> = {
-  tool: [
-    "Use the hit's runCmd for execution so runtime and working directory stay correct.",
-    "Use `akm show <openRef>` to inspect the tool before running it.",
-  ],
-  skill: [
-    "Read and apply the skill instructions as written, then adapt examples to your current repo state and task.",
-    "Use `akm show <openRef>` to read the full SKILL.md for required steps and constraints.",
-  ],
-  command: [
-    "Read the .md file, fill placeholders, and run it in the current repo context.",
-    "Use `akm show <openRef>` to retrieve the command template body.",
-  ],
-  agent: [
-    "Read the .md file and dispatch and agent using the content of the file. Use modelHint/toolPolicy when present to run the agent with compatible settings.",
-    "Use with `akm show <openRef>` to get the full prompt payload.",
-  ],
-  knowledge: [
-    "Use `akm show <openRef>` to read the document; start with `--view toc` for large files.",
-    "Use `--view section` or `--view lines` to load only the part you need.",
-  ],
-}
-
 export async function agentikitSearch(input: {
   query: string
   type?: AgentikitSearchType
@@ -64,7 +54,8 @@ export async function agentikitSearch(input: {
   const limit = normalizeLimit(input.limit)
   const usageMode = parseSearchUsageMode(input.usage)
   const source = parseSearchSource(input.source)
-  const stashDir = resolveStashDir()
+  const sources = resolveStashSources()
+  const stashDir = sources[0].path
   const localResult = source === "registry"
     ? undefined
     : await searchLocal({
@@ -73,11 +64,13 @@ export async function agentikitSearch(input: {
       limit,
       usageMode,
       stashDir,
+      sources,
     })
 
+  const config = loadConfig()
   const registryResult = source === "local"
     ? undefined
-    : await searchRegistry(query, { limit })
+    : await searchRegistry(query, { limit, registryUrls: config.registryUrls })
 
   if (source === "local") {
     return {
@@ -103,6 +96,7 @@ export async function agentikitSearch(input: {
       homepage: hit.homepage,
       score: hit.score,
       metadata: hit.metadata,
+      curated: hit.curated,
       installRef,
       installCmd: `akm add ${installRef}`,
     }
@@ -139,30 +133,41 @@ async function searchLocal(input: {
   limit: number
   usageMode: SearchUsageMode
   stashDir: string
+  sources: StashSource[]
 }): Promise<{ hits: LocalSearchHit[]; usageGuide?: Partial<Record<AgentikitAssetType, string[]>>; tip?: string; embedMs?: number; rankMs?: number }> {
-  const { query, searchType, limit, usageMode, stashDir } = input
-  const config = loadConfig(stashDir)
-  const allStashDirs = [
-    stashDir,
-    ...config.additionalStashDirs.filter((d) => {
-      try { return fs.statSync(d).isDirectory() } catch { return false }
-    }),
-  ]
+  const { query, searchType, limit, usageMode, stashDir, sources } = input
+  const config = loadConfig()
+  const allStashDirs = sources.map((s) => s.path)
 
-  const index = loadSearchIndex()
-  if (index && index.entries && index.entries.length > 0 && index.stashDir === stashDir) {
-    const { hits, usageGuide, embedMs, rankMs } = await searchIndex(index, query, searchType, limit, stashDir, allStashDirs, config, usageMode)
-    return {
-      hits,
-      usageGuide,
-      tip: hits.length === 0 ? "No matching stash assets were found. Try running 'akm index' to rebuild." : undefined,
-      embedMs,
-      rankMs,
+  // Try to open the database
+  const dbPath = getDbPath()
+  try {
+    if (fs.existsSync(dbPath)) {
+      const embeddingDim = config.embedding?.dimension
+      const db = openDatabase(dbPath, embeddingDim ? { embeddingDim } : undefined)
+      try {
+        const entryCount = getEntryCount(db)
+        const storedStashDir = getMeta(db, "stashDir")
+        if (entryCount > 0 && storedStashDir === stashDir) {
+          const { hits, usageGuide, embedMs, rankMs } = await searchDatabase(db, query, searchType, limit, stashDir, allStashDirs, config, usageMode, sources)
+          return {
+            hits,
+            usageGuide,
+            tip: hits.length === 0 ? "No matching stash assets were found. Try running 'akm index' to rebuild." : undefined,
+            embedMs,
+            rankMs,
+          }
+        }
+      } finally {
+        closeDatabase(db)
+      }
     }
+  } catch (error) {
+    console.warn("Search index unavailable, falling back to substring search:", error instanceof Error ? error.message : String(error))
   }
 
   const hits = allStashDirs
-    .flatMap((dir) => substringSearch(query, searchType, limit, dir))
+    .flatMap((dir) => substringSearch(query, searchType, limit, dir, sources))
     .slice(0, limit)
   const usageGuide = shouldIncludeUsageGuide(usageMode) ? buildUsageGuide(hits.map((hit) => hit.type), searchType) : undefined
   return {
@@ -172,10 +177,10 @@ async function searchLocal(input: {
   }
 }
 
-// ── Unified indexed search ──────────────────────────────────────────────────
+// ── Database search ─────────────────────────────────────────────────────────
 
-async function searchIndex(
-  index: import("./indexer").SearchIndex,
+async function searchDatabase(
+  db: import("bun:sqlite").Database,
   query: string,
   searchType: AgentikitSearchType,
   limit: number,
@@ -183,80 +188,136 @@ async function searchIndex(
   allStashDirs: string[],
   config: import("./config").AgentikitConfig,
   usageMode: SearchUsageMode,
+  sources: StashSource[],
 ): Promise<{ hits: LocalSearchHit[]; usageGuide?: Partial<Record<AgentikitAssetType, string[]>>; embedMs?: number; rankMs?: number }> {
-  // Filter candidates by type
-  let candidates = index.entries
-  if (searchType !== "any") {
-    candidates = candidates.filter((ie) => ie.entry.type === searchType)
-  }
-
-  if (candidates.length === 0) {
-    return {
-      hits: [],
-      usageGuide: shouldIncludeUsageGuide(usageMode) ? buildUsageGuide([], searchType) : undefined,
-    }
-  }
-
-  // Empty query: return all entries (no scoring needed)
+  // Empty query: return all entries
   if (!query) {
-    const selectedCandidates = candidates.slice(0, limit)
-    const hits = selectedCandidates.map((ie) =>
-      buildIndexedHit({
+    const typeFilter = searchType === "any" ? undefined : searchType
+    const allEntries = getAllEntries(db, typeFilter)
+    const selected = allEntries.slice(0, limit)
+    const hits = selected.map((ie) =>
+      buildDbHit({
         entry: ie.entry,
-        path: ie.path,
+        path: ie.filePath,
         score: 1,
         query,
-        rankingMode: "tfidf",
+        rankingMode: "fts",
         defaultStashDir: stashDir,
         allStashDirs,
+        sources,
         includeItemUsage: shouldIncludeItemUsage(usageMode),
       }),
     )
     return {
       hits,
       usageGuide: shouldIncludeUsageGuide(usageMode)
-        ? buildUsageGuideFromEntries(selectedCandidates.map((candidate) => candidate.entry), searchType)
+        ? buildUsageGuideFromEntries(selected.map((e) => e.entry), searchType)
         : undefined,
     }
   }
 
-  // Score each candidate using available signals
+  // Score using FTS5 (BM25) and optionally sqlite-vec
   const tEmbed0 = Date.now()
-  const embeddingScores = await tryEmbeddingScores(candidates, query, config)
+  const embeddingScores = await tryVecScores(db, query, limit * 3, config)
   const embedMs = Date.now() - tEmbed0
 
   const tRank0 = Date.now()
-  const tfidfScores = computeTfidfScores(index, candidates, query, searchType)
+  const typeFilter = searchType === "any" ? undefined : searchType
+  const ftsResults = searchFts(db, query, limit * 3, typeFilter)
+
+  // Build score map from FTS results (normalize BM25 scores)
+  const ftsScoreMap = new Map<number, { score: number; result: DbSearchResult }>()
+  for (const r of ftsResults) {
+    // BM25 returns negative scores (more negative = better match), normalize to 0-1
+    const absScore = Math.abs(r.bm25Score)
+    const normalized = absScore / (1 + absScore)
+    ftsScoreMap.set(r.id, { score: normalized, result: r })
+  }
 
-  const scored: Array<{ ie: IndexedEntry; score: number; rankingMode: "semantic" | "tfidf" }> = []
+  // Blend scores
+  const scored: Array<{ id: number; entry: import("./metadata").StashEntry; filePath: string; score: number; rankingMode: "semantic" | "fts" }> = []
+  const seenIds = new Set<number>()
 
-  for (const ie of candidates) {
-    const key = ie.path
-    const embScore = embeddingScores?.get(key)
-    const tfidfScore = tfidfScores.get(key) ?? 0
+  // Process FTS results
+  for (const [id, { score: ftsScore, result }] of ftsScoreMap) {
+    seenIds.add(id)
+    const embScore = embeddingScores?.get(id)
 
     if (embScore !== undefined) {
-      // Weighted blend: embedding dominates when available, TF-IDF boosts lexical matches
-      const blended = embScore * 0.7 + tfidfScore * 0.3
-      if (blended > 0) scored.push({ ie, score: blended, rankingMode: "semantic" })
-    } else if (tfidfScore > 0) {
-      scored.push({ ie, score: tfidfScore, rankingMode: "tfidf" })
+      const blended = embScore * 0.7 + ftsScore * 0.3
+      if (blended > 0) scored.push({ id, entry: result.entry, filePath: result.filePath, score: blended, rankingMode: "semantic" })
+    } else if (ftsScore > 0) {
+      scored.push({ id, entry: result.entry, filePath: result.filePath, score: ftsScore, rankingMode: "fts" })
+    }
+  }
+
+  // Add vec-only results not already in FTS results
+  if (embeddingScores) {
+    for (const [id, embScore] of embeddingScores) {
+      if (seenIds.has(id)) continue
+      const found = getEntryById(db, id)
+      if (found) {
+        if (typeFilter && found.entry.type !== typeFilter) continue
+        scored.push({
+          id,
+          entry: found.entry,
+          filePath: found.filePath,
+          score: embScore,
+          rankingMode: "semantic",
+        })
+      }
+    }
+  }
+
+  // Apply boosts (tag, intent, name matches)
+  const queryTokens = query.toLowerCase().split(/\s+/).filter(Boolean)
+  for (const item of scored) {
+    const entry = item.entry
+    // Tag boost
+    if (entry.tags) {
+      for (const tag of entry.tags) {
+        if (queryTokens.some((t) => tag.toLowerCase() === t)) {
+          item.score += 0.15
+        }
+      }
+    }
+    // Intent boost
+    if (entry.intents) {
+      for (const intent of entry.intents) {
+        const intentLower = intent.toLowerCase()
+        for (const token of queryTokens) {
+          if (intentLower.includes(token)) {
+            item.score += 0.12
+            break
+          }
+        }
+      }
+    }
+    // Name boost
+    const nameLower = entry.name.toLowerCase().replace(/[-_]/g, " ")
+    if (queryTokens.some((t) => nameLower.includes(t))) {
+      item.score += 0.1
     }
   }
 
+  for (const item of scored) {
+    item.score = Math.min(item.score, 1.0)
+  }
+
   scored.sort((a, b) => b.score - a.score)
   const rankMs = Date.now() - tRank0
 
   const selected = scored.slice(0, limit)
-  const hits = selected.map(({ ie, score, rankingMode }) =>
-    buildIndexedHit({
-      entry: ie.entry,
-      path: ie.path,
+  const hits = selected.map(({ entry, filePath, score, rankingMode }) =>
+    buildDbHit({
+      entry,
+      path: filePath,
       score: Math.round(score * 1000) / 1000,
       query,
       rankingMode,
       defaultStashDir: stashDir,
       allStashDirs,
+      sources,
       includeItemUsage: shouldIncludeItemUsage(usageMode),
     }),
   )
@@ -266,65 +327,41 @@ async function searchIndex(
     rankMs,
     hits,
     usageGuide: shouldIncludeUsageGuide(usageMode)
-      ? buildUsageGuideFromEntries(selected.map((item) => item.ie.entry), searchType)
+      ? buildUsageGuideFromEntries(selected.map((item) => item.entry), searchType)
       : undefined,
   }
 }
 
-// ── Embedding scorer ────────────────────────────────────────────────────────
+// ── Vector scorer ───────────────────────────────────────────────────────────
 
-async function tryEmbeddingScores(
-  candidates: IndexedEntry[],
+async function tryVecScores(
+  db: import("bun:sqlite").Database,
   query: string,
+  k: number,
   config: import("./config").AgentikitConfig,
-): Promise<Map<string, number> | null> {
-  if (!config.semanticSearch) return null
-
-  const withEmbeddings = candidates.filter((ie) => ie.embedding && ie.embedding.length > 0)
-  if (withEmbeddings.length === 0) return null
+): Promise<Map<number, number> | null> {
+  if (!config.semanticSearch || !isVecAvailable()) return null
+  const hasEmbeddings = getMeta(db, "hasEmbeddings")
+  if (hasEmbeddings !== "1") return null
 
   try {
-    const { embed, cosineSimilarity } = await import("./embedder.js")
+    const { embed } = await import("./embedder.js")
     const queryEmbedding = await embed(query, config.embedding)
-    const scores = new Map<string, number>()
-    for (const ie of withEmbeddings) {
-      scores.set(ie.path, cosineSimilarity(queryEmbedding, ie.embedding!))
+    const vecResults = searchVec(db, queryEmbedding, k)
+
+    const scores = new Map<number, number>()
+    for (const { id, distance } of vecResults) {
+      // Convert L2 distance to cosine similarity (vectors are normalized)
+      const cosineSim = 1 - (distance * distance) / 2
+      scores.set(id, Math.max(0, cosineSim))
     }
     return scores
-  } catch {
+  } catch (error) {
+    console.warn("Vector search failed, skipping:", error instanceof Error ? error.message : String(error))
     return null
   }
 }
 
-// ── TF-IDF scorer ───────────────────────────────────────────────────────────
-
-function computeTfidfScores(
-  index: import("./indexer").SearchIndex,
-  candidates: IndexedEntry[],
-  query: string,
-  searchType: AgentikitSearchType,
-): Map<string, number> {
-  const candidateScoredEntries = toScoredEntries(candidates)
-
-  let adapter: TfIdfAdapter
-  if (index.tfidf) {
-    const allScored = toScoredEntries(index.entries)
-    adapter = TfIdfAdapter.deserialize(index.tfidf, allScored)
-  } else {
-    adapter = new TfIdfAdapter()
-    adapter.buildIndex(candidateScoredEntries)
-  }
-
-  const typeFilter = searchType === "any" ? undefined : searchType
-  const results = adapter.search(query, candidates.length, typeFilter)
-
-  const scores = new Map<string, number>()
-  for (const r of results) {
-    scores.set(r.path, r.score)
-  }
-  return scores
-}
-
 // ── Substring fallback (no index) ───────────────────────────────────────────
 
 function substringSearch(
@@ -332,52 +369,50 @@ function substringSearch(
   searchType: AgentikitSearchType,
   limit: number,
   stashDir: string,
+  sources: StashSource[],
 ): LocalSearchHit[] {
   const assets = indexAssets(stashDir, searchType)
   return assets
     .filter((asset) => asset.name.toLowerCase().includes(query))
     .sort(compareAssets)
     .slice(0, limit)
-    .map((asset) => assetToSearchHit(asset, stashDir))
+    .map((asset) => assetToSearchHit(asset, stashDir, sources))
 }
 
 // ── Hit building ────────────────────────────────────────────────────────────
 
-function findStashDirForPath(filePath: string, stashDirs: string[]): string | undefined {
-  const resolved = path.resolve(filePath)
-  for (const dir of stashDirs) {
-    if (resolved.startsWith(path.resolve(dir) + path.sep)) return dir
-  }
-  return undefined
-}
-
-function buildIndexedHit(input: {
-  entry: IndexedEntry["entry"]
+function buildDbHit(input: {
+  entry: import("./metadata").StashEntry
   path: string
   score: number
   query: string
-  rankingMode: "semantic" | "tfidf"
+  rankingMode: "semantic" | "fts"
   defaultStashDir: string
   allStashDirs: string[]
+  sources: StashSource[]
   includeItemUsage: boolean
 }): LocalSearchHit {
-  const entryStashDir = findStashDirForPath(input.path, input.allStashDirs) ?? input.defaultStashDir
+  const entryStashDir = findSourceForPath(input.path, input.sources)?.path ?? input.defaultStashDir
   const typeRoot = path.join(entryStashDir, TYPE_DIRS[input.entry.type])
   const openRefName = deriveCanonicalAssetName(input.entry.type, typeRoot, input.path)
     ?? input.entry.name
 
   const qualityBoost = input.entry.generated === true ? 0 : 0.05
   const confidenceBoost = typeof input.entry.confidence === "number" ? Math.min(0.05, Math.max(0, input.entry.confidence) * 0.05) : 0
-  const score = Math.round((input.score + qualityBoost + confidenceBoost) * 1000) / 1000
+  const score = Math.min(Math.round((input.score + qualityBoost + confidenceBoost) * 1000) / 1000, 1.0)
 
   const whyMatched = buildWhyMatched(input.entry, input.query, input.rankingMode, qualityBoost, confidenceBoost)
 
+  const source = findSourceForPath(input.path, input.sources)
+
   const hit: LocalSearchHit = {
     hitSource: "local",
     type: input.entry.type,
     name: input.entry.name,
     path: input.path,
-    openRef: makeOpenRef(input.entry.type, openRefName),
+    openRef: makeAssetRef(input.entry.type, openRefName, source?.registryId),
+    registryId: source?.registryId,
+    editable: source?.writable ?? false,
     description: input.entry.description,
     tags: input.entry.tags,
     score,
@@ -388,27 +423,22 @@ function buildIndexedHit(input: {
     hit.usage = input.entry.usage
   }
 
-  if (input.entry.type === "tool") {
-    try {
-      const toolInfo = buildToolInfo(entryStashDir, input.path)
-      hit.runCmd = toolInfo.runCmd
-      hit.kind = toolInfo.kind
-    } catch (error: unknown) {
-      if (!hasErrnoCode(error, "ENOENT")) throw error
-    }
+  const handler = tryGetHandler(input.entry.type)
+  if (handler?.enrichSearchHit) {
+    handler.enrichSearchHit(hit, entryStashDir)
   }
 
   return hit
 }
 
 function buildWhyMatched(
-  entry: IndexedEntry["entry"],
+  entry: import("./metadata").StashEntry,
   query: string,
-  rankingMode: "semantic" | "tfidf",
+  rankingMode: "semantic" | "fts",
   qualityBoost: number,
   confidenceBoost: number,
 ): string[] {
-  const reasons: string[] = [rankingMode === "semantic" ? "semantic similarity" : "tf-idf lexical relevance"]
+  const reasons: string[] = [rankingMode === "semantic" ? "semantic similarity" : "fts bm25 relevance"]
   const tokens = query.toLowerCase().split(/\s+/).filter(Boolean)
 
   const name = entry.name.toLowerCase()
@@ -428,35 +458,22 @@ function buildWhyMatched(
 
 // ── Helpers ─────────────────────────────────────────────────────────────────
 
-function toScoredEntries(entries: IndexedEntry[]): ScoredEntry[] {
-  return entries.map((ie) => ({
-    id: `${ie.entry.type}:${ie.entry.name}`,
-    text: buildSearchText(ie.entry),
-    entry: ie.entry,
-    path: ie.path,
-  }))
-}
-
-function assetToSearchHit(asset: IndexedAsset, stashDir: string): LocalSearchHit {
-  if (asset.type !== "tool") {
-    return {
-      hitSource: "local",
-      type: asset.type,
-      name: asset.name,
-      path: asset.path,
-      openRef: makeOpenRef(asset.type, asset.name),
-    }
-  }
-  const toolInfo = buildToolInfo(stashDir, asset.path)
-  return {
+function assetToSearchHit(asset: IndexedAsset, stashDir: string, sources: StashSource[]): LocalSearchHit {
+  const source = findSourceForPath(asset.path, sources)
+  const hit: LocalSearchHit = {
     hitSource: "local",
-    type: "tool",
+    type: asset.type,
     name: asset.name,
     path: asset.path,
-    openRef: makeOpenRef("tool", asset.name),
-    runCmd: toolInfo.runCmd,
-    kind: toolInfo.kind,
+    openRef: makeAssetRef(asset.type, asset.name, source?.registryId),
+    registryId: source?.registryId,
+    editable: source?.writable ?? false,
+  }
+  const handler = tryGetHandler(asset.type)
+  if (handler?.enrichSearchHit) {
+    handler.enrichSearchHit(hit, stashDir)
   }
+  return hit
 }
 
 function normalizeLimit(limit?: number): number {
@@ -509,7 +526,7 @@ function shouldIncludeItemUsage(mode: SearchUsageMode): boolean {
 }
 
 function buildUsageGuideFromEntries(
-  entries: IndexedEntry["entry"][],
+  entries: import("./metadata").StashEntry[],
   searchType: AgentikitSearchType,
 ): Partial<Record<AgentikitAssetType, string[]>> | undefined {
   const types = entries.map((entry) => entry.type)
@@ -564,7 +581,8 @@ function resolveGuideTypes(hitTypes: AgentikitAssetType[], searchType: Agentikit
 }
 
 function usageGuideByType(type: AgentikitAssetType): string[] {
-  return DEFAULT_USAGE_GUIDE_BY_TYPE[type]
+  const handler = tryGetHandler(type)
+  return handler?.defaultUsageGuide ?? []
 }
 
 function fileToAsset(assetType: AgentikitAssetType, root: string, file: string): IndexedAsset | undefined {