mirror of
https://github.com/okieselbach/AutopilotMonitor.git
synced 2026-10-24 20:53:06 +00:00
feat(mcp): add vector search with semantic knowledge base and event search
Add two new MCP tools powered by local embeddings (all-MiniLM-L6-v2 via @xenova/transformers): - search_knowledge: Semantic search over 103 analysis rules, gather rules, and IME log patterns. Finds relevant rules by meaning, not exact match. - search_events_semantic: Semantic search over enrollment event messages within a session or across recent failed sessions. Matches symptoms by meaning (e.g. "timeout" finds "connection timed out", "waiting exceeded"). Implementation: - vector-store.ts: In-memory vector store with cosine similarity - knowledge-base.ts: Loads rules from disk, computes embeddings at startup - Embeddings computed locally — no external API key required https://claude.ai/code/session_01TBPsvvMBue9Ft7j3CfD4S2
This commit is contained in:
Claude
+852
-4
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "autopilot-monitor-mcp",
|
||||
"version": "1.0.0",
|
||||
"version": "1.1.0",
|
||||
"description": "MCP server for Autopilot-Monitor — enables AI agents to query enrollment sessions, device properties, events and vulnerabilities",
|
||||
"type": "module",
|
||||
"bin": {
|
||||
@@ -14,6 +14,7 @@
|
||||
"dependencies": {
|
||||
"@anthropic-ai/sdk": "^0.80.0",
|
||||
"@modelcontextprotocol/sdk": "^1.0.0",
|
||||
"@xenova/transformers": "^2.17.0",
|
||||
"zod": "^3.22.0"
|
||||
},
|
||||
"devDependencies": {
|
||||
|
||||
@@ -1,14 +1,27 @@
|
||||
import { resolve, dirname } from 'node:path';
|
||||
import { fileURLToPath } from 'node:url';
|
||||
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';
|
||||
|
||||
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.0.0',
|
||||
version: '1.1.0',
|
||||
});
|
||||
|
||||
registerTools(server);
|
||||
// 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);
|
||||
|
||||
registerTools(server, knowledgeBase);
|
||||
registerResources(server);
|
||||
|
||||
const transport = new StdioServerTransport();
|
||||
@@ -16,3 +29,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}`);
|
||||
|
||||
@@ -0,0 +1,210 @@
|
||||
/**
|
||||
* 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.
|
||||
*/
|
||||
|
||||
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';
|
||||
|
||||
const MODEL_NAME = 'Xenova/all-MiniLM-L6-v2';
|
||||
|
||||
let embedder: FeatureExtractionPipeline | null = null;
|
||||
|
||||
async function getEmbedder(): Promise<FeatureExtractionPipeline> {
|
||||
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<number[]> {
|
||||
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<number[][]> {
|
||||
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
|
||||
// ──────────────────────────────────────────────────────────────
|
||||
|
||||
interface AnalyzeRule {
|
||||
ruleId: string;
|
||||
title: string;
|
||||
description: string;
|
||||
severity: string;
|
||||
category: string;
|
||||
explanation?: string;
|
||||
remediation?: Array<{ title: string; steps: string[] }>;
|
||||
tags?: string[];
|
||||
}
|
||||
|
||||
interface GatherRule {
|
||||
ruleId: string;
|
||||
title: string;
|
||||
description: string;
|
||||
category: string;
|
||||
collectorType?: string;
|
||||
target?: string;
|
||||
trigger?: string;
|
||||
tags?: string[];
|
||||
}
|
||||
|
||||
interface ImeLogPattern {
|
||||
patternId: string;
|
||||
category: string;
|
||||
description: string;
|
||||
pattern: string;
|
||||
action?: string;
|
||||
}
|
||||
|
||||
async function loadJsonFiles<T>(dir: string): Promise<T[]> {
|
||||
let entries: string[];
|
||||
try {
|
||||
entries = await readdir(dir);
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
const results: T[] = [];
|
||||
for (const f of entries) {
|
||||
if (!f.endsWith('.json')) continue;
|
||||
try {
|
||||
const raw = await readFile(join(dir, f), 'utf-8');
|
||||
results.push(JSON.parse(raw) as T);
|
||||
} catch {
|
||||
// skip malformed files
|
||||
}
|
||||
}
|
||||
return results;
|
||||
}
|
||||
|
||||
/** Build the searchable text for each document type. */
|
||||
|
||||
function analyzeRuleText(r: AnalyzeRule): string {
|
||||
const parts = [
|
||||
`[${r.ruleId}] ${r.title}`,
|
||||
r.description,
|
||||
r.explanation ?? '',
|
||||
];
|
||||
if (r.remediation) {
|
||||
for (const rem of r.remediation) {
|
||||
parts.push(`Remediation: ${rem.title} — ${rem.steps.join('; ')}`);
|
||||
}
|
||||
}
|
||||
if (r.tags?.length) parts.push(`Tags: ${r.tags.join(', ')}`);
|
||||
return parts.filter(Boolean).join('\n');
|
||||
}
|
||||
|
||||
function gatherRuleText(r: GatherRule): string {
|
||||
return [
|
||||
`[${r.ruleId}] ${r.title}`,
|
||||
r.description,
|
||||
r.target ? `Target: ${r.target}` : '',
|
||||
r.tags?.length ? `Tags: ${r.tags.join(', ')}` : '',
|
||||
].filter(Boolean).join('\n');
|
||||
}
|
||||
|
||||
function imePatternText(p: ImeLogPattern): string {
|
||||
return [
|
||||
`[${p.patternId}] ${p.description}`,
|
||||
`Category: ${p.category}`,
|
||||
p.action ? `Action: ${p.action}` : '',
|
||||
].filter(Boolean).join('\n');
|
||||
}
|
||||
|
||||
// ──────────────────────────────────────────────────────────────
|
||||
// Public API
|
||||
// ──────────────────────────────────────────────────────────────
|
||||
|
||||
export async function buildKnowledgeBase(rulesRoot: string): Promise<VectorStore> {
|
||||
const store = new VectorStore();
|
||||
|
||||
console.error('Loading knowledge base documents…');
|
||||
|
||||
const [analyzeRules, gatherRules, imePatterns] = await Promise.all([
|
||||
loadJsonFiles<AnalyzeRule>(join(rulesRoot, 'analyze')),
|
||||
loadJsonFiles<GatherRule>(join(rulesRoot, 'gather')),
|
||||
loadJsonFiles<ImeLogPattern>(join(rulesRoot, 'ime-log-patterns')),
|
||||
]);
|
||||
|
||||
const docs: Array<{ id: string; text: string; metadata: Record<string, unknown> }> = [];
|
||||
|
||||
for (const r of analyzeRules) {
|
||||
docs.push({
|
||||
id: r.ruleId,
|
||||
text: analyzeRuleText(r),
|
||||
metadata: {
|
||||
type: 'analyze-rule',
|
||||
severity: r.severity,
|
||||
category: r.category,
|
||||
title: r.title,
|
||||
tags: r.tags ?? [],
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
for (const r of gatherRules) {
|
||||
docs.push({
|
||||
id: r.ruleId,
|
||||
text: gatherRuleText(r),
|
||||
metadata: {
|
||||
type: 'gather-rule',
|
||||
category: r.category,
|
||||
title: r.title,
|
||||
tags: r.tags ?? [],
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
for (const p of imePatterns) {
|
||||
docs.push({
|
||||
id: p.patternId,
|
||||
text: imePatternText(p),
|
||||
metadata: {
|
||||
type: 'ime-log-pattern',
|
||||
category: p.category,
|
||||
description: p.description,
|
||||
action: p.action ?? null,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
console.error(`Embedding ${docs.length} documents (model: ${MODEL_NAME})…`);
|
||||
|
||||
const texts = docs.map((d) => 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;
|
||||
}
|
||||
@@ -1,8 +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';
|
||||
|
||||
export function registerTools(server: McpServer): void {
|
||||
export function registerTools(server: McpServer, knowledgeBase?: VectorStore): void {
|
||||
// Tool 1: search_sessions
|
||||
server.tool(
|
||||
'search_sessions',
|
||||
@@ -167,4 +169,173 @@ export function registerTools(server: McpServer): void {
|
||||
return { content: [{ type: 'text' as const, text: JSON.stringify(data, null, 2) }] };
|
||||
}
|
||||
);
|
||||
|
||||
// Tool 8: search_events_semantic (vector search over live events)
|
||||
server.tool(
|
||||
'search_events_semantic',
|
||||
'Semantic 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.',
|
||||
{
|
||||
query: z.string().describe('Natural language description of what to find (e.g. "app download stuck", "certificate error", "disk space low")'),
|
||||
sessionId: z.string().optional().describe('Search within a specific session. If omitted, searches across recent failed sessions.'),
|
||||
tenantId: z.string().optional().describe('Tenant ID (global-scoped key only)'),
|
||||
topK: z.number().min(1).max(30).optional().default(10).describe('Number of matching events to return (1-30, default 10)'),
|
||||
minScore: z.number().min(0).max(1).optional().default(0.35)
|
||||
.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<string, unknown> };
|
||||
let events: EventEntry[] = [];
|
||||
let sessionIds: string[] = [];
|
||||
|
||||
if (sessionId) {
|
||||
// Single session mode
|
||||
const q = buildQuery({ tenantId } as Record<string, string | undefined>);
|
||||
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<string, string | number | undefined>);
|
||||
const sessions = await apiFetch(`/api/sessions/search${searchQ}`) as {
|
||||
sessions?: Array<{ sessionId?: string }>;
|
||||
};
|
||||
const ids = (sessions?.sessions ?? []).map((s) => s.sessionId).filter(Boolean) as string[];
|
||||
sessionIds = ids;
|
||||
const q = buildQuery({ tenantId } as Record<string, string | undefined>);
|
||||
const allEvents = await Promise.all(
|
||||
ids.map(async (sid) => {
|
||||
try {
|
||||
const d = await apiFetch(`/api/sessions/${sid}/events${q}`) as { events?: EventEntry[] };
|
||||
return (d?.events ?? []).map((e) => ({ ...e, _sessionId: sid }));
|
||||
} catch { return []; }
|
||||
})
|
||||
);
|
||||
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) {
|
||||
return {
|
||||
content: [{
|
||||
type: 'text' as const,
|
||||
text: JSON.stringify({ query, resultCount: 0, results: [], note: 'No events with messages found.' }),
|
||||
}],
|
||||
};
|
||||
}
|
||||
|
||||
// Build text representations and compute embeddings
|
||||
const texts = candidates.map((e) => {
|
||||
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(' | ');
|
||||
});
|
||||
|
||||
const queryEmbedding = await embed(query);
|
||||
|
||||
// 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) => ({
|
||||
score: Math.round(r.score * 1000) / 1000,
|
||||
sessionId: (r.event as Record<string, unknown>)._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,
|
||||
}));
|
||||
|
||||
return {
|
||||
content: [{
|
||||
type: 'text' as const,
|
||||
text: JSON.stringify({
|
||||
query,
|
||||
sessionsSearched: sessionIds,
|
||||
eventsAnalyzed: candidates.length,
|
||||
resultCount: results.length,
|
||||
results,
|
||||
}, null, 2),
|
||||
}],
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
// Tool 9: search_knowledge (vector search over rules)
|
||||
server.tool(
|
||||
'search_knowledge',
|
||||
'Semantic 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. ' +
|
||||
'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")'),
|
||||
topK: z.number().min(1).max(20).optional().default(5).describe('Number of results to return (1-20, default 5)'),
|
||||
type: z.enum(['all', 'analyze-rule', 'gather-rule', 'ime-log-pattern']).optional().default('all')
|
||||
.describe('Filter by document type. Default: search all types.'),
|
||||
minScore: z.number().min(0).max(1).optional().default(0.3)
|
||||
.describe('Minimum similarity score threshold (0-1, default 0.3). Lower = more results, higher = stricter matching.'),
|
||||
},
|
||||
async ({ query, topK, type, minScore }) => {
|
||||
if (!knowledgeBase || knowledgeBase.size === 0) {
|
||||
return {
|
||||
content: [{
|
||||
type: 'text' as const,
|
||||
text: JSON.stringify({ error: 'Knowledge base not initialized. The server may still be loading embeddings.' }),
|
||||
}],
|
||||
};
|
||||
}
|
||||
|
||||
const queryEmbedding = await embed(query);
|
||||
let results = knowledgeBase.search(queryEmbedding, type === 'all' ? topK * 3 : topK, minScore);
|
||||
|
||||
// Filter by type if specified
|
||||
if (type !== 'all') {
|
||||
results = results.filter((r) => r.metadata.type === type);
|
||||
}
|
||||
|
||||
results = results.slice(0, topK);
|
||||
|
||||
const formatted = results.map((r) => ({
|
||||
id: r.id,
|
||||
score: Math.round(r.score * 1000) / 1000,
|
||||
type: r.metadata.type,
|
||||
title: r.metadata.title ?? r.metadata.description ?? r.id,
|
||||
content: r.text,
|
||||
metadata: r.metadata,
|
||||
}));
|
||||
|
||||
return {
|
||||
content: [{
|
||||
type: 'text' as const,
|
||||
text: JSON.stringify({ query, resultCount: formatted.length, results: formatted }, null, 2),
|
||||
}],
|
||||
};
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
/**
|
||||
* 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<string, unknown>;
|
||||
embedding: number[];
|
||||
}
|
||||
|
||||
export interface SearchResult {
|
||||
id: string;
|
||||
text: string;
|
||||
metadata: Record<string, unknown>;
|
||||
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;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user