From a219affd91152491198378fbb3479c380aced04d Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 26 Mar 2026 17:30:21 +0000 Subject: [PATCH] refactor(mcp): abstract search behind SearchProvider interface MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce a pluggable search architecture with two backends: - VectorSearchProvider: Sentence-transformer embeddings (all-MiniLM-L6-v2) for true semantic search. Best quality, requires model download. - FuseSearchProvider: Fuse.js fuzzy text matching. Zero model download, instant startup, lightweight fallback. Auto-detection: if @xenova/transformers is available, uses vector; otherwise falls back to fuse. Override via SEARCH_BACKEND=vector|fuse env var. Architecture: SearchProvider (interface) ← tools.ts programs against this ├── VectorSearchProvider ← embeddings + cosine similarity └── FuseSearchProvider ← fuse.js weighted fuzzy match createSearchProvider() ← factory with auto-detect Also refactored knowledge-base.ts to be provider-agnostic (just loads documents) and removed the old monolithic vector-store.ts. https://claude.ai/code/session_01TBPsvvMBue9Ft7j3CfD4S2 --- .../autopilot-monitor-mcp/package-lock.json | 14 ++- .../autopilot-monitor-mcp/package.json | 3 +- .../src/fuse-search-provider.ts | 57 +++++++++++ .../autopilot-monitor-mcp/src/index.ts | 21 ++-- .../src/knowledge-base.ts | 87 ++++------------- .../src/search-factory.ts | 46 +++++++++ .../src/search-provider.ts | 54 +++++++++++ .../autopilot-monitor-mcp/src/tools.ts | 95 +++++++++--------- .../src/vector-search-provider.ts | 96 +++++++++++++++++++ .../autopilot-monitor-mcp/src/vector-store.ts | 61 ------------ 10 files changed, 340 insertions(+), 194 deletions(-) create mode 100644 src/McpServer/autopilot-monitor-mcp/src/fuse-search-provider.ts create mode 100644 src/McpServer/autopilot-monitor-mcp/src/search-factory.ts create mode 100644 src/McpServer/autopilot-monitor-mcp/src/search-provider.ts create mode 100644 src/McpServer/autopilot-monitor-mcp/src/vector-search-provider.ts delete mode 100644 src/McpServer/autopilot-monitor-mcp/src/vector-store.ts diff --git a/src/McpServer/autopilot-monitor-mcp/package-lock.json b/src/McpServer/autopilot-monitor-mcp/package-lock.json index 41862f41..75853ce0 100644 --- a/src/McpServer/autopilot-monitor-mcp/package-lock.json +++ b/src/McpServer/autopilot-monitor-mcp/package-lock.json @@ -1,16 +1,17 @@ { "name": "autopilot-monitor-mcp", - "version": "1.1.0", + "version": "1.2.0", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "autopilot-monitor-mcp", - "version": "1.1.0", + "version": "1.2.0", "dependencies": { "@anthropic-ai/sdk": "^0.80.0", "@modelcontextprotocol/sdk": "^1.0.0", "@xenova/transformers": "^2.17.0", + "fuse.js": "^7.1.0", "zod": "^3.22.0" }, "bin": { @@ -1420,6 +1421,15 @@ "url": "https://github.com/sponsors/ljharb" } }, + "node_modules/fuse.js": { + "version": "7.1.0", + "resolved": "https://registry.npmjs.org/fuse.js/-/fuse.js-7.1.0.tgz", + "integrity": "sha512-trLf4SzuuUxfusZADLINj+dE8clK1frKdmqiJNb1Es75fmI5oY6X2mxLVUciLLjxqw/xr72Dhy+lER6dGd02FQ==", + "license": "Apache-2.0", + "engines": { + "node": ">=10" + } + }, "node_modules/get-intrinsic": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/get-intrinsic/-/get-intrinsic-1.3.0.tgz", diff --git a/src/McpServer/autopilot-monitor-mcp/package.json b/src/McpServer/autopilot-monitor-mcp/package.json index ad5beac4..ba60345f 100644 --- a/src/McpServer/autopilot-monitor-mcp/package.json +++ b/src/McpServer/autopilot-monitor-mcp/package.json @@ -1,6 +1,6 @@ { "name": "autopilot-monitor-mcp", - "version": "1.1.0", + "version": "1.2.0", "description": "MCP server for Autopilot-Monitor — enables AI agents to query enrollment sessions, device properties, events and vulnerabilities", "type": "module", "bin": { @@ -15,6 +15,7 @@ "@anthropic-ai/sdk": "^0.80.0", "@modelcontextprotocol/sdk": "^1.0.0", "@xenova/transformers": "^2.17.0", + "fuse.js": "^7.1.0", "zod": "^3.22.0" }, "devDependencies": { diff --git a/src/McpServer/autopilot-monitor-mcp/src/fuse-search-provider.ts b/src/McpServer/autopilot-monitor-mcp/src/fuse-search-provider.ts new file mode 100644 index 00000000..01b1b320 --- /dev/null +++ b/src/McpServer/autopilot-monitor-mcp/src/fuse-search-provider.ts @@ -0,0 +1,57 @@ +/** + * SearchProvider backed by Fuse.js — lightweight fuzzy text search. + * + * Pros: Zero model download, instant startup, tiny memory footprint. + * Cons: Token/substring matching only — no true semantic understanding. + * "timeout" will NOT match "waiting exceeded" (use vector for that). + */ + +import Fuse from 'fuse.js'; +import type { SearchDocument, SearchOptions, SearchProvider, SearchResult } from './search-provider.js'; + +export class FuseSearchProvider implements SearchProvider { + readonly name = 'fuse'; + private documents: SearchDocument[] = []; + private fuse: Fuse | null = null; + + get size(): number { + return this.documents.length; + } + + async index(docs: SearchDocument[]): Promise { + this.documents.push(...docs); + this.fuse = new Fuse(this.documents, { + keys: [ + { name: 'text', weight: 0.7 }, + { name: 'id', weight: 0.15 }, + { name: 'metadata.title', weight: 0.1 }, + { name: 'metadata.category', weight: 0.05 }, + ], + includeScore: true, + threshold: 0.5, + ignoreLocation: true, + useExtendedSearch: true, + }); + } + + async search(query: string, options: SearchOptions = {}): Promise { + const { topK = 5, minScore = 0.3 } = options; + + if (!this.fuse || this.documents.length === 0) { + return []; + } + + const fuseResults = this.fuse.search(query, { limit: topK * 2 }); + + return fuseResults + .map((r) => ({ + id: r.item.id, + text: r.item.text, + metadata: r.item.metadata, + // Fuse score is 0 = perfect, 1 = no match. Invert to 0..1 where 1 = best. + score: 1 - (r.score ?? 1), + })) + .filter((r) => r.score >= minScore) + .slice(0, topK); + } +} diff --git a/src/McpServer/autopilot-monitor-mcp/src/index.ts b/src/McpServer/autopilot-monitor-mcp/src/index.ts index 05a3faf0..b1aa0a26 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/index.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/index.ts @@ -4,22 +4,27 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js'; import { registerTools } from './tools.js'; import { registerResources } from './resources.js'; -import { buildKnowledgeBase } from './knowledge-base.js'; +import { loadKnowledgeDocs } from './knowledge-base.js'; +import { createSearchProvider } from './search-factory.js'; const __dirname = dirname(fileURLToPath(import.meta.url)); -// Resolve the rules directory (relative to project root) const RULES_DIR = process.env.RULES_DIR ?? resolve(__dirname, '..', '..', '..', '..', 'rules'); const server = new McpServer({ name: 'autopilot-monitor', - version: '1.1.0', + version: '1.2.0', }); -// Build the knowledge base (embeddings) before registering tools, -// so the search_knowledge tool has immediate access. -console.error('Initializing vector search knowledge base…'); -const knowledgeBase = await buildKnowledgeBase(RULES_DIR); +// Initialize the knowledge base search provider +console.error('Loading knowledge base documents…'); +const docs = await loadKnowledgeDocs(RULES_DIR); + +console.error(`Initializing search provider (${docs.length} documents)…`); +const knowledgeBase = await createSearchProvider(); +await knowledgeBase.index(docs); + +console.error(`Search provider ready: ${knowledgeBase.name} — ${knowledgeBase.size} documents indexed.`); registerTools(server, knowledgeBase); registerResources(server); @@ -29,4 +34,4 @@ await server.connect(transport); console.error('Autopilot-Monitor MCP Server running (stdio transport)'); console.error(`API URL: ${process.env.AUTOPILOT_API_URL ?? 'https://autopilotmonitor-api.azurewebsites.net'}`); -console.error(`Knowledge base: ${knowledgeBase.size} documents indexed from ${RULES_DIR}`); +console.error(`Search backend: ${knowledgeBase.name} | Documents: ${knowledgeBase.size} | Rules dir: ${RULES_DIR}`); diff --git a/src/McpServer/autopilot-monitor-mcp/src/knowledge-base.ts b/src/McpServer/autopilot-monitor-mcp/src/knowledge-base.ts index 570c8d25..dd159d4c 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/knowledge-base.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/knowledge-base.ts @@ -1,56 +1,16 @@ /** * Loads analysis rules, gather rules, and IME log patterns from the rules/ - * directory, computes embeddings via a local transformer model, and populates - * the in-memory vector store for semantic search. + * directory and indexes them into a SearchProvider for semantic or fuzzy search. + * + * This module is backend-agnostic — it only deals with document loading and + * text preparation. The actual search strategy is determined by the provider. */ import { readdir, readFile } from 'node:fs/promises'; import { join } from 'node:path'; -import { pipeline, type FeatureExtractionPipeline } from '@xenova/transformers'; -import { VectorStore, type VectorDocument } from './vector-store.js'; +import type { SearchDocument } from './search-provider.js'; -const MODEL_NAME = 'Xenova/all-MiniLM-L6-v2'; - -let embedder: FeatureExtractionPipeline | null = null; - -async function getEmbedder(): Promise { - if (!embedder) { - embedder = await pipeline('feature-extraction', MODEL_NAME, { - quantized: true, - }) as FeatureExtractionPipeline; - } - return embedder; -} - -/** Compute a normalized embedding vector for a text string. */ -export async function embed(text: string): Promise { - const model = await getEmbedder(); - const output = await model(text, { pooling: 'mean', normalize: true }); - return Array.from(output.data as Float32Array); -} - -/** Batch-embed multiple texts (sequential to avoid OOM on large batches). */ -async function embedBatch(texts: string[]): Promise { - const model = await getEmbedder(); - const results: number[][] = []; - // Process in small batches to balance speed and memory - const batchSize = 16; - for (let i = 0; i < texts.length; i += batchSize) { - const batch = texts.slice(i, i + batchSize); - const outputs = await Promise.all( - batch.map(async (t) => { - const out = await model(t, { pooling: 'mean', normalize: true }); - return Array.from(out.data as Float32Array); - }) - ); - results.push(...outputs); - } - return results; -} - -// ────────────────────────────────────────────────────────────── -// Rule / pattern loading helpers -// ────────────────────────────────────────────────────────────── +// ── Rule / pattern types ───────────────────────────────────── interface AnalyzeRule { ruleId: string; @@ -82,6 +42,8 @@ interface ImeLogPattern { action?: string; } +// ── File helpers ───────────────────────────────────────────── + async function loadJsonFiles(dir: string): Promise { let entries: string[]; try { @@ -102,7 +64,7 @@ async function loadJsonFiles(dir: string): Promise { return results; } -/** Build the searchable text for each document type. */ +// ── Text builders ──────────────────────────────────────────── function analyzeRuleText(r: AnalyzeRule): string { const parts = [ @@ -136,22 +98,20 @@ function imePatternText(p: ImeLogPattern): string { ].filter(Boolean).join('\n'); } -// ────────────────────────────────────────────────────────────── -// Public API -// ────────────────────────────────────────────────────────────── - -export async function buildKnowledgeBase(rulesRoot: string): Promise { - const store = new VectorStore(); - - console.error('Loading knowledge base documents…'); +// ── Public API ─────────────────────────────────────────────── +/** + * Load all rule/pattern documents from the rules directory. + * Returns SearchDocument[] ready to be fed into any SearchProvider. + */ +export async function loadKnowledgeDocs(rulesRoot: string): Promise { const [analyzeRules, gatherRules, imePatterns] = await Promise.all([ loadJsonFiles(join(rulesRoot, 'analyze')), loadJsonFiles(join(rulesRoot, 'gather')), loadJsonFiles(join(rulesRoot, 'ime-log-patterns')), ]); - const docs: Array<{ id: string; text: string; metadata: Record }> = []; + const docs: SearchDocument[] = []; for (const r of analyzeRules) { docs.push({ @@ -193,18 +153,5 @@ export async function buildKnowledgeBase(rulesRoot: string): Promise d.text); - const embeddings = await embedBatch(texts); - - const vectorDocs: VectorDocument[] = docs.map((d, i) => ({ - ...d, - embedding: embeddings[i], - })); - - store.addMany(vectorDocs); - console.error(`Knowledge base ready — ${store.size} documents indexed.`); - - return store; + return docs; } diff --git a/src/McpServer/autopilot-monitor-mcp/src/search-factory.ts b/src/McpServer/autopilot-monitor-mcp/src/search-factory.ts new file mode 100644 index 00000000..9c794a72 --- /dev/null +++ b/src/McpServer/autopilot-monitor-mcp/src/search-factory.ts @@ -0,0 +1,46 @@ +/** + * Factory for creating SearchProvider instances. + * + * Selection logic: + * 1. Explicit env var SEARCH_BACKEND=vector|fuse → use that + * 2. Auto-detect: try to load @xenova/transformers + * - available → vector + * - missing → fuse (graceful fallback) + * + * This allows running the MCP server without the ~23 MB model download + * by setting SEARCH_BACKEND=fuse or simply not installing @xenova/transformers. + */ + +import type { SearchBackend, SearchProvider } from './search-provider.js'; + +async function isTransformersAvailable(): Promise { + try { + await import('@xenova/transformers'); + return true; + } catch { + return false; + } +} + +export async function resolveBackend(): Promise { + const explicit = process.env.SEARCH_BACKEND?.toLowerCase(); + if (explicit === 'vector' || explicit === 'fuse') return explicit; + + return (await isTransformersAvailable()) ? 'vector' : 'fuse'; +} + +/** + * Create a new SearchProvider instance. + * Pass a backend explicitly, or let it auto-detect. + */ +export async function createSearchProvider(backend?: SearchBackend): Promise { + const resolved = backend ?? await resolveBackend(); + + if (resolved === 'vector') { + const { VectorSearchProvider } = await import('./vector-search-provider.js'); + return new VectorSearchProvider(); + } + + const { FuseSearchProvider } = await import('./fuse-search-provider.js'); + return new FuseSearchProvider(); +} diff --git a/src/McpServer/autopilot-monitor-mcp/src/search-provider.ts b/src/McpServer/autopilot-monitor-mcp/src/search-provider.ts new file mode 100644 index 00000000..6300bc70 --- /dev/null +++ b/src/McpServer/autopilot-monitor-mcp/src/search-provider.ts @@ -0,0 +1,54 @@ +/** + * Abstract search provider interface. + * + * Implementations can use vector embeddings, fuzzy text matching, or any other + * ranking strategy. The MCP tools program against this interface so the backend + * is swappable without touching tool registration code. + */ + +// ── Shared types ───────────────────────────────────────────── + +export interface SearchDocument { + id: string; + text: string; + metadata: Record; +} + +export interface SearchResult { + id: string; + text: string; + metadata: Record; + /** Relevance score — always normalized to 0..1 regardless of backend. */ + score: number; +} + +export interface SearchOptions { + topK?: number; + minScore?: number; +} + +// ── Provider contract ──────────────────────────────────────── + +export interface SearchProvider { + /** Human-readable backend name (e.g. "vector/all-MiniLM-L6-v2", "fuse"). */ + readonly name: string; + + /** Number of currently indexed documents. */ + readonly size: number; + + /** + * Index a batch of documents. Can be called multiple times to add more docs. + * Implementation may precompute embeddings, build a Fuse index, etc. + */ + index(docs: SearchDocument[]): Promise; + + /** + * Search indexed documents by a natural-language query string. + * Returns results sorted by descending relevance. + */ + search(query: string, options?: SearchOptions): Promise; +} + +// ── Provider identifiers ───────────────────────────────────── + +export type SearchBackend = 'vector' | 'fuse'; diff --git a/src/McpServer/autopilot-monitor-mcp/src/tools.ts b/src/McpServer/autopilot-monitor-mcp/src/tools.ts index 1018e83b..09c1b9e4 100644 --- a/src/McpServer/autopilot-monitor-mcp/src/tools.ts +++ b/src/McpServer/autopilot-monitor-mcp/src/tools.ts @@ -1,10 +1,10 @@ import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { z } from 'zod'; import { apiFetch, buildQuery } from './client.js'; -import type { VectorStore } from './vector-store.js'; -import { embed } from './knowledge-base.js'; +import type { SearchProvider } from './search-provider.js'; +import { createSearchProvider } from './search-factory.js'; -export function registerTools(server: McpServer, knowledgeBase?: VectorStore): void { +export function registerTools(server: McpServer, knowledgeBase?: SearchProvider): void { // Tool 1: search_sessions server.tool( 'search_sessions', @@ -40,7 +40,6 @@ export function registerTools(server: McpServer, knowledgeBase?: VectorStore): v async (args) => { const { deviceProperties, ...rest } = args; const queryParams: Record = { ...rest }; - // Convert deviceProperties record into prop.* query parameters if (deviceProperties) { for (const [key, value] of Object.entries(deviceProperties)) { queryParams[`prop.${key}`] = value; @@ -107,7 +106,6 @@ export function registerTools(server: McpServer, knowledgeBase?: VectorStore): v events?: Array<{ eventType?: string; severity?: string; source?: string }>; count?: number; }; - // Client-side filter since backend doesn't support these filters on the events endpoint if (data?.events && (filters.eventType || filters.severity || filters.source)) { data.events = data.events.filter((e) => { if (filters.eventType && e.eventType !== filters.eventType) return false; @@ -170,10 +168,10 @@ export function registerTools(server: McpServer, knowledgeBase?: VectorStore): v } ); - // Tool 8: search_events_semantic (vector search over live events) + // Tool 8: search_events_semantic server.tool( 'search_events_semantic', - 'Semantic search over enrollment event messages within a session or across recent failed sessions. ' + + 'Semantic/fuzzy search over enrollment event messages within a session or across recent failed sessions. ' + 'Unlike exact filters, this finds events by MEANING — e.g. "network timeout" also matches "connection timed out", "request failed after waiting". ' + 'Use this when you need to find events that match a symptom description rather than an exact event type. ' + 'Provide a sessionId to search within one session, or omit it to search across recent failed sessions.', @@ -186,19 +184,16 @@ export function registerTools(server: McpServer, knowledgeBase?: VectorStore): v .describe('Minimum similarity score (0-1, default 0.35)'), }, async ({ query, sessionId, tenantId, topK, minScore }) => { - // Fetch events — either from a specific session or from recent failed sessions type EventEntry = { eventType?: string; severity?: string; source?: string; message?: string; timestamp?: string; phase?: string; data?: Record }; let events: EventEntry[] = []; let sessionIds: string[] = []; if (sessionId) { - // Single session mode const q = buildQuery({ tenantId } as Record); const data = await apiFetch(`/api/sessions/${sessionId}/events${q}`) as { events?: EventEntry[] }; events = data?.events ?? []; sessionIds = [sessionId]; } else { - // Cross-session mode: fetch recent failed sessions, then their events const searchQ = buildQuery({ status: 'Failed', tenantId, limit: 5 } as Record); const sessions = await apiFetch(`/api/sessions/search${searchQ}`) as { sessions?: Array<{ sessionId?: string }>; @@ -217,7 +212,6 @@ export function registerTools(server: McpServer, knowledgeBase?: VectorStore): v events = allEvents.flat(); } - // Filter to events that have a meaningful message const candidates = events.filter((e) => e.message && e.message.length > 5); if (candidates.length === 0) { @@ -229,53 +223,45 @@ export function registerTools(server: McpServer, knowledgeBase?: VectorStore): v }; } - // Build text representations and compute embeddings - const texts = candidates.map((e) => { + // Build searchable documents from events + const docs = candidates.map((e, i) => { const parts = [e.message]; if (e.eventType) parts.push(`Event: ${e.eventType}`); if (e.severity) parts.push(`Severity: ${e.severity}`); if (e.source) parts.push(`Source: ${e.source}`); - return parts.join(' | '); + return { + id: `event-${i}`, + text: parts.join(' | '), + metadata: { index: i } as Record, + }; }); - const queryEmbedding = await embed(query); + // Create a temporary search provider for this batch of events + const provider = await createSearchProvider(); + await provider.index(docs); + const searchResults = await provider.search(query, { topK, minScore }); - // Embed events in batches - const eventEmbeddings = await Promise.all(texts.map((t) => embed(t))); - - // Score and rank - const scored = candidates.map((e, i) => { - let dot = 0, normA = 0, normB = 0; - for (let j = 0; j < queryEmbedding.length; j++) { - dot += queryEmbedding[j] * eventEmbeddings[i][j]; - normA += queryEmbedding[j] * queryEmbedding[j]; - normB += eventEmbeddings[i][j] * eventEmbeddings[i][j]; - } - const denom = Math.sqrt(normA) * Math.sqrt(normB); - const score = denom === 0 ? 0 : dot / denom; - return { event: e, score, text: texts[i] }; - }); - - const results = scored - .filter((r) => r.score >= minScore) - .sort((a, b) => b.score - a.score) - .slice(0, topK) - .map((r) => ({ + const results = searchResults.map((r) => { + const idx = r.metadata.index as number; + const e = candidates[idx]; + return { score: Math.round(r.score * 1000) / 1000, - sessionId: (r.event as Record)._sessionId ?? sessionId, - eventType: r.event.eventType, - severity: r.event.severity, - source: r.event.source, - phase: r.event.phase, - timestamp: r.event.timestamp, - message: r.event.message, - })); + sessionId: (e as Record)._sessionId ?? sessionId, + eventType: e.eventType, + severity: e.severity, + source: e.source, + phase: e.phase, + timestamp: e.timestamp, + message: e.message, + }; + }); return { content: [{ type: 'text' as const, text: JSON.stringify({ query, + searchBackend: provider.name, sessionsSearched: sessionIds, eventsAnalyzed: candidates.length, resultCount: results.length, @@ -286,12 +272,12 @@ export function registerTools(server: McpServer, knowledgeBase?: VectorStore): v } ); - // Tool 9: search_knowledge (vector search over rules) + // Tool 9: search_knowledge server.tool( 'search_knowledge', - 'Semantic search over the Autopilot Monitor knowledge base: analysis rules, gather rules, and IME log patterns. ' + + 'Semantic/fuzzy search over the Autopilot Monitor knowledge base: analysis rules, gather rules, and IME log patterns. ' + 'Use natural language queries like "app install timeout", "BitLocker issues", "detection script failure". ' + - 'Returns the most relevant rules and patterns ranked by semantic similarity. ' + + 'Returns the most relevant rules and patterns ranked by similarity. ' + 'Great for finding remediation steps, understanding error patterns, or discovering relevant diagnostic rules.', { query: z.string().describe('Natural language search query (e.g. "app download timeout", "TPM not ready", "ESP stuck")'), @@ -306,15 +292,15 @@ export function registerTools(server: McpServer, knowledgeBase?: VectorStore): v return { content: [{ type: 'text' as const, - text: JSON.stringify({ error: 'Knowledge base not initialized. The server may still be loading embeddings.' }), + text: JSON.stringify({ error: 'Knowledge base not initialized. The server may still be loading.' }), }], }; } - const queryEmbedding = await embed(query); - let results = knowledgeBase.search(queryEmbedding, type === 'all' ? topK * 3 : topK, minScore); + // Over-fetch when filtering by type, then trim + const fetchK = type === 'all' ? topK : topK * 3; + let results = await knowledgeBase.search(query, { topK: fetchK, minScore }); - // Filter by type if specified if (type !== 'all') { results = results.filter((r) => r.metadata.type === type); } @@ -333,7 +319,12 @@ export function registerTools(server: McpServer, knowledgeBase?: VectorStore): v return { content: [{ type: 'text' as const, - text: JSON.stringify({ query, resultCount: formatted.length, results: formatted }, null, 2), + text: JSON.stringify({ + query, + searchBackend: knowledgeBase.name, + resultCount: formatted.length, + results: formatted, + }, null, 2), }], }; } diff --git a/src/McpServer/autopilot-monitor-mcp/src/vector-search-provider.ts b/src/McpServer/autopilot-monitor-mcp/src/vector-search-provider.ts new file mode 100644 index 00000000..1a1feb79 --- /dev/null +++ b/src/McpServer/autopilot-monitor-mcp/src/vector-search-provider.ts @@ -0,0 +1,96 @@ +/** + * SearchProvider backed by sentence-transformer embeddings + cosine similarity. + * Uses @xenova/transformers with the all-MiniLM-L6-v2 model (quantized, ~23 MB). + * + * Pros: True semantic understanding — "timeout" matches "waiting exceeded". + * Cons: First-run downloads the model; higher memory & CPU at index time. + */ + +import { pipeline, type FeatureExtractionPipeline } from '@xenova/transformers'; +import type { SearchDocument, SearchOptions, SearchProvider, SearchResult } from './search-provider.js'; + +const MODEL_NAME = 'Xenova/all-MiniLM-L6-v2'; + +// ── Shared singleton embedder ──────────────────────────────── + +let embedder: FeatureExtractionPipeline | null = null; + +async function getEmbedder(): Promise { + if (!embedder) { + embedder = await pipeline('feature-extraction', MODEL_NAME, { + quantized: true, + }) as FeatureExtractionPipeline; + } + return embedder; +} + +/** Compute a normalized embedding for a single text. */ +export async function embed(text: string): Promise { + const model = await getEmbedder(); + const output = await model(text, { pooling: 'mean', normalize: true }); + return Array.from(output.data as Float32Array); +} + +// ── Helpers ────────────────────────────────────────────────── + +function cosineSimilarity(a: number[], b: number[]): number { + let dot = 0; + let normA = 0; + let normB = 0; + for (let i = 0; i < a.length; i++) { + dot += a[i] * b[i]; + normA += a[i] * a[i]; + normB += b[i] * b[i]; + } + const denom = Math.sqrt(normA) * Math.sqrt(normB); + return denom === 0 ? 0 : dot / denom; +} + +// ── Provider implementation ────────────────────────────────── + +interface StoredDocument extends SearchDocument { + embedding: number[]; +} + +export class VectorSearchProvider implements SearchProvider { + readonly name = `vector/${MODEL_NAME}`; + private documents: StoredDocument[] = []; + + get size(): number { + return this.documents.length; + } + + async index(docs: SearchDocument[]): Promise { + const model = await getEmbedder(); + const batchSize = 16; + for (let i = 0; i < docs.length; i += batchSize) { + const batch = docs.slice(i, i + batchSize); + const embeddings = await Promise.all( + batch.map(async (d) => { + const out = await model(d.text, { pooling: 'mean', normalize: true }); + return Array.from(out.data as Float32Array); + }) + ); + for (let j = 0; j < batch.length; j++) { + this.documents.push({ ...batch[j], embedding: embeddings[j] }); + } + } + } + + async search(query: string, options: SearchOptions = {}): Promise { + const { topK = 5, minScore = 0.3 } = options; + const queryEmbedding = await embed(query); + + const scored = this.documents.map((doc) => ({ + id: doc.id, + text: doc.text, + metadata: doc.metadata, + score: cosineSimilarity(queryEmbedding, doc.embedding), + })); + + return scored + .filter((r) => r.score >= minScore) + .sort((a, b) => b.score - a.score) + .slice(0, topK); + } +} diff --git a/src/McpServer/autopilot-monitor-mcp/src/vector-store.ts b/src/McpServer/autopilot-monitor-mcp/src/vector-store.ts deleted file mode 100644 index c3828562..00000000 --- a/src/McpServer/autopilot-monitor-mcp/src/vector-store.ts +++ /dev/null @@ -1,61 +0,0 @@ -/** - * Lightweight in-memory vector store using cosine similarity. - * Each document is stored alongside its embedding vector. - */ - -export interface VectorDocument { - id: string; - text: string; - metadata: Record; - embedding: number[]; -} - -export interface SearchResult { - id: string; - text: string; - metadata: Record; - score: number; -} - -function cosineSimilarity(a: number[], b: number[]): number { - let dot = 0; - let normA = 0; - let normB = 0; - for (let i = 0; i < a.length; i++) { - dot += a[i] * b[i]; - normA += a[i] * a[i]; - normB += b[i] * b[i]; - } - const denom = Math.sqrt(normA) * Math.sqrt(normB); - return denom === 0 ? 0 : dot / denom; -} - -export class VectorStore { - private documents: VectorDocument[] = []; - - add(doc: VectorDocument): void { - this.documents.push(doc); - } - - addMany(docs: VectorDocument[]): void { - this.documents.push(...docs); - } - - search(queryEmbedding: number[], topK = 5, minScore = 0.3): SearchResult[] { - const scored = this.documents.map((doc) => ({ - id: doc.id, - text: doc.text, - metadata: doc.metadata, - score: cosineSimilarity(queryEmbedding, doc.embedding), - })); - - return scored - .filter((r) => r.score >= minScore) - .sort((a, b) => b.score - a.score) - .slice(0, topK); - } - - get size(): number { - return this.documents.length; - } -}