@rkarim08/sia 1.0.0

package/src/retrieval/query-classifier.ts

@@ -0,0 +1,106 @@
+// Module: query-classifier — Local vs global query routing + task-type boosts
+
+import type { SiaDb } from "@/graph/db-interface";
+
+export type QueryMode = "local" | "global";
+
+export interface ClassificationResult {
+	mode: QueryMode;
+	globalUnavailable: boolean;
+}
+
+const GLOBAL_KEYWORDS: string[] = [
+	"architecture",
+	"overview",
+	"explain",
+	"structure",
+	"high-level",
+	"design",
+	"modules",
+	"subsystems",
+];
+
+const LOCAL_KEYWORDS: string[] = [
+	"function",
+	"class",
+	"method",
+	"variable",
+	"import",
+	"error",
+	"bug",
+	"fix",
+	"implement",
+	"where is",
+	"how does",
+	"what does",
+];
+
+/**
+ * Classify a query as local (three-stage pipeline) or global (community summaries).
+ *
+ * Keyword-based classification: count matches against global and local keyword
+ * lists, default to local when tied. If the graph is too small for meaningful
+ * community summaries (fewer than `config.communityMinGraphSize` active
+ * entities), force local and set `globalUnavailable: true`.
+ */
+export async function classifyQuery(
+	db: SiaDb,
+	query: string,
+	config: { communityMinGraphSize: number },
+): Promise<ClassificationResult> {
+	const lower = query.toLowerCase();
+
+	let globalScore = 0;
+	for (const kw of GLOBAL_KEYWORDS) {
+		if (lower.includes(kw)) {
+			globalScore++;
+		}
+	}
+
+	let localScore = 0;
+	for (const kw of LOCAL_KEYWORDS) {
+		if (lower.includes(kw)) {
+			localScore++;
+		}
+	}
+
+	// Default to local when tied (localScore >= globalScore means local wins on tie)
+	let mode: QueryMode = globalScore > localScore ? "global" : "local";
+	let globalUnavailable = false;
+
+	// Check graph size — force local if too few entities for community summaries
+	if (mode === "global") {
+		const result = await db.execute(
+			"SELECT COUNT(*) AS cnt FROM graph_nodes WHERE t_valid_until IS NULL AND archived_at IS NULL",
+		);
+		const count = Number((result.rows[0] as { cnt: number }).cnt);
+		if (count < config.communityMinGraphSize) {
+			mode = "local";
+			globalUnavailable = true;
+		}
+	}
+
+	return { mode, globalUnavailable };
+}
+
+/**
+ * Task-type boost vectors: maps task type strings to sets of entity types
+ * that should receive a scoring boost during reranking.
+ */
+export const TASK_TYPE_BOOSTS: Record<string, Set<string>> = {
+	"bug-fix": new Set(["Bug", "Solution"]),
+	regression: new Set(["Bug", "Solution"]),
+	feature: new Set(["Concept", "Decision"]),
+	review: new Set(["Convention"]),
+};
+
+/**
+ * Package-path boost: returns 0.15 when the entity's package matches the
+ * active package, 0 otherwise.
+ */
+export function packagePathBoost(entityPkg: string | null, activePkg: string | null): number {
+	if (entityPkg != null && activePkg != null && entityPkg === activePkg) {
+		return 0.15;
+	}
+	return 0;
+}

package/src/retrieval/reranker.ts

@@ -0,0 +1,156 @@
+// Module: reranker — RRF combination + trust-weighted scoring
+
+import type { SiaDb } from "@/graph/db-interface";
+import type { Entity } from "@/graph/entities";
+import type { SiaSearchResult } from "@/mcp/tools/sia-search";
+import { packagePathBoost, TASK_TYPE_BOOSTS } from "@/retrieval/query-classifier";
+
+/** A candidate with an entity ID and a score from a single retrieval signal. */
+export interface RankedCandidate {
+	entityId: string;
+	score: number;
+}
+
+/** Options for the rerank function. */
+export interface RerankOpts {
+	/** Task type for task-type boosting (e.g. "bug-fix", "feature"). */
+	taskType?: string;
+	/** Active package path for same-package boosting. */
+	packagePath?: string;
+	/** If true, exclude Tier 4 entities. */
+	paranoid?: boolean;
+	/** Maximum number of results to return. */
+	limit?: number;
+	/** If true, include extraction_method in results. */
+	includeProvenance?: boolean;
+}
+
+/** Trust weights keyed by tier number (1-4). No index-0. */
+const TRUST_WEIGHTS: Record<number, number> = {
+	1: 1.0,
+	2: 0.9,
+	3: 0.7,
+	4: 0.5,
+};
+
+/**
+ * Combine multiple ranked candidate lists via Reciprocal Rank Fusion (k=60).
+ *
+ * Each list is sorted by score DESC. For each entity in each list,
+ * the contribution is `1 / (k + rank + 1)` where rank is 0-based.
+ * Scores are summed across all lists.
+ */
+export function rrfCombine(...lists: RankedCandidate[][]): Map<string, number> {
+	const k = 60;
+	const scores = new Map<string, number>();
+
+	for (const list of lists) {
+		// Sort by score descending to establish rank order
+		const sorted = [...list].sort((a, b) => b.score - a.score);
+
+		for (let rank = 0; rank < sorted.length; rank++) {
+			const candidate = sorted[rank];
+			const rrfScore = 1 / (k + rank + 1);
+			const current = scores.get(candidate.entityId) ?? 0;
+			scores.set(candidate.entityId, current + rrfScore);
+		}
+	}
+
+	return scores;
+}
+
+/** Batch size for fetching entities from the database. */
+const ENTITY_BATCH_SIZE = 500;
+
+/**
+ * Rerank entities by combining RRF scores with trust weights, importance,
+ * confidence, task-type boosts, and package-path boosts.
+ *
+ * Formula: rrf_score * importance * confidence * trust_weight[tier] * (1 + task_boost * 0.3) + package_boost
+ */
+export async function rerank(
+	db: SiaDb,
+	rrfScores: Map<string, number>,
+	opts?: RerankOpts,
+): Promise<SiaSearchResult[]> {
+	if (rrfScores.size === 0) {
+		return [];
+	}
+
+	const entityIds = Array.from(rrfScores.keys());
+
+	// Fetch entities in batches of 500
+	const entities = new Map<string, Entity>();
+	for (let i = 0; i < entityIds.length; i += ENTITY_BATCH_SIZE) {
+		const batch = entityIds.slice(i, i + ENTITY_BATCH_SIZE);
+		const placeholders = batch.map(() => "?").join(", ");
+		const sql = `SELECT * FROM graph_nodes WHERE id IN (${placeholders}) AND t_valid_until IS NULL AND archived_at IS NULL`;
+		const result = await db.execute(sql, batch);
+		for (const row of result.rows) {
+			entities.set(row.id as string, row as unknown as Entity);
+		}
+	}
+
+	// Determine task-type boosted entity types
+	const boostedTypes: Set<string> | undefined = opts?.taskType
+		? TASK_TYPE_BOOSTS[opts.taskType]
+		: undefined;
+
+	// Score and filter
+	const scored: Array<{ entity: Entity; finalScore: number }> = [];
+
+	for (const [entityId, rrfScore] of rrfScores) {
+		const entity = entities.get(entityId);
+		if (!entity) {
+			// Entity was invalidated, archived, or doesn't exist
+			continue;
+		}
+
+		// Paranoid filter: remove Tier 4
+		if (opts?.paranoid && entity.trust_tier === 4) {
+			continue;
+		}
+
+		const trustWeight = TRUST_WEIGHTS[entity.trust_tier] ?? 0.5;
+		const taskBoost = boostedTypes?.has(entity.type) ? 1 : 0;
+		const pkgBoost = packagePathBoost(entity.package_path, opts?.packagePath ?? null);
+
+		const finalScore =
+			rrfScore * entity.importance * entity.confidence * trustWeight * (1 + taskBoost * 0.3) +
+			pkgBoost;
+
+		scored.push({ entity, finalScore });
+	}
+
+	// Sort by finalScore DESC
+	scored.sort((a, b) => b.finalScore - a.finalScore);
+
+	// Apply limit
+	const limit = opts?.limit ?? 15;
+	const top = scored.slice(0, limit);
+
+	// Map to SiaSearchResult
+	return top.map(({ entity, finalScore: _finalScore }) => {
+		const base: SiaSearchResult = {
+			id: entity.id,
+			type: entity.type,
+			name: entity.name,
+			summary: entity.summary,
+			content: entity.content,
+			trust_tier: entity.trust_tier,
+			confidence: entity.confidence,
+			importance: entity.importance,
+			tags: entity.tags,
+			file_paths: entity.file_paths,
+			conflict_group_id: entity.conflict_group_id ?? null,
+			t_valid_from: entity.t_valid_from ?? null,
+			source_repo_name: null,
+		};
+
+		if (opts?.includeProvenance) {
+			base.extraction_method = entity.extraction_method ?? null;
+		}
+
+		return base;
+	});
+}

package/src/retrieval/search.ts

@@ -0,0 +1,236 @@
+// Module: search — Three-stage pipeline orchestration
+//
+// Stage 1: Parallel BM25 + graph traversal + vector search
+// Stage 2: 1-hop neighbor expansion for candidates
+// Stage 3: RRF combination + trust-weighted reranking
+// Global queries bypass the pipeline and return community summaries.
+
+import type { Embedder } from "@/capture/embedder";
+import type { SiaDb } from "@/graph/db-interface";
+import type { SiaSearchResult } from "@/mcp/tools/sia-search";
+import { bm25Search } from "@/retrieval/bm25-search";
+import { graphTraversalSearch } from "@/retrieval/graph-traversal";
+import { classifyQuery } from "@/retrieval/query-classifier";
+import { type RankedCandidate, rerank, rrfCombine } from "@/retrieval/reranker";
+import { vectorSearch } from "@/retrieval/vector-search";
+
+/** Options accepted by hybridSearch. */
+export interface SearchOptions {
+	query: string;
+	taskType?: string;
+	nodeTypes?: string[];
+	packagePath?: string;
+	paranoid?: boolean;
+	limit?: number;
+	includeProvenance?: boolean;
+	communityMinGraphSize?: number;
+}
+
+/** Result returned by hybridSearch. */
+export interface SearchResult {
+	results: SiaSearchResult[];
+	mode: "local" | "global";
+	globalUnavailable: boolean;
+}
+
+/** Default minimum graph size before community summaries are available. */
+const DEFAULT_COMMUNITY_MIN_GRAPH_SIZE = 100;
+
+/**
+ * Three-stage hybrid retrieval pipeline.
+ *
+ * 1. Classify query as local or global.
+ * 2. If global, return community summaries from the `communities` table.
+ * 3. Stage 1: parallel BM25 + graph traversal + vector search.
+ * 4. Stage 2: expand 1-hop neighbors for every candidate.
+ * 5. Stage 3: RRF combine + trust-weighted rerank.
+ * 6. Post-filter by nodeTypes if specified.
+ * 7. Attach extraction_method if includeProvenance is set.
+ *
+ * The `embedder` parameter is nullable -- when null, vector search is skipped
+ * and the pipeline runs on BM25 + graph traversal only.
+ */
+export async function hybridSearch(
+	db: SiaDb,
+	embedder: Embedder | null,
+	opts: SearchOptions,
+): Promise<SearchResult> {
+	const limit = opts.limit ?? 15;
+	const communityMinGraphSize = opts.communityMinGraphSize ?? DEFAULT_COMMUNITY_MIN_GRAPH_SIZE;
+
+	// --- Classify query ---------------------------------------------------
+	const classification = await classifyQuery(db, opts.query, {
+		communityMinGraphSize,
+	});
+
+	// --- Global mode: return community summaries --------------------------
+	if (classification.mode === "global") {
+		const communities = await fetchCommunitySummaries(db, limit);
+		return {
+			results: communities,
+			mode: "global",
+			globalUnavailable: false,
+		};
+	}
+
+	// --- Stage 1: parallel retrieval signals ------------------------------
+	const searchOpts = {
+		limit: limit * 3, // over-fetch to leave room for reranking
+		paranoid: opts.paranoid,
+		packagePath: opts.packagePath,
+	};
+
+	const [bm25Results, graphResults, vecResults] = await Promise.all([
+		bm25Search(db, opts.query, searchOpts),
+		graphTraversalSearch(db, opts.query, searchOpts),
+		embedder ? vectorSearch(db, opts.query, embedder, searchOpts) : Promise.resolve([]),
+	]);
+
+	// --- Stage 2: expand 1-hop neighbors ----------------------------------
+	const expandedGraphResults = await expandNeighbors(db, graphResults, opts.paranoid);
+
+	// --- Stage 3: RRF combine + rerank ------------------------------------
+	const bm25Candidates: RankedCandidate[] = bm25Results.map((r) => ({
+		entityId: r.entityId,
+		score: r.score,
+	}));
+	const graphCandidates: RankedCandidate[] = expandedGraphResults.map((r) => ({
+		entityId: r.entityId,
+		score: r.score,
+	}));
+	const vecCandidates: RankedCandidate[] = vecResults.map((r) => ({
+		entityId: r.entityId,
+		score: r.score,
+	}));
+
+	const rrfScores = rrfCombine(bm25Candidates, graphCandidates, vecCandidates);
+
+	let results = await rerank(db, rrfScores, {
+		taskType: opts.taskType,
+		packagePath: opts.packagePath,
+		paranoid: opts.paranoid,
+		limit,
+		includeProvenance: opts.includeProvenance,
+	});
+
+	// --- Post-filter by nodeTypes ------------------------------------------
+	if (opts.nodeTypes && opts.nodeTypes.length > 0) {
+		const allowed = new Set(opts.nodeTypes);
+		results = results.filter((r) => allowed.has(r.type));
+	}
+
+	// --- Provenance --------------------------------------------------------
+	if (opts.includeProvenance) {
+		await attachProvenance(db, results);
+	}
+
+	return {
+		results,
+		mode: "local",
+		globalUnavailable: classification.globalUnavailable,
+	};
+}
+
+// ---------------------------------------------------------------------------
+// Internal helpers
+// ---------------------------------------------------------------------------
+
+/**
+ * Fetch community summaries for global-mode queries.
+ * Returns communities that have a non-NULL summary, ordered by member_count DESC.
+ */
+async function fetchCommunitySummaries(db: SiaDb, limit: number): Promise<SiaSearchResult[]> {
+	const result = await db.execute(
+		"SELECT * FROM communities WHERE summary IS NOT NULL ORDER BY member_count DESC LIMIT ?",
+		[limit],
+	);
+
+	return (result.rows as Record<string, unknown>[]).map((row) => ({
+		id: row.id as string,
+		type: "Community",
+		name: row.id as string,
+		summary: (row.summary as string) ?? "",
+		content: (row.summary as string) ?? "",
+		trust_tier: 1,
+		confidence: 1.0,
+		importance: 1.0,
+		tags: "[]",
+		file_paths: "[]",
+		conflict_group_id: null,
+		t_valid_from: null,
+		source_repo_name: null,
+	}));
+}
+
+/**
+ * Stage 2: expand 1-hop neighbors for each candidate entity.
+ *
+ * For each entity in the input list, query the `edges` table for active
+ * 1-hop neighbors. Neighbors not already present in the result set are
+ * added at score 0.7.
+ */
+async function expandNeighbors(
+	db: SiaDb,
+	results: Array<{ entityId: string; score: number }>,
+	paranoid?: boolean,
+): Promise<Array<{ entityId: string; score: number }>> {
+	const scoreMap = new Map<string, number>();
+
+	// Seed with existing results
+	for (const r of results) {
+		const existing = scoreMap.get(r.entityId);
+		if (existing === undefined || r.score > existing) {
+			scoreMap.set(r.entityId, r.score);
+		}
+	}
+
+	// Expand each candidate
+	const candidateIds = results.map((r) => r.entityId);
+	for (const entityId of candidateIds) {
+		const edgeResult = await db.execute(
+			"SELECT from_id, to_id FROM graph_edges WHERE (from_id = ? OR to_id = ?) AND t_valid_until IS NULL",
+			[entityId, entityId],
+		);
+
+		for (const row of edgeResult.rows) {
+			const fromId = row.from_id as string;
+			const toId = row.to_id as string;
+			const neighborId = fromId === entityId ? toId : fromId;
+
+			// Skip if already in the result set
+			if (scoreMap.has(neighborId)) continue;
+
+			// Validate the neighbor is active (and paranoid-safe)
+			const paranoidClause = paranoid ? " AND trust_tier < 4" : "";
+			const check = await db.execute(
+				`SELECT id FROM graph_nodes WHERE id = ? AND t_valid_until IS NULL AND archived_at IS NULL${paranoidClause}`,
+				[neighborId],
+			);
+			if (check.rows.length === 0) continue;
+
+			scoreMap.set(neighborId, 0.7);
+		}
+	}
+
+	return [...scoreMap.entries()].map(([entityId, score]) => ({
+		entityId,
+		score,
+	}));
+}
+
+/**
+ * Attach extraction_method to results that don't already have it set.
+ * Only queries the DB for results where extraction_method is undefined.
+ */
+async function attachProvenance(db: SiaDb, results: SiaSearchResult[]): Promise<void> {
+	for (const result of results) {
+		if (result.extraction_method === undefined) {
+			const row = await db.execute("SELECT extraction_method FROM graph_nodes WHERE id = ?", [
+				result.id,
+			]);
+			if (row.rows.length > 0) {
+				result.extraction_method = (row.rows[0].extraction_method as string | null) ?? null;
+			}
+		}
+	}
+}

package/src/retrieval/throttle.ts

@@ -0,0 +1,102 @@
+// Module: throttle — Progressive rate limiting for MCP tool calls
+
+import type { SiaDb } from "@/graph/db-interface";
+
+export type ThrottleMode = "normal" | "reduced" | "blocked";
+
+export interface ThrottleResult {
+	mode: ThrottleMode;
+	callCount: number;
+	warning?: string;
+}
+
+export interface ThrottleConfig {
+	normalMax: number;
+	reducedMax: number;
+}
+
+const DEFAULT_THROTTLE_CONFIG: ThrottleConfig = {
+	normalMax: 3,
+	reducedMax: 8,
+};
+
+export const THROTTLED_TOOLS = new Set([
+	"sia_search",
+	"sia_execute",
+	"sia_execute_file",
+	"sia_fetch_and_index",
+	"sia_by_file",
+	"sia_expand",
+	"sia_at_time",
+	"sia_backlinks",
+]);
+
+/**
+ * Progressive throttle backed by the search_throttle table in graph.db.
+ *
+ * Thresholds (inclusive):
+ *  - count <= normalMax  → "normal"
+ *  - count <= reducedMax → "reduced" (with warning)
+ *  - count >  reducedMax → "blocked" (with warning mentioning sia_batch_execute)
+ */
+export class ProgressiveThrottle {
+	private config: ThrottleConfig;
+
+	constructor(
+		private db: SiaDb,
+		config?: Partial<ThrottleConfig>,
+	) {
+		this.config = { ...DEFAULT_THROTTLE_CONFIG, ...config };
+	}
+
+	/**
+	 * Record a tool call for the given session/tool pair and return the throttle mode.
+	 */
+	async check(sessionId: string, toolName: string): Promise<ThrottleResult> {
+		const now = Date.now();
+
+		// Upsert: insert or increment call_count
+		await this.db.execute(
+			`INSERT INTO search_throttle (session_id, tool_name, call_count, last_called_at)
+       VALUES (?, ?, 1, ?)
+       ON CONFLICT(session_id, tool_name) DO UPDATE SET
+         call_count = call_count + 1,
+         last_called_at = ?`,
+			[sessionId, toolName, now, now],
+		);
+
+		// Read back the updated count
+		const { rows } = await this.db.execute(
+			"SELECT call_count FROM search_throttle WHERE session_id = ? AND tool_name = ?",
+			[sessionId, toolName],
+		);
+
+		const callCount = (rows[0]?.call_count as number) ?? 1;
+		const { normalMax, reducedMax } = this.config;
+
+		if (callCount <= normalMax) {
+			return { mode: "normal", callCount };
+		}
+
+		if (callCount <= reducedMax) {
+			return {
+				mode: "reduced",
+				callCount,
+				warning: `Reducing results (${callCount} calls). Consider sia_batch_execute for batch operations.`,
+			};
+		}
+
+		return {
+			mode: "blocked",
+			callCount,
+			warning: `Tool blocked for this session (${callCount} calls). Use sia_batch_execute instead.`,
+		};
+	}
+
+	/**
+	 * Clear all throttle entries for the given session.
+	 */
+	async reset(sessionId: string): Promise<void> {
+		await this.db.execute("DELETE FROM search_throttle WHERE session_id = ?", [sessionId]);
+	}
+}