opencode-config/.opencode/tools/memory-search.ts

import { spawnSync } from "child_process"
import fs from "fs"
import path from "path"
import { tool } from "@opencode-ai/plugin"
import {
  parseFrontmatter,
  resolveMemoryDir,
  resolveRgBinary,
  type FrontmatterMeta,
} from "./_memory-shared"

const STOP_WORDS = new Set(["the", "a", "an", "is", "of", "to", "in", "for", "on", "and", "or"])

function parseSearchTerms(query: string): string[] {
  const terms: string[] = []
  const regex = /"([^"]+)"|(\S+)/g
  let match
  while ((match = regex.exec(query)) !== null) {
    const original = match[1] || match[2]
    const lower = original.toLowerCase()
    if (lower.length < 2) continue
    if (STOP_WORDS.has(lower)) continue
    terms.push(original)
  }
  return terms
}

function countTermMatches(text: string, terms: string[]): number {
  if (terms.length === 0) return 0
  const lowerText = text.toLowerCase()
  return terms.filter((t) => lowerText.includes(t.toLowerCase())).length
}

interface ScoreInput {
  rgMatch: boolean
  ragScore?: number
  termMatches: number
  totalTerms: number
  tags?: string[]
  path: string
  importance?: string
  accessCount?: number
  terms: string[]
}

function scoreCandidate(input: ScoreInput): number {
  let score = 0
  if (input.rgMatch) {
    const termCoverage = input.totalTerms > 0 ? input.termMatches / input.totalTerms : 1
    score += 0.15 + 0.35 * termCoverage
  }
  if (input.ragScore) score += input.ragScore * 1.4
  if (input.rgMatch && input.ragScore) score += 0.1
  if (input.tags && input.tags.length > 0 && input.terms.length > 0) {
    const lowerTags = input.tags.map((t) => t.toLowerCase())
    const tagHits = input.terms.filter((t) =>
      lowerTags.some((tag) => tag.includes(t.toLowerCase()))
    ).length
    if (tagHits > 0) score += 0.2 * (tagHits / input.terms.length)
  }
  const pathLower = input.path.toLowerCase().replace(/[-_/.]/g, " ")
  const pathHits = input.terms.filter((t) => pathLower.includes(t.toLowerCase())).length
  if (pathHits > 0 && input.terms.length > 0) {
    score += 0.15 * (pathHits / input.terms.length)
  }
  if (input.importance === "high") score += 0.15
  else if (input.importance === "low") score -= 0.1
  if (input.accessCount && input.accessCount >= 5) score += 0.1
  else if (input.accessCount && input.accessCount >= 2) score += 0.05
  return score
}

interface RagHit {
  source: string
  score: number
  text: string
}

function parseRagHits(ragText: string): RagHit[] {
  if (!ragText) return []
  try {
    const parsed = JSON.parse(ragText)
    if (!Array.isArray(parsed)) return []
    return parsed.filter(
      (hit): hit is RagHit =>
        hit &&
        typeof hit.source === "string" &&
        typeof hit.score === "number" &&
        typeof hit.text === "string"
    )
  } catch (e) {
    process.stderr.write(`[memory-search] parseRagHits JSON parse failed: ${(e as Error).message}\n`)
    return []
  }
}

function buildRgArgs(terms: string[]): string[] {
  const args = ["-il", "--glob", "*.md", "--glob", "!.git", "--glob", "!.rag", "--glob", "!**/INDEX.md"]
  for (const term of terms) args.push("-e", term)
  return args
}

function toRelPath(memoryDir: string, absPath: string): string {
  return path.relative(memoryDir, path.normalize(absPath))
}

interface CandInfo {
  rgMatch: boolean
  ragScore?: number
  ragText?: string
}

function runKeywordSearch(
  rgBin: string,
  terms: string[],
  searchDir: string,
  memoryDir: string
): Map<string, CandInfo> {
  const resultMap = new Map<string, CandInfo>()
  const rgR = spawnSync(rgBin, [...buildRgArgs(terms), searchDir], { encoding: "utf-8" })
  if (rgR.status === 0) {
    for (const line of (rgR.stdout || "").trim().split("\n").filter(Boolean)) {
      const rel = toRelPath(memoryDir, line)
      if (rel.endsWith(".md") && !rel.startsWith(".")) {
        resultMap.set(rel, { rgMatch: true })
      }
    }
  }
  return resultMap
}

function runSemanticSearch(
  query: string,
  indexDir: string,
  worktree: string
): RagHit[] {
  if (!process.env.OPENAI_BASE_URL) return []
  if (!fs.existsSync(path.join(indexDir, "index.json"))) return []
  const r = spawnSync(
    "python3",
    ["-m", "src.memory", "search", query, "-i", indexDir, "-k", "15", "--json"],
    { cwd: worktree, encoding: "utf-8" }
  )
  if (r.status !== 0) {
    if (r.stderr) process.stderr.write(`[memory-search] rag search failed (exit ${r.status}): ${r.stderr.trim()}\n`)
    return []
  }
  return parseRagHits(r.stdout || "")
}

function mergeRagHits(
  resultMap: Map<string, CandInfo>,
  ragHits: RagHit[],
  category?: string
): void {
  for (const hit of ragHits) {
    if (hit.source.endsWith("INDEX.md")) continue
    if (category && !hit.source.startsWith(category + "/")) continue
    const existing = resultMap.get(hit.source) || { rgMatch: false }
    if (!existing.ragScore || hit.score > existing.ragScore) {
      existing.ragScore = hit.score
      existing.ragText = hit.text
    }
    resultMap.set(hit.source, existing)
  }
}

interface ResultRow {
  path: string
  meta: FrontmatterMeta
  score: number
  termMatches: number
  rgMatch: boolean
  ragScore?: number
  ragText?: string
}

function scoreAndCombine(
  resultMap: Map<string, CandInfo>,
  terms: string[],
  memoryDir: string
): ResultRow[] {
  const results: ResultRow[] = []
  for (const [p, info] of resultMap) {
    try {
      const content = fs.readFileSync(path.join(memoryDir, p), "utf-8")
      const meta = parseFrontmatter(content)
      const termMatches =
        terms.length > 0 ? countTermMatches(content, terms) : info.rgMatch ? 1 : 0
      const score = scoreCandidate({
        rgMatch: info.rgMatch,
        ragScore: info.ragScore,
        termMatches,
        totalTerms: terms.length,
        tags: meta.tags,
        path: p,
        importance: meta.importance,
        accessCount: meta.access_count,
        terms,
      })
      results.push({
        path: p,
        meta,
        score,
        termMatches,
        rgMatch: info.rgMatch,
        ragScore: info.ragScore,
        ragText: info.ragText,
      })
    } catch (e) {
      process.stderr.write(`[memory-search] score read failed ${p}: ${(e as Error).message}\n`)
    }
  }
  results.sort((a, b) => b.score - a.score)
  return results
}

function formatResults(
  results: ResultRow[],
  query: string,
  terms: string[],
  category: string | undefined,
  limit: number,
  memoryDir: string
): string {
  const lines = [`## Results for "${query}" (${results.length} matches)`]
  if (category && results.length > 0) {
    lines.push(`_Filtered to **${category}/**_`)
  }
  if (terms.length > 1) lines.push(`_Searching for: ${terms.join(", ")}_`)

  const filtered = results.filter((r) => r.score >= 0.2)
  const topResults = filtered.length > 0 ? filtered : results.slice(0, 3)
  const MAX = Math.min(limit, 7)
  for (const [i, r] of topResults.slice(0, MAX).entries()) {
    const isDirectHit =
      r.rgMatch && r.ragScore !== undefined && r.ragScore > 0.4 && r.termMatches === terms.length
    const hitLabel = isDirectHit ? " ★ DIRECT HIT" : ""
    const importance = r.meta.importance || "medium"
    lines.push(`${i + 1}. **${r.path}** [${importance}]${hitLabel}`)
    if (r.meta.tags?.length) lines.push(`   Tags: ${r.meta.tags.join(", ")}`)
    if (r.meta.summary) lines.push(`   ${r.meta.summary}`)
    const sources: string[] = []
    if (r.rgMatch) {
      const termInfo = terms.length > 1 ? ` (${r.termMatches}/${terms.length} terms)` : ""
      sources.push(`keyword${termInfo}`)
    }
    if (r.ragScore) sources.push(`semantic: ${r.ragScore.toFixed(2)}`)
    lines.push(`   Match: ${sources.join(" + ")} | score: ${r.score.toFixed(2)}`)
    if (r.meta.related?.length) lines.push(`   Related: ${r.meta.related.join(", ")}`)
    if (r.ragText) lines.push(`   Preview: "...${r.ragText.slice(0, 200).trim()}..."`)
    lines.push("")
  }

  if (topResults.length > 0) {
    lines.push(`_Read the top result: ${topResults[0].path}_`)
    lines.push("")
  }

  const shownPaths = new Set(topResults.slice(0, 7).map((r) => r.path))
  const relatedSuggestions: string[] = []
  for (const r of topResults.slice(0, 7)) {
    for (const rel of r.meta.related || []) {
      if (!shownPaths.has(rel) && !relatedSuggestions.includes(rel)) relatedSuggestions.push(rel)
    }
  }
  const verified: string[] = []
  for (const rel of relatedSuggestions) {
    try {
      if (fs.existsSync(path.join(memoryDir, rel))) verified.push(rel)
    } catch (e) {
      process.stderr.write(`[memory-search] related exists check failed ${rel}: ${(e as Error).message}\n`)
    }
  }
  if (verified.length > 0) {
    lines.push("---")
    lines.push(`**Related files** (not in results — use the Read tool):`)
    for (const rel of verified) lines.push(`  - ${rel}`)
    lines.push("")
  }
  return lines.join("\n")
}

export default tool({
  description:
    "Search memories in the memory directory using both keyword (ripgrep) and semantic (Python rag) search. " +
    "Multi-term queries match files containing ANY search term (OR logic); files matching more terms rank higher. " +
    "Results are summaries only (path, tags, importance, short context). Use the Read tool on the memory path to get the full content. " +
    "WHEN TO SEARCH (do this BEFORE starting work): starting work on any repo, using unfamiliar tool/API, encountering unfamiliar codebase, debugging a problem, looking up people/team info, before writing new memory. " +
    "FALLBACK: if semantic search is unavailable (OPENAI_BASE_URL not set, Python crash, no index), keyword-only results are returned — never fails.",
  args: {
    query: tool.schema.string().describe("Search terms or natural language query"),
    category: tool.schema
      .string()
      .optional()
      .describe("Filter to a category: preferences, repos, technical, people, workflows, snippets, notes"),
    limit: tool.schema.number().optional().describe("Max results (default 15)"),
  },
  async execute(args, context) {
    const memoryDir = resolveMemoryDir()
    const searchDir = args.category ? path.join(memoryDir, args.category) : memoryDir
    const terms = parseSearchTerms(args.query)
    const rgTerms = terms.length > 0 ? terms : [args.query]
    const limit = args.limit ?? 15
    const rgBin = resolveRgBinary({ allowSystemFallback: true })

    if (!fs.existsSync(memoryDir)) {
      return `No memories found for: "${args.query}" (memory directory not initialized — run memory-save to auto-setup)`
    }

    let resultMap = new Map<string, CandInfo>()
    if (rgBin) {
      resultMap = runKeywordSearch(rgBin, rgTerms, searchDir, memoryDir)
    } else {
      process.stderr.write(
        "[memory-search] ripgrep not resolvable — keyword search disabled\n"
      )
    }

    const indexDir = path.join(memoryDir, ".rag")
    const ragHits = runSemanticSearch(args.query, indexDir, context.worktree)
    mergeRagHits(resultMap, ragHits, args.category)

    if (resultMap.size === 0 && args.category) {
      if (rgBin) {
        const fallback = runKeywordSearch(rgBin, rgTerms, memoryDir, memoryDir)
        for (const [k, v] of fallback) resultMap.set(k, v)
      } else {
        process.stderr.write(
          "[memory-search] ripgrep not resolvable — category fallback keyword search disabled\n"
        )
      }
      mergeRagHits(resultMap, ragHits)
    }

    if (resultMap.size === 0) {
      return `No memories found for: "${args.query}"`
    }

    const results = scoreAndCombine(resultMap, terms, memoryDir)
    return formatResults(results, args.query, terms, args.category, limit, memoryDir)
  },
})