brainbank 0.1.0-beta.1

package/src/cli/tui/stats-search.ts

@@ -0,0 +1,262 @@
+/**
+ * BrainBank — Stats TUI Search Session
+ *
+ * Wraps a BrainBank instance to provide staged search pipeline results
+ * for the interactive TUI. Captures each stage (raw → pruned → expanded)
+ * separately so the UI can display them side-by-side.
+ *
+ * Lazy initialization: the brain is created on first search.
+ */
+
+import type { SearchResult, Pruner, Expander, ExpanderManifestItem } from '@/types.ts';
+import type { BrainBank } from '@/brainbank.ts';
+
+import { createBrain } from '@/cli/factory/index.ts';
+import { pruneResults } from '@/lib/prune.ts';
+import { isExpandablePlugin } from '@/plugin.ts';
+
+// ── Types ─────────────────────────────────────────
+
+/** Timings for each pipeline stage (ms). */
+export interface PipelineTimings {
+    init: number;
+    search: number;
+    prune: number;
+    expand: number;
+    total: number;
+}
+
+/** Full result of a staged search pipeline run. */
+export interface SearchPipelineResult {
+    /** All results from vector search (before pruning). */
+    raw: SearchResult[];
+    /** Results after LLM pruner filtered noise. Same order as pruner returned. */
+    pruned: SearchResult[];
+    /** Results dropped by pruner (raw minus pruned). */
+    dropped: SearchResult[];
+    /** Additional results discovered by LLM expander. */
+    expanded: SearchResult[];
+    /** Pipeline stage timings. */
+    timings: PipelineTimings;
+    /** Name of the active pruner (or null if none). */
+    prunerName: string | null;
+    /** Name of the active expander (or null if none). */
+    expanderName: string | null;
+}
+
+/** Source filter option — discovered from registered plugins. */
+export interface SourceOption {
+    /** Source key: 'code', 'git', 'docs', etc. */
+    key: string;
+    /** Display label. */
+    label: string;
+    /** Whether this source is enabled for the current search. */
+    enabled: boolean;
+    /** Max results for this source (default: 20). */
+    k: number;
+}
+
+// ── Session ───────────────────────────────────────
+
+export class BrainSearchSession {
+    private _brain: BrainBank | null = null;
+    private _repoPath: string;
+    private _initPromise: Promise<void> | null = null;
+    private _sources: SourceOption[] = [];
+
+    constructor(repoPath: string) {
+        this._repoPath = repoPath;
+    }
+
+    /** Whether the brain has been initialized. */
+    get initialized(): boolean {
+        return this._brain !== null;
+    }
+
+    /** Available source filters (populated after init). */
+    get sources(): readonly SourceOption[] {
+        return this._sources;
+    }
+
+    /** Pruner name (or null). */
+    get prunerName(): string | null {
+        return this._brain?.config.pruner?.constructor?.name ?? null;
+    }
+
+    /** Expander name (or null). */
+    get expanderName(): string | null {
+        return this._brain?.config.expander?.constructor?.name ?? null;
+    }
+
+    /**
+     * Initialize the BrainBank instance (lazy, idempotent).
+     * Creates the brain, loads HNSW + FTS indices.
+     */
+    async init(): Promise<void> {
+        if (this._brain) return;
+        if (this._initPromise) return this._initPromise;
+
+        this._initPromise = this._doInit();
+        return this._initPromise;
+    }
+
+    private async _doInit(): Promise<void> {
+        this._brain = await createBrain(this._repoPath);
+        await this._brain.initialize();
+
+        // Discover installed sources from plugin names (no 'all' — user selects individually)
+        this._sources = [];
+        const seen = new Set<string>();
+        for (const name of this._brain.plugins) {
+            const base = name.split(':')[0];
+            if (seen.has(base)) continue;
+            seen.add(base);
+            this._sources.push({
+                key: base,
+                label: base.charAt(0).toUpperCase() + base.slice(1),
+                enabled: true,
+                k: 20,
+            });
+        }
+    }
+
+    /**
+     * Run the full search pipeline and return staged results.
+     *
+     * @param query  - Search query
+     * @param activeSourceKeys - Set of active source keys (e.g. {'code', 'git'})
+     * @param usePruner - Whether to run the pruner stage
+     * @param useExpander - Whether to run the expander stage
+     */
+    async search(
+        query: string,
+        activeSourceKeys?: Set<string>,
+        usePruner = true,
+        useExpander = true,
+    ): Promise<SearchPipelineResult> {
+        if (!this._brain) throw new Error('BrainSearchSession not initialized');
+
+        const tTotal = Date.now();
+        const pruner = usePruner ? (this._brain.config.pruner as Pruner | undefined) : undefined;
+        const expander = useExpander ? (this._brain.config.expander as Expander | undefined) : undefined;
+
+        // Build sources filter — always pass K values so we don't fall back to DEFAULT_K=6
+        const sources: Record<string, number> = {};
+        for (const src of this._sources) {
+            const enabled = activeSourceKeys ? activeSourceKeys.has(src.key) : src.enabled;
+            sources[src.key] = enabled ? src.k : 0;
+        }
+
+        // Stage 1: Vector search
+        const tSearch0 = Date.now();
+        const raw = await this._brain.search(query, {
+            sources,
+            source: 'cli',
+        });
+        const tSearch = Date.now() - tSearch0;
+
+        // Stage 2: Prune
+        let pruned = raw;
+        let dropped: SearchResult[] = [];
+        let tPrune = 0;
+        if (pruner && raw.length > 1) {
+            const pt0 = Date.now();
+            pruned = await pruneResults(query, raw, pruner);
+            tPrune = Date.now() - pt0;
+            const prunedSet = new Set(pruned);
+            dropped = raw.filter(r => !prunedSet.has(r));
+        }
+
+        // Stage 3: Expand — filter results to only include active sources
+        let expanded: SearchResult[] = [];
+        let tExpand = 0;
+        if (expander && pruned.length > 0) {
+            const et0 = Date.now();
+            let expandedRaw = await this._runExpansion(query, pruned, expander);
+            // Filter expanded results by active sources
+            if (activeSourceKeys) {
+                expandedRaw = expandedRaw.filter(r => activeSourceKeys.has(r.type));
+            }
+            expanded = expandedRaw;
+            tExpand = Date.now() - et0;
+        }
+
+        return {
+            raw,
+            pruned,
+            dropped,
+            expanded,
+            timings: {
+                init: 0,
+                search: tSearch,
+                prune: tPrune,
+                expand: tExpand,
+                total: Date.now() - tTotal,
+            },
+            prunerName: usePruner ? (pruner?.constructor?.name ?? null) : null,
+            expanderName: useExpander ? (expander?.constructor?.name ?? null) : null,
+        };
+    }
+
+    /** Run LLM expansion — mirrors ContextBuilder._expand logic. */
+    private async _runExpansion(
+        query: string,
+        results: SearchResult[],
+        expander: Expander,
+    ): Promise<SearchResult[]> {
+        if (!this._brain) return [];
+
+        // Collect file paths and IDs already in results
+        const excludeFilePaths = [...new Set(
+            results.filter(r => r.filePath).map(r => r.filePath as string),
+        )];
+        const excludeIds: number[] = [];
+        for (const r of results) {
+            const meta = r.metadata as Record<string, unknown> | undefined;
+            const id = meta?.id as number | undefined;
+            if (id !== undefined) excludeIds.push(id);
+        }
+
+        // Build manifest from ExpandablePlugins
+        const manifest: ExpanderManifestItem[] = [];
+        let resolver: ((ids: number[]) => SearchResult[]) | undefined;
+
+        // Access plugins via the brain's plugin list
+        // We need the registry... which is private. Use a workaround:
+        // BrainBank exposes `plugins` (names) but not instances.
+        // However, we can access it through the search functionality.
+        // The cleanest approach: use the brain's internal registry via
+        // a property we know exists on the class.
+        const registry = (this._brain as unknown as { _registry: { all: unknown[] } })._registry;
+        if (registry) {
+            for (const mod of registry.all) {
+                if (!isExpandablePlugin(mod as never)) continue;
+                const plugin = mod as { buildManifest: (fp: string[], ids: number[], rfp?: string[]) => ExpanderManifestItem[]; resolveChunks: (ids: number[]) => SearchResult[] };
+                manifest.push(...plugin.buildManifest(excludeFilePaths, excludeIds, excludeFilePaths));
+                if (!resolver) {
+                    resolver = (ids: number[]) => plugin.resolveChunks(ids);
+                }
+            }
+        }
+
+        if (manifest.length === 0 || !resolver) return [];
+
+        try {
+            const expandResult = await expander.expand(query, excludeIds, manifest);
+            if (expandResult.ids.length === 0) return [];
+            return resolver(expandResult.ids);
+        } catch {
+            // Fail-open: expansion errors are non-fatal
+            return [];
+        }
+    }
+
+    /** Cleanup. */
+    close(): void {
+        if (this._brain) {
+            this._brain.close();
+            this._brain = null;
+        }
+        this._initPromise = null;
+    }
+}