Refactor: extract store, LLM, and formatter modules with comprehensive tests
- Extract store.ts: database operations, search, document retrieval - createStore() factory pattern for clean DB lifecycle management - Unified DocumentResult type with optional body loading - Snippet extraction with diff-style headers (@@ -line,count @@) - Extract llm.ts: LLM abstraction layer with Ollama implementation - Clean interface for embed, generate, rerank operations - High-level rerankerLogprobsCheck with logprob-based scoring - Query expansion support - Extract formatter.ts: output formatting utilities - Support for CLI, JSON, CSV, MD, XML formats - MCP-specific CSV formatting - Extract mcp.ts: MCP server using createStore() pattern - Single DB connection for server lifetime (fixes closed DB errors) - URL-decode resource paths for proper space/special char handling - Add comprehensive test suites (215 tests total) - store.test.ts: 96 tests covering all store operations - llm.test.ts: 60 tests for LLM abstraction - mcp.test.ts: 59 tests for MCP endpoints and resources - All tests use mocked Ollama (errors on unmocked calls) - Add bun run inspector script for MCP debugging 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Tobi Lutke
parent
542509a098
commit
bab46dacb2
20
CLAUDE.md
20
CLAUDE.md
@ -12,6 +12,26 @@ qmd embed # Generate vector embeddings (requires Ollama)
|
||||
qmd search <query> # BM25 full-text search
|
||||
qmd vsearch <query> # Vector similarity search
|
||||
qmd query <query> # Hybrid search with reranking (best quality)
|
||||
qmd get <file> # Get document content (fuzzy matches if not found)
|
||||
qmd multi-get <pattern> # Get multiple docs by glob or comma-separated list
|
||||
```
|
||||
|
||||
## Options
|
||||
|
||||
```sh
|
||||
# Search & retrieval
|
||||
-c, --collection <name> # Restrict search to a collection (matches pwd suffix)
|
||||
-n <num> # Number of results
|
||||
--all # Return all matches
|
||||
--min-score <num> # Minimum score threshold
|
||||
--full # Show full document content
|
||||
|
||||
# Multi-get specific
|
||||
-l <num> # Maximum lines per file
|
||||
--max-bytes <num> # Skip files larger than this (default 10KB)
|
||||
|
||||
# Output formats (search and multi-get)
|
||||
--json, --csv, --md, --xml, --files
|
||||
```
|
||||
|
||||
## Development
|
||||
|
||||
45
README.md
45
README.md
@ -31,6 +31,12 @@ qmd query "quarterly planning process" # Hybrid + reranking (best quality)
|
||||
# Get a specific document
|
||||
qmd get "meetings/2024-01-15.md"
|
||||
|
||||
# Get multiple documents by glob pattern
|
||||
qmd multi-get "journals/2025-05*.md"
|
||||
|
||||
# Search within a specific collection
|
||||
qmd search "API" -c notes
|
||||
|
||||
# Export all matches for an agent
|
||||
qmd search "API" --all --files --min-score 0.3
|
||||
```
|
||||
@ -55,10 +61,11 @@ qmd get "docs/api-reference.md" --full
|
||||
Although the tool works perfectly fine when you just tell your agent to use it on the command line, it also exposes an MCP (Model Context Protocol) server for tighter integration.
|
||||
|
||||
**Tools exposed:**
|
||||
- `qmd_search` - Fast BM25 keyword search
|
||||
- `qmd_vsearch` - Semantic vector search
|
||||
- `qmd_query` - Hybrid search with reranking (best quality)
|
||||
- `qmd_get` - Retrieve document content
|
||||
- `qmd_search` - Fast BM25 keyword search (supports collection filter)
|
||||
- `qmd_vsearch` - Semantic vector search (supports collection filter)
|
||||
- `qmd_query` - Hybrid search with reranking (supports collection filter)
|
||||
- `qmd_get` - Retrieve document content (with fuzzy matching suggestions)
|
||||
- `qmd_multi_get` - Retrieve multiple documents by glob pattern or list
|
||||
- `qmd_status` - Index health and collection info
|
||||
|
||||
**Claude Desktop configuration** (`~/Library/Application Support/Claude/claude_desktop_config.json`):
|
||||
@ -278,16 +285,24 @@ qmd query "user authentication"
|
||||
### Options
|
||||
|
||||
```sh
|
||||
# Search options
|
||||
-n <num> # Number of results (default: 5, or 20 for --files/--json)
|
||||
-c, --collection # Restrict search to a specific collection
|
||||
--all # Return all matches (use with --min-score to filter)
|
||||
--min-score <num> # Minimum score threshold (default: 0)
|
||||
--full # Show full document content
|
||||
--files # Output: score,filepath,context
|
||||
--json # JSON output with snippets
|
||||
--csv # CSV output with snippets
|
||||
--index <name> # Use named index
|
||||
|
||||
# Output formats (for search and multi-get)
|
||||
--files # Output: score,filepath,context (search) or filepath,context (multi-get)
|
||||
--json # JSON output
|
||||
--csv # CSV output
|
||||
--md # Markdown output
|
||||
--xml # XML output
|
||||
--index <name> # Use named index
|
||||
|
||||
# Multi-get options
|
||||
-l <num> # Maximum lines per file
|
||||
--max-bytes <num> # Skip files larger than N bytes (default: 10KB)
|
||||
```
|
||||
|
||||
### Output Format
|
||||
@ -345,9 +360,21 @@ qmd status
|
||||
# Re-index all collections
|
||||
qmd update-all
|
||||
|
||||
# Get document body by filepath
|
||||
# Get document body by filepath (with fuzzy matching)
|
||||
qmd get ~/notes/meeting.md
|
||||
|
||||
# Get multiple documents by glob pattern
|
||||
qmd multi-get "journals/2025-05*.md"
|
||||
|
||||
# Get multiple documents by comma-separated list
|
||||
qmd multi-get "doc1.md, doc2.md, doc3.md"
|
||||
|
||||
# Limit multi-get to files under 20KB
|
||||
qmd multi-get "docs/*.md" --max-bytes 20480
|
||||
|
||||
# Output multi-get as JSON for agent processing
|
||||
qmd multi-get "docs/*.md" --json
|
||||
|
||||
# Clean up cache and orphaned data
|
||||
qmd cleanup
|
||||
```
|
||||
|
||||
359
formatter.ts
Normal file
359
formatter.ts
Normal file
@ -0,0 +1,359 @@
|
||||
/**
|
||||
* formatter.ts - Output formatting utilities for QMD
|
||||
*
|
||||
* Provides methods to format search results and documents into various output formats:
|
||||
* JSON, CSV, XML, Markdown, files list, and CLI (colored terminal output).
|
||||
*/
|
||||
|
||||
import { extractSnippet } from "./store.js";
|
||||
import type { SearchResult, MultiGetFile, MultiGetResult, DocumentResult } from "./store.js";
|
||||
|
||||
// =============================================================================
|
||||
// Types
|
||||
// =============================================================================
|
||||
|
||||
// Re-export store types for convenience
|
||||
export type { SearchResult, MultiGetFile, MultiGetResult, DocumentResult };
|
||||
|
||||
export type OutputFormat = "cli" | "csv" | "md" | "xml" | "files" | "json";
|
||||
|
||||
export type FormatOptions = {
|
||||
full?: boolean; // Show full document content instead of snippet
|
||||
query?: string; // Query for snippet extraction and highlighting
|
||||
useColor?: boolean; // Enable terminal colors (default: false for non-CLI)
|
||||
};
|
||||
|
||||
// =============================================================================
|
||||
// Escape Helpers
|
||||
// =============================================================================
|
||||
|
||||
export function escapeCSV(value: string | null | number): string {
|
||||
if (value === null || value === undefined) return "";
|
||||
const str = String(value);
|
||||
if (str.includes(",") || str.includes('"') || str.includes("\n")) {
|
||||
return `"${str.replace(/"/g, '""')}"`;
|
||||
}
|
||||
return str;
|
||||
}
|
||||
|
||||
export function escapeXml(str: string): string {
|
||||
return str
|
||||
.replace(/&/g, "&")
|
||||
.replace(/</g, "<")
|
||||
.replace(/>/g, ">")
|
||||
.replace(/"/g, """)
|
||||
.replace(/'/g, "'");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Search Results Formatters
|
||||
// =============================================================================
|
||||
|
||||
/**
|
||||
* Format search results as JSON
|
||||
*/
|
||||
export function searchResultsToJson(
|
||||
results: SearchResult[],
|
||||
opts: FormatOptions = {}
|
||||
): string {
|
||||
const output = results.map(row => ({
|
||||
score: Math.round(row.score * 100) / 100,
|
||||
file: row.displayPath,
|
||||
title: row.title,
|
||||
...(row.context && { context: row.context }),
|
||||
...(opts.full && { body: row.body }),
|
||||
...(!opts.full && opts.query && { snippet: extractSnippet(row.body, opts.query, 300, row.chunkPos).snippet }),
|
||||
}));
|
||||
return JSON.stringify(output, null, 2);
|
||||
}
|
||||
|
||||
/**
|
||||
* Format search results as CSV
|
||||
*/
|
||||
export function searchResultsToCsv(
|
||||
results: SearchResult[],
|
||||
opts: FormatOptions = {}
|
||||
): string {
|
||||
const query = opts.query || "";
|
||||
const header = "score,file,title,context,line,snippet";
|
||||
const rows = results.map(row => {
|
||||
const { line, snippet } = extractSnippet(row.body, query, 500, row.chunkPos);
|
||||
const content = opts.full ? row.body : snippet;
|
||||
return [
|
||||
row.score.toFixed(4),
|
||||
escapeCSV(row.displayPath),
|
||||
escapeCSV(row.title),
|
||||
escapeCSV(row.context || ""),
|
||||
line,
|
||||
escapeCSV(content),
|
||||
].join(",");
|
||||
});
|
||||
return [header, ...rows].join("\n");
|
||||
}
|
||||
|
||||
/**
|
||||
* Format search results as simple files list (score,filepath,context)
|
||||
*/
|
||||
export function searchResultsToFiles(results: SearchResult[]): string {
|
||||
return results.map(row => {
|
||||
const ctx = row.context ? `,"${row.context.replace(/"/g, '""')}"` : "";
|
||||
return `${row.score.toFixed(2)},${row.displayPath}${ctx}`;
|
||||
}).join("\n");
|
||||
}
|
||||
|
||||
/**
|
||||
* Format search results as Markdown
|
||||
*/
|
||||
export function searchResultsToMarkdown(
|
||||
results: SearchResult[],
|
||||
opts: FormatOptions = {}
|
||||
): string {
|
||||
const query = opts.query || "";
|
||||
return results.map(row => {
|
||||
const heading = row.title || row.displayPath;
|
||||
if (opts.full) {
|
||||
return `---\n# ${heading}\n\n${row.body}\n`;
|
||||
} else {
|
||||
const { snippet } = extractSnippet(row.body, query, 500, row.chunkPos);
|
||||
return `---\n# ${heading}\n\n${snippet}\n`;
|
||||
}
|
||||
}).join("\n");
|
||||
}
|
||||
|
||||
/**
|
||||
* Format search results as XML
|
||||
*/
|
||||
export function searchResultsToXml(
|
||||
results: SearchResult[],
|
||||
opts: FormatOptions = {}
|
||||
): string {
|
||||
const query = opts.query || "";
|
||||
const items = results.map(row => {
|
||||
const titleAttr = row.title ? ` title="${escapeXml(row.title)}"` : "";
|
||||
const content = opts.full ? row.body : extractSnippet(row.body, query, 500, row.chunkPos).snippet;
|
||||
return `<file name="${escapeXml(row.displayPath)}"${titleAttr}>\n${escapeXml(content)}\n</file>`;
|
||||
});
|
||||
return items.join("\n\n");
|
||||
}
|
||||
|
||||
/**
|
||||
* Format search results for MCP (simpler CSV format with pre-extracted snippets)
|
||||
*/
|
||||
export function searchResultsToMcpCsv(
|
||||
results: { file: string; title: string; score: number; context: string | null; snippet: string }[]
|
||||
): string {
|
||||
const header = "file,title,score,context,snippet";
|
||||
const rows = results.map(r =>
|
||||
[r.file, r.title, r.score, r.context || "", r.snippet].map(escapeCSV).join(",")
|
||||
);
|
||||
return [header, ...rows].join("\n");
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Document Formatters (for multi-get using MultiGetFile from store)
|
||||
// =============================================================================
|
||||
|
||||
/**
|
||||
* Format documents as JSON
|
||||
*/
|
||||
export function documentsToJson(results: MultiGetFile[]): string {
|
||||
const output = results.map(r => ({
|
||||
file: r.displayPath,
|
||||
title: r.title,
|
||||
...(r.context && { context: r.context }),
|
||||
...(r.skipped ? { skipped: true, reason: r.skipReason } : { body: r.body }),
|
||||
}));
|
||||
return JSON.stringify(output, null, 2);
|
||||
}
|
||||
|
||||
/**
|
||||
* Format documents as CSV
|
||||
*/
|
||||
export function documentsToCsv(results: MultiGetFile[]): string {
|
||||
const header = "file,title,context,skipped,body";
|
||||
const rows = results.map(r =>
|
||||
[
|
||||
r.displayPath,
|
||||
r.title,
|
||||
r.context || "",
|
||||
r.skipped ? "true" : "false",
|
||||
r.skipped ? (r.skipReason || "") : r.body
|
||||
].map(escapeCSV).join(",")
|
||||
);
|
||||
return [header, ...rows].join("\n");
|
||||
}
|
||||
|
||||
/**
|
||||
* Format documents as files list
|
||||
*/
|
||||
export function documentsToFiles(results: MultiGetFile[]): string {
|
||||
return results.map(r => {
|
||||
const ctx = r.context ? `,"${r.context.replace(/"/g, '""')}"` : "";
|
||||
const status = r.skipped ? ",[SKIPPED]" : "";
|
||||
return `${r.displayPath}${ctx}${status}`;
|
||||
}).join("\n");
|
||||
}
|
||||
|
||||
/**
|
||||
* Format documents as Markdown
|
||||
*/
|
||||
export function documentsToMarkdown(results: MultiGetFile[]): string {
|
||||
return results.map(r => {
|
||||
let md = `## ${r.displayPath}\n\n`;
|
||||
if (r.title && r.title !== r.displayPath) md += `**Title:** ${r.title}\n\n`;
|
||||
if (r.context) md += `**Context:** ${r.context}\n\n`;
|
||||
if (r.skipped) {
|
||||
md += `> ${r.skipReason}\n`;
|
||||
} else {
|
||||
md += "```\n" + r.body + "\n```\n";
|
||||
}
|
||||
return md;
|
||||
}).join("\n");
|
||||
}
|
||||
|
||||
/**
|
||||
* Format documents as XML
|
||||
*/
|
||||
export function documentsToXml(results: MultiGetFile[]): string {
|
||||
const items = results.map(r => {
|
||||
let xml = " <document>\n";
|
||||
xml += ` <file>${escapeXml(r.displayPath)}</file>\n`;
|
||||
xml += ` <title>${escapeXml(r.title)}</title>\n`;
|
||||
if (r.context) xml += ` <context>${escapeXml(r.context)}</context>\n`;
|
||||
if (r.skipped) {
|
||||
xml += ` <skipped>true</skipped>\n`;
|
||||
xml += ` <reason>${escapeXml(r.skipReason || "")}</reason>\n`;
|
||||
} else {
|
||||
xml += ` <body>${escapeXml(r.body)}</body>\n`;
|
||||
}
|
||||
xml += " </document>";
|
||||
return xml;
|
||||
});
|
||||
return `<?xml version="1.0" encoding="UTF-8"?>\n<documents>\n${items.join("\n")}\n</documents>`;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Single Document Formatters
|
||||
// =============================================================================
|
||||
|
||||
/**
|
||||
* Format a single DocumentResult as JSON
|
||||
*/
|
||||
export function documentToJson(doc: DocumentResult): string {
|
||||
return JSON.stringify({
|
||||
file: doc.displayPath,
|
||||
title: doc.title,
|
||||
...(doc.context && { context: doc.context }),
|
||||
hash: doc.hash,
|
||||
modifiedAt: doc.modifiedAt,
|
||||
bodyLength: doc.bodyLength,
|
||||
...(doc.body !== undefined && { body: doc.body }),
|
||||
}, null, 2);
|
||||
}
|
||||
|
||||
/**
|
||||
* Format a single DocumentResult as Markdown
|
||||
*/
|
||||
export function documentToMarkdown(doc: DocumentResult): string {
|
||||
let md = `# ${doc.title || doc.displayPath}\n\n`;
|
||||
if (doc.context) md += `**Context:** ${doc.context}\n\n`;
|
||||
md += `**File:** ${doc.displayPath}\n`;
|
||||
md += `**Modified:** ${doc.modifiedAt}\n\n`;
|
||||
if (doc.body !== undefined) {
|
||||
md += "---\n\n" + doc.body + "\n";
|
||||
}
|
||||
return md;
|
||||
}
|
||||
|
||||
/**
|
||||
* Format a single DocumentResult as XML
|
||||
*/
|
||||
export function documentToXml(doc: DocumentResult): string {
|
||||
let xml = `<?xml version="1.0" encoding="UTF-8"?>\n<document>\n`;
|
||||
xml += ` <file>${escapeXml(doc.displayPath)}</file>\n`;
|
||||
xml += ` <title>${escapeXml(doc.title)}</title>\n`;
|
||||
if (doc.context) xml += ` <context>${escapeXml(doc.context)}</context>\n`;
|
||||
xml += ` <hash>${escapeXml(doc.hash)}</hash>\n`;
|
||||
xml += ` <modifiedAt>${escapeXml(doc.modifiedAt)}</modifiedAt>\n`;
|
||||
xml += ` <bodyLength>${doc.bodyLength}</bodyLength>\n`;
|
||||
if (doc.body !== undefined) {
|
||||
xml += ` <body>${escapeXml(doc.body)}</body>\n`;
|
||||
}
|
||||
xml += `</document>`;
|
||||
return xml;
|
||||
}
|
||||
|
||||
/**
|
||||
* Format a single document to the specified format
|
||||
*/
|
||||
export function formatDocument(doc: DocumentResult, format: OutputFormat): string {
|
||||
switch (format) {
|
||||
case "json":
|
||||
return documentToJson(doc);
|
||||
case "md":
|
||||
return documentToMarkdown(doc);
|
||||
case "xml":
|
||||
return documentToXml(doc);
|
||||
default:
|
||||
// Default to markdown for CLI and other formats
|
||||
return documentToMarkdown(doc);
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Universal Format Function
|
||||
// =============================================================================
|
||||
|
||||
/**
|
||||
* Format search results to the specified output format
|
||||
*/
|
||||
export function formatSearchResults(
|
||||
results: SearchResult[],
|
||||
format: OutputFormat,
|
||||
opts: FormatOptions = {}
|
||||
): string {
|
||||
switch (format) {
|
||||
case "json":
|
||||
return searchResultsToJson(results, opts);
|
||||
case "csv":
|
||||
return searchResultsToCsv(results, opts);
|
||||
case "files":
|
||||
return searchResultsToFiles(results);
|
||||
case "md":
|
||||
return searchResultsToMarkdown(results, opts);
|
||||
case "xml":
|
||||
return searchResultsToXml(results, opts);
|
||||
case "cli":
|
||||
// CLI format should be handled separately with colors
|
||||
// Return a simple text version as fallback
|
||||
return searchResultsToMarkdown(results, opts);
|
||||
default:
|
||||
return searchResultsToJson(results, opts);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Format documents to the specified output format
|
||||
*/
|
||||
export function formatDocuments(
|
||||
results: MultiGetFile[],
|
||||
format: OutputFormat
|
||||
): string {
|
||||
switch (format) {
|
||||
case "json":
|
||||
return documentsToJson(results);
|
||||
case "csv":
|
||||
return documentsToCsv(results);
|
||||
case "files":
|
||||
return documentsToFiles(results);
|
||||
case "md":
|
||||
return documentsToMarkdown(results);
|
||||
case "xml":
|
||||
return documentsToXml(results);
|
||||
case "cli":
|
||||
// CLI format should be handled separately with colors
|
||||
return documentsToMarkdown(results);
|
||||
default:
|
||||
return documentsToJson(results);
|
||||
}
|
||||
}
|
||||
902
llm.test.ts
Normal file
902
llm.test.ts
Normal file
@ -0,0 +1,902 @@
|
||||
/**
|
||||
* llm.test.ts - Comprehensive unit tests for the LLM abstraction layer
|
||||
*
|
||||
* Run with: bun test llm.test.ts
|
||||
*
|
||||
* Tests use a mock HTTP server to simulate Ollama responses.
|
||||
*/
|
||||
|
||||
import { describe, test, expect, beforeAll, afterAll, beforeEach, afterEach } from "bun:test";
|
||||
import {
|
||||
Ollama,
|
||||
getDefaultOllama,
|
||||
setDefaultOllama,
|
||||
formatQueryForEmbedding,
|
||||
formatDocForEmbedding,
|
||||
type EmbeddingResult,
|
||||
type GenerateResult,
|
||||
type RerankDocumentResult,
|
||||
type TokenLogProb,
|
||||
} from "./llm.js";
|
||||
|
||||
// =============================================================================
|
||||
// Mock Server Setup
|
||||
// =============================================================================
|
||||
|
||||
type MockHandler = (body: unknown) => {
|
||||
status: number;
|
||||
body: unknown;
|
||||
};
|
||||
|
||||
const mockHandlers: Map<string, MockHandler> = new Map();
|
||||
let mockServerUrl: string;
|
||||
let mockCallLog: Array<{ path: string; body: unknown }> = [];
|
||||
|
||||
// Track original fetch
|
||||
const originalFetch = globalThis.fetch;
|
||||
|
||||
function installMockFetch(): void {
|
||||
globalThis.fetch = async (input: RequestInfo | URL, init?: RequestInit): Promise<Response> => {
|
||||
const url = typeof input === "string" ? input : input instanceof URL ? input.href : input.url;
|
||||
|
||||
// Only intercept calls to our mock server URL
|
||||
if (!url.startsWith(mockServerUrl)) {
|
||||
throw new Error(`TEST ERROR: Unexpected fetch to: ${url}`);
|
||||
}
|
||||
|
||||
const path = url.replace(mockServerUrl, "");
|
||||
const body = init?.body ? JSON.parse(init.body as string) : {};
|
||||
|
||||
// Log the call
|
||||
mockCallLog.push({ path, body });
|
||||
|
||||
const handler = mockHandlers.get(path);
|
||||
if (!handler) {
|
||||
return new Response(JSON.stringify({ error: "Not found" }), {
|
||||
status: 404,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
});
|
||||
}
|
||||
|
||||
const result = handler(body);
|
||||
return new Response(JSON.stringify(result.body), {
|
||||
status: result.status,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
});
|
||||
};
|
||||
}
|
||||
|
||||
function restoreFetch(): void {
|
||||
globalThis.fetch = originalFetch;
|
||||
}
|
||||
|
||||
// Setup before all tests
|
||||
beforeAll(() => {
|
||||
mockServerUrl = "http://mock-ollama:11434";
|
||||
installMockFetch();
|
||||
});
|
||||
|
||||
// Restore after all tests
|
||||
afterAll(() => {
|
||||
restoreFetch();
|
||||
});
|
||||
|
||||
// Clear call log and handlers before each test
|
||||
beforeEach(() => {
|
||||
mockCallLog = [];
|
||||
mockHandlers.clear();
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Helper Functions
|
||||
// =============================================================================
|
||||
|
||||
function createOllama(): Ollama {
|
||||
return new Ollama({ baseUrl: mockServerUrl });
|
||||
}
|
||||
|
||||
function setEmbedHandler(embeddings: number[][]): void {
|
||||
mockHandlers.set("/api/embed", () => ({
|
||||
status: 200,
|
||||
body: { embeddings },
|
||||
}));
|
||||
}
|
||||
|
||||
function setGenerateHandler(
|
||||
response: string,
|
||||
logprobs?: { tokens: string[]; token_logprobs: number[] }
|
||||
): void {
|
||||
mockHandlers.set("/api/generate", () => ({
|
||||
status: 200,
|
||||
body: {
|
||||
response,
|
||||
done: true,
|
||||
...(logprobs && { logprobs }),
|
||||
},
|
||||
}));
|
||||
}
|
||||
|
||||
function setModelShowHandler(exists: boolean, size?: number): void {
|
||||
mockHandlers.set("/api/show", () => {
|
||||
if (exists) {
|
||||
return {
|
||||
status: 200,
|
||||
body: { size: size ?? 1000000, modified_at: "2024-01-01T00:00:00Z" },
|
||||
};
|
||||
}
|
||||
return { status: 404, body: { error: "model not found" } };
|
||||
});
|
||||
}
|
||||
|
||||
function setPullHandler(success: boolean): void {
|
||||
mockHandlers.set("/api/pull", () => ({
|
||||
status: success ? 200 : 500,
|
||||
body: success ? { status: "success" } : { error: "failed" },
|
||||
}));
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Formatting Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("Formatting Functions", () => {
|
||||
test("formatQueryForEmbedding adds search task prefix", () => {
|
||||
const result = formatQueryForEmbedding("how to deploy");
|
||||
expect(result).toBe("task: search result | query: how to deploy");
|
||||
});
|
||||
|
||||
test("formatQueryForEmbedding handles empty query", () => {
|
||||
const result = formatQueryForEmbedding("");
|
||||
expect(result).toBe("task: search result | query: ");
|
||||
});
|
||||
|
||||
test("formatDocForEmbedding adds title and text prefix", () => {
|
||||
const result = formatDocForEmbedding("Document content", "My Title");
|
||||
expect(result).toBe("title: My Title | text: Document content");
|
||||
});
|
||||
|
||||
test("formatDocForEmbedding handles missing title", () => {
|
||||
const result = formatDocForEmbedding("Document content");
|
||||
expect(result).toBe("title: none | text: Document content");
|
||||
});
|
||||
|
||||
test("formatDocForEmbedding handles empty content", () => {
|
||||
const result = formatDocForEmbedding("", "Title");
|
||||
expect(result).toBe("title: Title | text: ");
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Ollama Constructor Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("Ollama Constructor", () => {
|
||||
test("uses default URL when not specified", () => {
|
||||
const ollama = new Ollama();
|
||||
expect(ollama.getBaseUrl()).toBe("http://localhost:11434");
|
||||
});
|
||||
|
||||
test("uses custom URL when specified", () => {
|
||||
const ollama = new Ollama({ baseUrl: "http://custom:9999" });
|
||||
expect(ollama.getBaseUrl()).toBe("http://custom:9999");
|
||||
});
|
||||
|
||||
test("respects OLLAMA_URL environment variable", () => {
|
||||
const originalEnv = process.env.OLLAMA_URL;
|
||||
process.env.OLLAMA_URL = "http://env-url:8888";
|
||||
|
||||
const ollama = new Ollama();
|
||||
expect(ollama.getBaseUrl()).toBe("http://env-url:8888");
|
||||
|
||||
// Restore
|
||||
if (originalEnv) {
|
||||
process.env.OLLAMA_URL = originalEnv;
|
||||
} else {
|
||||
delete process.env.OLLAMA_URL;
|
||||
}
|
||||
});
|
||||
|
||||
test("explicit baseUrl overrides environment variable", () => {
|
||||
const originalEnv = process.env.OLLAMA_URL;
|
||||
process.env.OLLAMA_URL = "http://env-url:8888";
|
||||
|
||||
const ollama = new Ollama({ baseUrl: "http://explicit:7777" });
|
||||
expect(ollama.getBaseUrl()).toBe("http://explicit:7777");
|
||||
|
||||
// Restore
|
||||
if (originalEnv) {
|
||||
process.env.OLLAMA_URL = originalEnv;
|
||||
} else {
|
||||
delete process.env.OLLAMA_URL;
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Embed Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("Ollama.embed", () => {
|
||||
test("returns embedding for query", async () => {
|
||||
const ollama = createOllama();
|
||||
const embedding = [0.1, 0.2, 0.3, 0.4, 0.5];
|
||||
setEmbedHandler([embedding]);
|
||||
|
||||
const result = await ollama.embed("test query", { model: "test-model", isQuery: true });
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result!.embedding).toEqual(embedding);
|
||||
expect(result!.model).toBe("test-model");
|
||||
|
||||
// Verify the request was formatted correctly
|
||||
expect(mockCallLog).toHaveLength(1);
|
||||
expect(mockCallLog[0].path).toBe("/api/embed");
|
||||
expect((mockCallLog[0].body as { input: string }).input).toContain("task: search result");
|
||||
});
|
||||
|
||||
test("returns embedding for document", async () => {
|
||||
const ollama = createOllama();
|
||||
const embedding = [0.5, 0.4, 0.3, 0.2, 0.1];
|
||||
setEmbedHandler([embedding]);
|
||||
|
||||
const result = await ollama.embed("doc content", {
|
||||
model: "test-model",
|
||||
isQuery: false,
|
||||
title: "Doc Title",
|
||||
});
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result!.embedding).toEqual(embedding);
|
||||
|
||||
// Verify document formatting
|
||||
expect((mockCallLog[0].body as { input: string }).input).toContain("title: Doc Title");
|
||||
expect((mockCallLog[0].body as { input: string }).input).toContain("text: doc content");
|
||||
});
|
||||
|
||||
test("returns null on API error", async () => {
|
||||
const ollama = createOllama();
|
||||
mockHandlers.set("/api/embed", () => ({ status: 500, body: { error: "Server error" } }));
|
||||
|
||||
const result = await ollama.embed("test", { model: "test-model" });
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
test("returns null on empty embeddings", async () => {
|
||||
const ollama = createOllama();
|
||||
setEmbedHandler([]);
|
||||
|
||||
const result = await ollama.embed("test", { model: "test-model" });
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
test("returns null on network error", async () => {
|
||||
const ollama = new Ollama({ baseUrl: "http://nonexistent:99999" });
|
||||
|
||||
// This will throw because our mock doesn't handle this URL
|
||||
const result = await ollama.embed("test", { model: "test-model" }).catch(() => null);
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
test("handles high-dimensional embeddings", async () => {
|
||||
const ollama = createOllama();
|
||||
const embedding = Array(768).fill(0).map((_, i) => i / 768);
|
||||
setEmbedHandler([embedding]);
|
||||
|
||||
const result = await ollama.embed("test", { model: "test-model" });
|
||||
expect(result!.embedding).toHaveLength(768);
|
||||
expect(result!.embedding[0]).toBeCloseTo(0, 5);
|
||||
expect(result!.embedding[767]).toBeCloseTo(767 / 768, 5);
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Generate Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("Ollama.generate", () => {
|
||||
test("returns generated text", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("Generated response text");
|
||||
|
||||
const result = await ollama.generate("prompt", { model: "test-model" });
|
||||
|
||||
expect(result).not.toBeNull();
|
||||
expect(result!.text).toBe("Generated response text");
|
||||
expect(result!.model).toBe("test-model");
|
||||
expect(result!.done).toBe(true);
|
||||
});
|
||||
|
||||
test("includes logprobs when requested", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", {
|
||||
tokens: ["yes"],
|
||||
token_logprobs: [-0.1],
|
||||
});
|
||||
|
||||
const result = await ollama.generate("prompt", { model: "test-model", logprobs: true });
|
||||
|
||||
expect(result!.logprobs).toBeDefined();
|
||||
expect(result!.logprobs).toHaveLength(1);
|
||||
expect(result!.logprobs![0].token).toBe("yes");
|
||||
expect(result!.logprobs![0].logprob).toBe(-0.1);
|
||||
});
|
||||
|
||||
test("handles multiple logprob tokens", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("hello world", {
|
||||
tokens: ["hello", " world"],
|
||||
token_logprobs: [-0.5, -0.3],
|
||||
});
|
||||
|
||||
const result = await ollama.generate("prompt", { model: "test-model", logprobs: true });
|
||||
|
||||
expect(result!.logprobs).toHaveLength(2);
|
||||
expect(result!.logprobs![0]).toEqual({ token: "hello", logprob: -0.5 });
|
||||
expect(result!.logprobs![1]).toEqual({ token: " world", logprob: -0.3 });
|
||||
});
|
||||
|
||||
test("sends maxTokens option", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("response");
|
||||
|
||||
await ollama.generate("prompt", { model: "test-model", maxTokens: 50 });
|
||||
|
||||
const body = mockCallLog[0].body as { options: { num_predict: number } };
|
||||
expect(body.options.num_predict).toBe(50);
|
||||
});
|
||||
|
||||
test("sends temperature option", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("response");
|
||||
|
||||
await ollama.generate("prompt", { model: "test-model", temperature: 0.7 });
|
||||
|
||||
const body = mockCallLog[0].body as { options: { temperature: number } };
|
||||
expect(body.options.temperature).toBe(0.7);
|
||||
});
|
||||
|
||||
test("sends raw option", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("response");
|
||||
|
||||
await ollama.generate("prompt", { model: "test-model", raw: true });
|
||||
|
||||
const body = mockCallLog[0].body as { raw: boolean };
|
||||
expect(body.raw).toBe(true);
|
||||
});
|
||||
|
||||
test("returns null on API error", async () => {
|
||||
const ollama = createOllama();
|
||||
mockHandlers.set("/api/generate", () => ({ status: 500, body: { error: "Error" } }));
|
||||
|
||||
const result = await ollama.generate("prompt", { model: "test-model" });
|
||||
expect(result).toBeNull();
|
||||
});
|
||||
|
||||
test("handles empty response", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("");
|
||||
|
||||
const result = await ollama.generate("prompt", { model: "test-model" });
|
||||
expect(result!.text).toBe("");
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Model Management Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("Ollama.modelExists", () => {
|
||||
test("returns true for existing model", async () => {
|
||||
const ollama = createOllama();
|
||||
setModelShowHandler(true, 5000000);
|
||||
|
||||
const result = await ollama.modelExists("test-model");
|
||||
|
||||
expect(result.exists).toBe(true);
|
||||
expect(result.name).toBe("test-model");
|
||||
expect(result.size).toBe(5000000);
|
||||
expect(result.modifiedAt).toBeDefined();
|
||||
});
|
||||
|
||||
test("returns false for non-existing model", async () => {
|
||||
const ollama = createOllama();
|
||||
setModelShowHandler(false);
|
||||
|
||||
const result = await ollama.modelExists("nonexistent-model");
|
||||
|
||||
expect(result.exists).toBe(false);
|
||||
expect(result.name).toBe("nonexistent-model");
|
||||
});
|
||||
|
||||
test("sends correct model name in request", async () => {
|
||||
const ollama = createOllama();
|
||||
setModelShowHandler(true);
|
||||
|
||||
await ollama.modelExists("specific-model:v1");
|
||||
|
||||
expect(mockCallLog[0].path).toBe("/api/show");
|
||||
expect((mockCallLog[0].body as { name: string }).name).toBe("specific-model:v1");
|
||||
});
|
||||
});
|
||||
|
||||
describe("Ollama.pullModel", () => {
|
||||
test("returns true on successful pull", async () => {
|
||||
const ollama = createOllama();
|
||||
setPullHandler(true);
|
||||
|
||||
const result = await ollama.pullModel("new-model");
|
||||
|
||||
expect(result).toBe(true);
|
||||
expect(mockCallLog[0].path).toBe("/api/pull");
|
||||
expect((mockCallLog[0].body as { name: string }).name).toBe("new-model");
|
||||
});
|
||||
|
||||
test("returns false on failed pull", async () => {
|
||||
const ollama = createOllama();
|
||||
setPullHandler(false);
|
||||
|
||||
const result = await ollama.pullModel("bad-model");
|
||||
expect(result).toBe(false);
|
||||
});
|
||||
|
||||
test("calls progress callback", async () => {
|
||||
const ollama = createOllama();
|
||||
setPullHandler(true);
|
||||
|
||||
let progressCalled = false;
|
||||
await ollama.pullModel("model", (progress) => {
|
||||
progressCalled = true;
|
||||
expect(progress).toBe(100);
|
||||
});
|
||||
|
||||
expect(progressCalled).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Query Expansion Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("Ollama.expandQuery", () => {
|
||||
test("returns original query plus expansions", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("variation one\nvariation two");
|
||||
|
||||
const result = await ollama.expandQuery("original query", "test-model");
|
||||
|
||||
expect(result).toContain("original query");
|
||||
expect(result[0]).toBe("original query");
|
||||
expect(result.length).toBeGreaterThanOrEqual(1);
|
||||
});
|
||||
|
||||
test("returns only original query on API failure", async () => {
|
||||
const ollama = createOllama();
|
||||
mockHandlers.set("/api/generate", () => ({ status: 500, body: { error: "Error" } }));
|
||||
|
||||
const result = await ollama.expandQuery("query", "test-model");
|
||||
|
||||
expect(result).toEqual(["query"]);
|
||||
});
|
||||
|
||||
test("filters out thinking tags from response", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("<think>some thinking</think>\nvariation one\nvariation two");
|
||||
|
||||
const result = await ollama.expandQuery("query", "test-model");
|
||||
|
||||
expect(result).not.toContain("<think>");
|
||||
expect(result.some((r) => r.includes("think"))).toBe(false);
|
||||
});
|
||||
|
||||
test("filters out very long variations", async () => {
|
||||
const ollama = createOllama();
|
||||
const longLine = "a".repeat(150);
|
||||
setGenerateHandler(`short variation\n${longLine}\nanother short`);
|
||||
|
||||
const result = await ollama.expandQuery("query", "test-model");
|
||||
|
||||
// Long variations (>100 chars) should be filtered
|
||||
expect(result.every((r) => r.length < 100)).toBe(true);
|
||||
});
|
||||
|
||||
test("respects numVariations parameter", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("one\ntwo\nthree\nfour\nfive");
|
||||
|
||||
const result = await ollama.expandQuery("query", "test-model", 3);
|
||||
|
||||
// Original + up to 3 variations
|
||||
expect(result.length).toBeLessThanOrEqual(4);
|
||||
});
|
||||
|
||||
test("sends correct prompt format", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("variation");
|
||||
|
||||
await ollama.expandQuery("test query", "test-model", 2);
|
||||
|
||||
const body = mockCallLog[0].body as { prompt: string };
|
||||
expect(body.prompt).toContain('Query: "test query"');
|
||||
expect(body.prompt).toContain("generate 2 alternative queries");
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Reranking Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("Ollama.rerankerLogprobsCheck", () => {
|
||||
test("returns relevance judgments for documents", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
const docs = [
|
||||
{ file: "doc1.md", text: "Relevant content" },
|
||||
{ file: "doc2.md", text: "Other content" },
|
||||
];
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck("query", docs, { model: "test-model" });
|
||||
|
||||
expect(results).toHaveLength(2);
|
||||
expect(results[0].file).toBe("doc1.md");
|
||||
expect(results[0].relevant).toBe(true);
|
||||
expect(results[0].rawToken).toBe("yes");
|
||||
});
|
||||
|
||||
test("parses yes with high confidence correctly", async () => {
|
||||
const ollama = createOllama();
|
||||
// -0.1 logprob = ~0.905 confidence
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].relevant).toBe(true);
|
||||
expect(results[0].confidence).toBeCloseTo(Math.exp(-0.1), 3);
|
||||
expect(results[0].score).toBeGreaterThan(0.9);
|
||||
expect(results[0].logprob).toBe(-0.1);
|
||||
});
|
||||
|
||||
test("parses yes with low confidence correctly", async () => {
|
||||
const ollama = createOllama();
|
||||
// -2.0 logprob = ~0.135 confidence
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-2.0] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].relevant).toBe(true);
|
||||
expect(results[0].confidence).toBeCloseTo(Math.exp(-2.0), 3);
|
||||
expect(results[0].score).toBeLessThan(0.6);
|
||||
});
|
||||
|
||||
test("parses no with high confidence correctly", async () => {
|
||||
const ollama = createOllama();
|
||||
// -0.05 logprob = ~0.95 confidence
|
||||
setGenerateHandler("no", { tokens: ["no"], token_logprobs: [-0.05] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].relevant).toBe(false);
|
||||
expect(results[0].confidence).toBeCloseTo(Math.exp(-0.05), 3);
|
||||
expect(results[0].score).toBeLessThan(0.1); // Low score for confident "no"
|
||||
});
|
||||
|
||||
test("parses no with low confidence correctly", async () => {
|
||||
const ollama = createOllama();
|
||||
// -1.5 logprob = ~0.22 confidence
|
||||
setGenerateHandler("no", { tokens: ["no"], token_logprobs: [-1.5] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].relevant).toBe(false);
|
||||
expect(results[0].score).toBeGreaterThan(0.3); // Higher score for uncertain "no"
|
||||
});
|
||||
|
||||
test("handles unknown token", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("maybe", { tokens: ["maybe"], token_logprobs: [-0.5] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].relevant).toBe(false);
|
||||
expect(results[0].score).toBe(0.3); // Neutral score
|
||||
});
|
||||
|
||||
test("handles API failure gracefully", async () => {
|
||||
const ollama = createOllama();
|
||||
mockHandlers.set("/api/generate", () => ({ status: 500, body: { error: "Error" } }));
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].relevant).toBe(false);
|
||||
expect(results[0].score).toBe(0);
|
||||
expect(results[0].confidence).toBe(0);
|
||||
});
|
||||
|
||||
test("respects batchSize option", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
const docs = Array(10).fill(null).map((_, i) => ({
|
||||
file: `doc${i}.md`,
|
||||
text: `content ${i}`,
|
||||
}));
|
||||
|
||||
await ollama.rerankerLogprobsCheck("query", docs, { model: "test-model", batchSize: 3 });
|
||||
|
||||
// Should process in batches: 3 + 3 + 3 + 1 = 10 calls
|
||||
expect(mockCallLog).toHaveLength(10);
|
||||
});
|
||||
|
||||
test("sends correct prompt format", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
await ollama.rerankerLogprobsCheck(
|
||||
"search query",
|
||||
[{ file: "test.md", text: "document content", title: "Test Doc" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
const body = mockCallLog[0].body as { prompt: string; raw: boolean; logprobs: boolean };
|
||||
expect(body.prompt).toContain("<Query>: search query");
|
||||
expect(body.prompt).toContain("<Document Title>: Test Doc");
|
||||
expect(body.prompt).toContain("document content");
|
||||
expect(body.raw).toBe(true);
|
||||
expect(body.logprobs).toBe(true);
|
||||
});
|
||||
|
||||
test("uses filename as title when title not provided", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "path/to/document.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
const body = mockCallLog[0].body as { prompt: string };
|
||||
expect(body.prompt).toContain("<Document Title>: document");
|
||||
});
|
||||
|
||||
test("truncates long documents", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
const longText = "x".repeat(10000);
|
||||
await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: longText }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
const body = mockCallLog[0].body as { prompt: string };
|
||||
// Should be truncated to ~4000 chars + "..."
|
||||
expect(body.prompt.length).toBeLessThan(10000);
|
||||
expect(body.prompt).toContain("...");
|
||||
});
|
||||
});
|
||||
|
||||
describe("Ollama.rerank", () => {
|
||||
test("returns sorted results by score", async () => {
|
||||
const ollama = createOllama();
|
||||
|
||||
// First call returns "no", second returns "yes"
|
||||
let callCount = 0;
|
||||
mockHandlers.set("/api/generate", () => {
|
||||
callCount++;
|
||||
if (callCount === 1) {
|
||||
return { status: 200, body: { response: "no", done: true, logprobs: { tokens: ["no"], token_logprobs: [-0.1] } } };
|
||||
}
|
||||
return { status: 200, body: { response: "yes", done: true, logprobs: { tokens: ["yes"], token_logprobs: [-0.1] } } };
|
||||
});
|
||||
|
||||
const docs = [
|
||||
{ file: "low.md", text: "irrelevant" },
|
||||
{ file: "high.md", text: "relevant" },
|
||||
];
|
||||
|
||||
const result = await ollama.rerank("query", docs, { model: "test-model" });
|
||||
|
||||
expect(result.results).toHaveLength(2);
|
||||
expect(result.results[0].file).toBe("high.md"); // Higher score first
|
||||
expect(result.results[0].score).toBeGreaterThan(result.results[1].score);
|
||||
});
|
||||
|
||||
test("includes model in result", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
const result = await ollama.rerank("query", [{ file: "doc.md", text: "content" }], {
|
||||
model: "custom-reranker",
|
||||
});
|
||||
|
||||
expect(result.model).toBe("custom-reranker");
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Default Ollama Singleton Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("Default Ollama Singleton", () => {
|
||||
afterEach(() => {
|
||||
setDefaultOllama(null);
|
||||
});
|
||||
|
||||
test("getDefaultOllama creates instance on first call", () => {
|
||||
const ollama = getDefaultOllama();
|
||||
expect(ollama).toBeInstanceOf(Ollama);
|
||||
});
|
||||
|
||||
test("getDefaultOllama returns same instance on subsequent calls", () => {
|
||||
const ollama1 = getDefaultOllama();
|
||||
const ollama2 = getDefaultOllama();
|
||||
expect(ollama1).toBe(ollama2);
|
||||
});
|
||||
|
||||
test("setDefaultOllama allows replacing the singleton", () => {
|
||||
const custom = new Ollama({ baseUrl: "http://custom:1234" });
|
||||
setDefaultOllama(custom);
|
||||
|
||||
const result = getDefaultOllama();
|
||||
expect(result).toBe(custom);
|
||||
expect(result.getBaseUrl()).toBe("http://custom:1234");
|
||||
});
|
||||
|
||||
test("setDefaultOllama with null resets singleton", () => {
|
||||
const original = getDefaultOllama();
|
||||
setDefaultOllama(null);
|
||||
const newInstance = getDefaultOllama();
|
||||
|
||||
expect(newInstance).not.toBe(original);
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Logprob Math Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("Logprob Mathematics", () => {
|
||||
test("logprob 0 = 100% confidence", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [0] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].confidence).toBe(1.0);
|
||||
expect(results[0].score).toBe(1.0); // 0.5 + 0.5 * 1.0
|
||||
});
|
||||
|
||||
test("logprob -ln(2) ≈ 50% confidence", async () => {
|
||||
const ollama = createOllama();
|
||||
const logprob = -Math.log(2); // ≈ -0.693
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [logprob] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].confidence).toBeCloseTo(0.5, 3);
|
||||
expect(results[0].score).toBeCloseTo(0.75, 3); // 0.5 + 0.5 * 0.5
|
||||
});
|
||||
|
||||
test("very negative logprob = very low confidence", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-10] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].confidence).toBeLessThan(0.0001);
|
||||
expect(results[0].score).toBeCloseTo(0.5, 2); // Nearly just the base 0.5
|
||||
});
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Edge Cases
|
||||
// =============================================================================
|
||||
|
||||
describe("Edge Cases", () => {
|
||||
test("handles empty document list", async () => {
|
||||
const ollama = createOllama();
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck("query", [], { model: "test-model" });
|
||||
expect(results).toHaveLength(0);
|
||||
});
|
||||
|
||||
test("handles very short document text", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "x" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(1);
|
||||
});
|
||||
|
||||
test("handles unicode in queries and documents", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"日本語クエリ",
|
||||
[{ file: "doc.md", text: "日本語コンテンツ 🎉" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results).toHaveLength(1);
|
||||
|
||||
const body = mockCallLog[0].body as { prompt: string };
|
||||
expect(body.prompt).toContain("日本語クエリ");
|
||||
expect(body.prompt).toContain("日本語コンテンツ");
|
||||
});
|
||||
|
||||
test("handles special characters in file paths", async () => {
|
||||
const ollama = createOllama();
|
||||
setGenerateHandler("yes", { tokens: ["yes"], token_logprobs: [-0.1] });
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "path/to/file with spaces.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
expect(results[0].file).toBe("path/to/file with spaces.md");
|
||||
});
|
||||
|
||||
test("handles missing logprobs in response", async () => {
|
||||
const ollama = createOllama();
|
||||
// Response without logprobs
|
||||
mockHandlers.set("/api/generate", () => ({
|
||||
status: 200,
|
||||
body: { response: "yes", done: true },
|
||||
}));
|
||||
|
||||
const results = await ollama.rerankerLogprobsCheck(
|
||||
"query",
|
||||
[{ file: "doc.md", text: "content" }],
|
||||
{ model: "test-model" }
|
||||
);
|
||||
|
||||
// Should still work, with logprob defaulting to 0
|
||||
expect(results[0].logprob).toBe(0);
|
||||
});
|
||||
});
|
||||
539
llm.ts
Normal file
539
llm.ts
Normal file
@ -0,0 +1,539 @@
|
||||
/**
|
||||
* llm.ts - LLM abstraction layer for QMD
|
||||
*
|
||||
* Provides a clean interface for LLM operations with an Ollama implementation.
|
||||
* All raw fetch calls to LLM APIs should go through this module.
|
||||
*/
|
||||
|
||||
// =============================================================================
|
||||
// Types
|
||||
// =============================================================================
|
||||
|
||||
/**
|
||||
* Token with log probability
|
||||
*/
|
||||
export type TokenLogProb = {
|
||||
token: string;
|
||||
logprob: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* Embedding result
|
||||
*/
|
||||
export type EmbeddingResult = {
|
||||
embedding: number[];
|
||||
model: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Generation result with optional logprobs
|
||||
*/
|
||||
export type GenerateResult = {
|
||||
text: string;
|
||||
model: string;
|
||||
logprobs?: TokenLogProb[];
|
||||
done: boolean;
|
||||
};
|
||||
|
||||
/**
|
||||
* Rerank result for a single document
|
||||
*/
|
||||
export type RerankDocumentResult = {
|
||||
file: string;
|
||||
relevant: boolean;
|
||||
confidence: number;
|
||||
score: number;
|
||||
rawToken: string;
|
||||
logprob: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* Batch rerank result
|
||||
*/
|
||||
export type RerankResult = {
|
||||
results: RerankDocumentResult[];
|
||||
model: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Model info
|
||||
*/
|
||||
export type ModelInfo = {
|
||||
name: string;
|
||||
exists: boolean;
|
||||
size?: number;
|
||||
modifiedAt?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Options for embedding
|
||||
*/
|
||||
export type EmbedOptions = {
|
||||
model: string;
|
||||
isQuery?: boolean;
|
||||
title?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* Options for text generation
|
||||
*/
|
||||
export type GenerateOptions = {
|
||||
model: string;
|
||||
maxTokens?: number;
|
||||
temperature?: number;
|
||||
logprobs?: boolean;
|
||||
raw?: boolean;
|
||||
stop?: string[];
|
||||
};
|
||||
|
||||
/**
|
||||
* Options for reranking
|
||||
*/
|
||||
export type RerankOptions = {
|
||||
model: string;
|
||||
batchSize?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* Document to rerank
|
||||
*/
|
||||
export type RerankDocument = {
|
||||
file: string;
|
||||
text: string;
|
||||
title?: string;
|
||||
};
|
||||
|
||||
// =============================================================================
|
||||
// LLM Interface
|
||||
// =============================================================================
|
||||
|
||||
/**
|
||||
* Abstract LLM interface - implement this for different backends
|
||||
*/
|
||||
export interface LLM {
|
||||
/**
|
||||
* Get embeddings for text
|
||||
*/
|
||||
embed(text: string, options: EmbedOptions): Promise<EmbeddingResult | null>;
|
||||
|
||||
/**
|
||||
* Generate text completion
|
||||
*/
|
||||
generate(prompt: string, options: GenerateOptions): Promise<GenerateResult | null>;
|
||||
|
||||
/**
|
||||
* Check if a model exists
|
||||
*/
|
||||
modelExists(model: string): Promise<ModelInfo>;
|
||||
|
||||
/**
|
||||
* Pull a model (download if not available)
|
||||
*/
|
||||
pullModel(model: string, onProgress?: (progress: number) => void): Promise<boolean>;
|
||||
|
||||
// ==========================================================================
|
||||
// High-level abstractions
|
||||
// ==========================================================================
|
||||
|
||||
/**
|
||||
* Expand a search query into multiple variations
|
||||
*/
|
||||
expandQuery(query: string, model: string, numVariations?: number): Promise<string[]>;
|
||||
|
||||
/**
|
||||
* Rerank documents by relevance to a query
|
||||
* Returns list of documents with relevance scores and boolean judgments
|
||||
*/
|
||||
rerank(query: string, documents: RerankDocument[], options: RerankOptions): Promise<RerankResult>;
|
||||
|
||||
/**
|
||||
* Quick relevance check - returns just boolean judgments with logprobs
|
||||
* More efficient than full rerank when you just need yes/no
|
||||
*/
|
||||
rerankerLogprobsCheck(query: string, documents: RerankDocument[], options: RerankOptions): Promise<RerankDocumentResult[]>;
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Ollama Implementation
|
||||
// =============================================================================
|
||||
|
||||
export type OllamaConfig = {
|
||||
baseUrl?: string;
|
||||
defaultEmbedModel?: string;
|
||||
defaultGenerateModel?: string;
|
||||
defaultRerankModel?: string;
|
||||
};
|
||||
|
||||
const DEFAULT_OLLAMA_URL = "http://localhost:11434";
|
||||
const DEFAULT_EMBED_MODEL = "embeddinggemma";
|
||||
const DEFAULT_GENERATE_MODEL = "qwen3:0.6b";
|
||||
const DEFAULT_RERANK_MODEL = "ExpedientFalcon/qwen3-reranker:0.6b-q8_0";
|
||||
|
||||
/**
|
||||
* Format text for embedding query
|
||||
*/
|
||||
export function formatQueryForEmbedding(query: string): string {
|
||||
return `task: search result | query: ${query}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Format text for embedding document
|
||||
*/
|
||||
export function formatDocForEmbedding(text: string, title?: string): string {
|
||||
return `title: ${title || "none"} | text: ${text}`;
|
||||
}
|
||||
|
||||
/**
|
||||
* Ollama LLM implementation
|
||||
*/
|
||||
export class Ollama implements LLM {
|
||||
private baseUrl: string;
|
||||
private defaultEmbedModel: string;
|
||||
private defaultGenerateModel: string;
|
||||
private defaultRerankModel: string;
|
||||
|
||||
constructor(config: OllamaConfig = {}) {
|
||||
this.baseUrl = config.baseUrl || process.env.OLLAMA_URL || DEFAULT_OLLAMA_URL;
|
||||
this.defaultEmbedModel = config.defaultEmbedModel || DEFAULT_EMBED_MODEL;
|
||||
this.defaultGenerateModel = config.defaultGenerateModel || DEFAULT_GENERATE_MODEL;
|
||||
this.defaultRerankModel = config.defaultRerankModel || DEFAULT_RERANK_MODEL;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the base URL for this Ollama instance
|
||||
*/
|
||||
getBaseUrl(): string {
|
||||
return this.baseUrl;
|
||||
}
|
||||
|
||||
// ==========================================================================
|
||||
// Core API methods
|
||||
// ==========================================================================
|
||||
|
||||
async embed(text: string, options: EmbedOptions): Promise<EmbeddingResult | null> {
|
||||
const model = options.model || this.defaultEmbedModel;
|
||||
const formatted = options.isQuery
|
||||
? formatQueryForEmbedding(text)
|
||||
: formatDocForEmbedding(text, options.title);
|
||||
|
||||
try {
|
||||
const response = await fetch(`${this.baseUrl}/api/embed`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ model, input: formatted }),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const data = await response.json() as { embeddings?: number[][] };
|
||||
if (!data.embeddings?.[0]) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
embedding: data.embeddings[0],
|
||||
model,
|
||||
};
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
async generate(prompt: string, options: GenerateOptions): Promise<GenerateResult | null> {
|
||||
const model = options.model || this.defaultGenerateModel;
|
||||
|
||||
const requestBody: Record<string, unknown> = {
|
||||
model,
|
||||
prompt,
|
||||
stream: false,
|
||||
options: {
|
||||
num_predict: options.maxTokens ?? 150,
|
||||
temperature: options.temperature ?? 0,
|
||||
},
|
||||
};
|
||||
|
||||
if (options.logprobs) {
|
||||
requestBody.logprobs = true;
|
||||
}
|
||||
|
||||
if (options.raw) {
|
||||
requestBody.raw = true;
|
||||
}
|
||||
|
||||
if (options.stop) {
|
||||
(requestBody.options as Record<string, unknown>).stop = options.stop;
|
||||
}
|
||||
|
||||
try {
|
||||
const response = await fetch(`${this.baseUrl}/api/generate`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(requestBody),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const data = await response.json() as {
|
||||
response?: string;
|
||||
done?: boolean;
|
||||
logprobs?: { tokens?: string[]; token_logprobs?: number[] };
|
||||
};
|
||||
|
||||
// Parse logprobs if present
|
||||
let logprobs: TokenLogProb[] | undefined;
|
||||
if (data.logprobs?.tokens && data.logprobs?.token_logprobs) {
|
||||
logprobs = data.logprobs.tokens.map((token, i) => ({
|
||||
token,
|
||||
logprob: data.logprobs!.token_logprobs![i],
|
||||
}));
|
||||
}
|
||||
|
||||
return {
|
||||
text: data.response || "",
|
||||
model,
|
||||
logprobs,
|
||||
done: data.done ?? true,
|
||||
};
|
||||
} catch {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
async modelExists(model: string): Promise<ModelInfo> {
|
||||
try {
|
||||
const response = await fetch(`${this.baseUrl}/api/show`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ name: model }),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
return { name: model, exists: false };
|
||||
}
|
||||
|
||||
const data = await response.json() as {
|
||||
size?: number;
|
||||
modified_at?: string;
|
||||
};
|
||||
|
||||
return {
|
||||
name: model,
|
||||
exists: true,
|
||||
size: data.size,
|
||||
modifiedAt: data.modified_at,
|
||||
};
|
||||
} catch {
|
||||
return { name: model, exists: false };
|
||||
}
|
||||
}
|
||||
|
||||
async pullModel(model: string, onProgress?: (progress: number) => void): Promise<boolean> {
|
||||
try {
|
||||
const response = await fetch(`${this.baseUrl}/api/pull`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ name: model, stream: false }),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
return false;
|
||||
}
|
||||
|
||||
// For non-streaming, we just wait for completion
|
||||
await response.json();
|
||||
onProgress?.(100);
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// ==========================================================================
|
||||
// High-level abstractions
|
||||
// ==========================================================================
|
||||
|
||||
async expandQuery(query: string, model?: string, numVariations: number = 2): Promise<string[]> {
|
||||
const useModel = model || this.defaultGenerateModel;
|
||||
|
||||
const prompt = `You are a search query expander. Given a search query, generate ${numVariations} alternative queries that would help find relevant documents.
|
||||
|
||||
Rules:
|
||||
- Use synonyms and related terminology (e.g., "craft" → "craftsmanship", "quality", "excellence")
|
||||
- Rephrase to capture different angles (e.g., "engineering culture" → "technical excellence", "developer practices")
|
||||
- Keep proper nouns and named concepts exactly as written (e.g., "Build a Business", "Stripe", "Shopify")
|
||||
- Each variation should be 3-8 words, natural search terms
|
||||
- Do NOT just append words like "search" or "find" or "documents"
|
||||
|
||||
Query: "${query}"
|
||||
|
||||
Output exactly ${numVariations} variations, one per line, no numbering or bullets:`;
|
||||
|
||||
const result = await this.generate(prompt, {
|
||||
model: useModel,
|
||||
maxTokens: 150,
|
||||
temperature: 0,
|
||||
});
|
||||
|
||||
if (!result) {
|
||||
return [query];
|
||||
}
|
||||
|
||||
// Parse response - filter out thinking tags and clean up
|
||||
const cleanText = result.text.replace(/<think>[\s\S]*?<\/think>/g, "").trim();
|
||||
const lines = cleanText
|
||||
.split("\n")
|
||||
.map((l) => l.trim())
|
||||
.filter((l) => l.length > 2 && l.length < 100 && !l.startsWith("<"));
|
||||
|
||||
return [query, ...lines.slice(0, numVariations)];
|
||||
}
|
||||
|
||||
async rerank(
|
||||
query: string,
|
||||
documents: RerankDocument[],
|
||||
options: RerankOptions
|
||||
): Promise<RerankResult> {
|
||||
const results = await this.rerankerLogprobsCheck(query, documents, options);
|
||||
|
||||
return {
|
||||
results: results.sort((a, b) => b.score - a.score),
|
||||
model: options.model || this.defaultRerankModel,
|
||||
};
|
||||
}
|
||||
|
||||
async rerankerLogprobsCheck(
|
||||
query: string,
|
||||
documents: RerankDocument[],
|
||||
options: RerankOptions
|
||||
): Promise<RerankDocumentResult[]> {
|
||||
const model = options.model || this.defaultRerankModel;
|
||||
const batchSize = options.batchSize || 5;
|
||||
|
||||
const results: RerankDocumentResult[] = [];
|
||||
|
||||
// Process in batches
|
||||
for (let i = 0; i < documents.length; i += batchSize) {
|
||||
const batch = documents.slice(i, i + batchSize);
|
||||
const batchResults = await Promise.all(
|
||||
batch.map((doc) => this.rerankSingle(query, doc, model))
|
||||
);
|
||||
results.push(...batchResults);
|
||||
}
|
||||
|
||||
return results;
|
||||
}
|
||||
|
||||
/**
|
||||
* Rerank a single document - internal helper
|
||||
*/
|
||||
private async rerankSingle(
|
||||
query: string,
|
||||
doc: RerankDocument,
|
||||
model: string
|
||||
): Promise<RerankDocumentResult> {
|
||||
const systemPrompt = `Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".`;
|
||||
|
||||
const instruct = `Given a search query, determine if the following document is relevant to the query. Consider both direct matches and related concepts.`;
|
||||
|
||||
const docTitle = doc.title || doc.file.split("/").pop()?.replace(/\.md$/, "") || doc.file;
|
||||
const docPreview = doc.text.length > 4000 ? doc.text.substring(0, 4000) + "..." : doc.text;
|
||||
|
||||
// Qwen3-reranker prompt format with empty think tags
|
||||
const prompt = `<|im_start|>system
|
||||
${systemPrompt}<|im_end|>
|
||||
<|im_start|>user
|
||||
<Instruct>: ${instruct}
|
||||
<Query>: ${query}
|
||||
<Document Title>: ${docTitle}
|
||||
<Document>: ${docPreview}<|im_end|>
|
||||
<|im_start|>assistant
|
||||
<think>
|
||||
|
||||
</think>
|
||||
|
||||
`;
|
||||
|
||||
const result = await this.generate(prompt, {
|
||||
model,
|
||||
maxTokens: 1,
|
||||
temperature: 0,
|
||||
logprobs: true,
|
||||
raw: true,
|
||||
});
|
||||
|
||||
if (!result) {
|
||||
return {
|
||||
file: doc.file,
|
||||
relevant: false,
|
||||
confidence: 0,
|
||||
score: 0,
|
||||
rawToken: "",
|
||||
logprob: 0,
|
||||
};
|
||||
}
|
||||
|
||||
return this.parseRerankResponse(doc.file, result);
|
||||
}
|
||||
|
||||
/**
|
||||
* Parse rerank response into structured result
|
||||
*/
|
||||
private parseRerankResponse(file: string, result: GenerateResult): RerankDocumentResult {
|
||||
const token = result.text.toLowerCase().trim();
|
||||
const logprob = result.logprobs?.[0]?.logprob ?? 0;
|
||||
const confidence = Math.exp(logprob);
|
||||
|
||||
let relevant: boolean;
|
||||
let score: number;
|
||||
|
||||
if (token.startsWith("yes")) {
|
||||
relevant = true;
|
||||
// Score: 0.5 base + up to 0.5 from confidence
|
||||
score = 0.5 + 0.5 * confidence;
|
||||
} else if (token.startsWith("no")) {
|
||||
relevant = false;
|
||||
// Score: up to 0.5 based on uncertainty (1 - confidence)
|
||||
score = 0.5 * (1 - confidence);
|
||||
} else {
|
||||
// Unknown token - neutral score
|
||||
relevant = false;
|
||||
score = 0.3;
|
||||
}
|
||||
|
||||
return {
|
||||
file,
|
||||
relevant,
|
||||
confidence,
|
||||
score,
|
||||
rawToken: result.logprobs?.[0]?.token ?? token,
|
||||
logprob,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Singleton for default Ollama instance
|
||||
// =============================================================================
|
||||
|
||||
let defaultOllama: Ollama | null = null;
|
||||
|
||||
/**
|
||||
* Get the default Ollama instance (creates one if needed)
|
||||
*/
|
||||
export function getDefaultOllama(): Ollama {
|
||||
if (!defaultOllama) {
|
||||
defaultOllama = new Ollama();
|
||||
}
|
||||
return defaultOllama;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set a custom default Ollama instance (useful for testing)
|
||||
*/
|
||||
export function setDefaultOllama(ollama: Ollama | null): void {
|
||||
defaultOllama = ollama;
|
||||
}
|
||||
870
mcp.test.ts
Normal file
870
mcp.test.ts
Normal file
@ -0,0 +1,870 @@
|
||||
/**
|
||||
* MCP Server Tests
|
||||
*
|
||||
* Tests all MCP tools, resources, and prompts.
|
||||
* Uses mocked Ollama responses and a test database.
|
||||
*/
|
||||
|
||||
import { describe, test, expect, beforeAll, afterAll, beforeEach, afterEach } from "bun:test";
|
||||
import { Database } from "bun:sqlite";
|
||||
import * as sqliteVec from "sqlite-vec";
|
||||
import { McpServer, ResourceTemplate } from "@modelcontextprotocol/sdk/server/mcp.js";
|
||||
import { z } from "zod";
|
||||
import { setDefaultOllama, Ollama } from "./llm";
|
||||
|
||||
// =============================================================================
|
||||
// Mock Ollama
|
||||
// =============================================================================
|
||||
|
||||
const OLLAMA_URL = "http://localhost:11434";
|
||||
const originalFetch = globalThis.fetch;
|
||||
|
||||
const mockOllamaResponses: Record<string, (body: unknown) => Response> = {
|
||||
"/api/embed": () => {
|
||||
const embedding = Array(768).fill(0).map(() => Math.random());
|
||||
return new Response(JSON.stringify({ embeddings: [embedding] }), {
|
||||
status: 200,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
});
|
||||
},
|
||||
"/api/generate": (body: unknown) => {
|
||||
const reqBody = body as { prompt?: string };
|
||||
if (reqBody.prompt?.includes("Judge") || reqBody.prompt?.includes("Document")) {
|
||||
return new Response(JSON.stringify({
|
||||
response: "yes",
|
||||
done: true,
|
||||
logprobs: { tokens: ["yes"], token_logprobs: [-0.1] },
|
||||
}), { status: 200, headers: { "Content-Type": "application/json" } });
|
||||
} else {
|
||||
return new Response(JSON.stringify({
|
||||
response: "expanded query variation 1\nexpanded query variation 2",
|
||||
done: true,
|
||||
}), { status: 200, headers: { "Content-Type": "application/json" } });
|
||||
}
|
||||
},
|
||||
"/api/show": () => {
|
||||
return new Response(JSON.stringify({ size: 1000000 }), {
|
||||
status: 200,
|
||||
headers: { "Content-Type": "application/json" },
|
||||
});
|
||||
},
|
||||
};
|
||||
|
||||
function mockFetch(input: RequestInfo | URL, init?: RequestInit): Promise<Response> {
|
||||
const url = typeof input === "string" ? input : input.toString();
|
||||
|
||||
if (url.startsWith(OLLAMA_URL)) {
|
||||
const path = url.replace(OLLAMA_URL, "");
|
||||
const handler = mockOllamaResponses[path];
|
||||
if (handler) {
|
||||
const body = init?.body ? JSON.parse(init.body as string) : {};
|
||||
return Promise.resolve(handler(body));
|
||||
}
|
||||
throw new Error(`Unmocked Ollama endpoint: ${path}`);
|
||||
}
|
||||
|
||||
throw new Error(`Unexpected fetch call to: ${url}`);
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Test Database Setup
|
||||
// =============================================================================
|
||||
|
||||
let testDb: Database;
|
||||
let testDbPath: string;
|
||||
|
||||
function initTestDatabase(db: Database): void {
|
||||
sqliteVec.load(db);
|
||||
db.exec("PRAGMA journal_mode = WAL");
|
||||
|
||||
db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS collections (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
pwd TEXT NOT NULL,
|
||||
glob_pattern TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
context TEXT,
|
||||
UNIQUE(pwd, glob_pattern)
|
||||
)
|
||||
`);
|
||||
|
||||
db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS path_contexts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
path_prefix TEXT NOT NULL UNIQUE,
|
||||
context TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
)
|
||||
`);
|
||||
|
||||
db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS ollama_cache (
|
||||
hash TEXT PRIMARY KEY,
|
||||
result TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL
|
||||
)
|
||||
`);
|
||||
|
||||
db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
collection_id INTEGER NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
hash TEXT NOT NULL,
|
||||
filepath TEXT NOT NULL,
|
||||
display_path TEXT NOT NULL DEFAULT '',
|
||||
body TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL,
|
||||
modified_at TEXT NOT NULL,
|
||||
active INTEGER NOT NULL DEFAULT 1,
|
||||
FOREIGN KEY (collection_id) REFERENCES collections(id)
|
||||
)
|
||||
`);
|
||||
|
||||
db.exec(`
|
||||
CREATE TABLE IF NOT EXISTS content_vectors (
|
||||
hash TEXT NOT NULL,
|
||||
seq INTEGER NOT NULL DEFAULT 0,
|
||||
pos INTEGER NOT NULL DEFAULT 0,
|
||||
model TEXT NOT NULL,
|
||||
embedded_at TEXT NOT NULL,
|
||||
PRIMARY KEY (hash, seq)
|
||||
)
|
||||
`);
|
||||
|
||||
db.exec(`
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
|
||||
name, body,
|
||||
content='documents',
|
||||
content_rowid='id',
|
||||
tokenize='porter unicode61'
|
||||
)
|
||||
`);
|
||||
|
||||
db.exec(`
|
||||
CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN
|
||||
INSERT INTO documents_fts(rowid, name, body) VALUES (new.id, new.name, new.body);
|
||||
END
|
||||
`);
|
||||
|
||||
// Create vector table
|
||||
db.exec(`CREATE VIRTUAL TABLE IF NOT EXISTS vectors_vec USING vec0(hash_seq TEXT PRIMARY KEY, embedding float[768])`);
|
||||
}
|
||||
|
||||
function seedTestData(db: Database): void {
|
||||
const now = new Date().toISOString();
|
||||
|
||||
// Create a collection
|
||||
db.prepare(`INSERT INTO collections (pwd, glob_pattern, created_at, context) VALUES (?, ?, ?, ?)`).run(
|
||||
"/test/docs",
|
||||
"**/*.md",
|
||||
now,
|
||||
"Test documentation collection"
|
||||
);
|
||||
|
||||
// Add path context
|
||||
db.prepare(`INSERT INTO path_contexts (path_prefix, context, created_at) VALUES (?, ?, ?)`).run(
|
||||
"/test/docs/meetings",
|
||||
"Meeting notes and transcripts",
|
||||
now
|
||||
);
|
||||
|
||||
// Add test documents
|
||||
const docs = [
|
||||
{
|
||||
name: "readme.md",
|
||||
title: "Project README",
|
||||
hash: "hash1",
|
||||
filepath: "/test/docs/readme.md",
|
||||
display_path: "readme.md",
|
||||
body: "# Project README\n\nThis is the main readme file for the project.\n\nIt contains important information about setup and usage.",
|
||||
},
|
||||
{
|
||||
name: "api.md",
|
||||
title: "API Documentation",
|
||||
hash: "hash2",
|
||||
filepath: "/test/docs/api.md",
|
||||
display_path: "api.md",
|
||||
body: "# API Documentation\n\nThis document describes the REST API endpoints.\n\n## Authentication\n\nUse Bearer tokens for auth.",
|
||||
},
|
||||
{
|
||||
name: "meeting-2024-01.md",
|
||||
title: "January Meeting Notes",
|
||||
hash: "hash3",
|
||||
filepath: "/test/docs/meetings/meeting-2024-01.md",
|
||||
display_path: "meetings/meeting-2024-01.md",
|
||||
body: "# January Meeting Notes\n\nDiscussed Q1 goals and roadmap.\n\n## Action Items\n\n- Review budget\n- Hire new team members",
|
||||
},
|
||||
{
|
||||
name: "meeting-2024-02.md",
|
||||
title: "February Meeting Notes",
|
||||
hash: "hash4",
|
||||
filepath: "/test/docs/meetings/meeting-2024-02.md",
|
||||
display_path: "meetings/meeting-2024-02.md",
|
||||
body: "# February Meeting Notes\n\nFollowed up on Q1 progress.\n\n## Updates\n\n- Budget approved\n- Two candidates interviewed",
|
||||
},
|
||||
{
|
||||
name: "large-file.md",
|
||||
title: "Large Document",
|
||||
hash: "hash5",
|
||||
filepath: "/test/docs/large-file.md",
|
||||
display_path: "large-file.md",
|
||||
body: "# Large Document\n\n" + "Lorem ipsum ".repeat(2000), // ~24KB
|
||||
},
|
||||
];
|
||||
|
||||
for (const doc of docs) {
|
||||
db.prepare(`
|
||||
INSERT INTO documents (collection_id, name, title, hash, filepath, display_path, body, created_at, modified_at, active)
|
||||
VALUES (1, ?, ?, ?, ?, ?, ?, ?, ?, 1)
|
||||
`).run(doc.name, doc.title, doc.hash, doc.filepath, doc.display_path, doc.body, now, now);
|
||||
}
|
||||
|
||||
// Add embeddings for vector search
|
||||
const embedding = new Float32Array(768);
|
||||
for (let i = 0; i < 768; i++) embedding[i] = Math.random();
|
||||
|
||||
for (const doc of docs.slice(0, 4)) { // Skip large file for embeddings
|
||||
db.prepare(`INSERT INTO content_vectors (hash, seq, pos, model, embedded_at) VALUES (?, 0, 0, 'embeddinggemma', ?)`).run(doc.hash, now);
|
||||
db.prepare(`INSERT INTO vectors_vec (hash_seq, embedding) VALUES (?, ?)`).run(`${doc.hash}_0`, embedding);
|
||||
}
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// MCP Server Test Helpers
|
||||
// =============================================================================
|
||||
|
||||
// We need to create a testable version of the MCP handlers
|
||||
// Since McpServer uses internal routing, we'll test the handler functions directly
|
||||
|
||||
import {
|
||||
searchFTS,
|
||||
searchVec,
|
||||
expandQuery,
|
||||
rerank,
|
||||
reciprocalRankFusion,
|
||||
extractSnippet,
|
||||
getContextForFile,
|
||||
getCollectionIdByName,
|
||||
getDocument,
|
||||
getMultipleDocuments,
|
||||
getStatus,
|
||||
DEFAULT_EMBED_MODEL,
|
||||
DEFAULT_QUERY_MODEL,
|
||||
DEFAULT_RERANK_MODEL,
|
||||
DEFAULT_MULTI_GET_MAX_BYTES,
|
||||
} from "./store";
|
||||
import type { RankedResult } from "./store";
|
||||
import { searchResultsToMcpCsv } from "./formatter";
|
||||
|
||||
// =============================================================================
|
||||
// Tests
|
||||
// =============================================================================
|
||||
|
||||
describe("MCP Server", () => {
|
||||
beforeAll(() => {
|
||||
globalThis.fetch = mockFetch as typeof fetch;
|
||||
setDefaultOllama(new Ollama({ baseUrl: OLLAMA_URL }));
|
||||
|
||||
testDbPath = `/tmp/qmd-mcp-test-${Date.now()}.sqlite`;
|
||||
testDb = new Database(testDbPath);
|
||||
initTestDatabase(testDb);
|
||||
seedTestData(testDb);
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
globalThis.fetch = originalFetch;
|
||||
setDefaultOllama(null);
|
||||
testDb.close();
|
||||
try {
|
||||
require("fs").unlinkSync(testDbPath);
|
||||
} catch {}
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// Tool: qmd_search (BM25)
|
||||
// ===========================================================================
|
||||
|
||||
describe("qmd_search tool", () => {
|
||||
test("returns results for matching query", () => {
|
||||
const results = searchFTS(testDb, "readme", 10);
|
||||
expect(results.length).toBeGreaterThan(0);
|
||||
expect(results[0].displayPath).toBe("readme.md");
|
||||
});
|
||||
|
||||
test("returns empty for non-matching query", () => {
|
||||
const results = searchFTS(testDb, "xyznonexistent", 10);
|
||||
expect(results.length).toBe(0);
|
||||
});
|
||||
|
||||
test("respects limit parameter", () => {
|
||||
const results = searchFTS(testDb, "meeting", 1);
|
||||
expect(results.length).toBe(1);
|
||||
});
|
||||
|
||||
test("filters by collection", () => {
|
||||
const collectionId = getCollectionIdByName(testDb, "docs");
|
||||
expect(collectionId).toBe(1);
|
||||
const results = searchFTS(testDb, "meeting", 10, collectionId!);
|
||||
expect(results.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
test("returns null for non-existent collection", () => {
|
||||
const collectionId = getCollectionIdByName(testDb, "nonexistent");
|
||||
expect(collectionId).toBeNull();
|
||||
});
|
||||
|
||||
test("formats results as CSV", () => {
|
||||
const results = searchFTS(testDb, "api", 10);
|
||||
const filtered = results.map(r => ({
|
||||
file: r.displayPath,
|
||||
title: r.title,
|
||||
score: Math.round(r.score * 100) / 100,
|
||||
context: getContextForFile(testDb, r.file),
|
||||
snippet: extractSnippet(r.body, "api", 300, r.chunkPos).snippet,
|
||||
}));
|
||||
const csv = searchResultsToMcpCsv(filtered);
|
||||
expect(csv).toContain("file,title,score,context,snippet");
|
||||
expect(csv).toContain("api.md");
|
||||
});
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// Tool: qmd_vsearch (Vector)
|
||||
// ===========================================================================
|
||||
|
||||
describe("qmd_vsearch tool", () => {
|
||||
test("returns results for semantic query", async () => {
|
||||
const results = await searchVec(testDb, "project documentation", DEFAULT_EMBED_MODEL, 10);
|
||||
expect(results.length).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
test("respects limit parameter", async () => {
|
||||
const results = await searchVec(testDb, "documentation", DEFAULT_EMBED_MODEL, 2);
|
||||
expect(results.length).toBeLessThanOrEqual(2);
|
||||
});
|
||||
|
||||
test("returns empty when no vector table exists", async () => {
|
||||
const emptyDb = new Database(":memory:");
|
||||
initTestDatabase(emptyDb);
|
||||
emptyDb.exec("DROP TABLE IF EXISTS vectors_vec");
|
||||
|
||||
const results = await searchVec(emptyDb, "test", DEFAULT_EMBED_MODEL, 10);
|
||||
expect(results.length).toBe(0);
|
||||
emptyDb.close();
|
||||
});
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// Tool: qmd_query (Hybrid)
|
||||
// ===========================================================================
|
||||
|
||||
describe("qmd_query tool", () => {
|
||||
test("expands query with variations", async () => {
|
||||
const queries = await expandQuery("api documentation", DEFAULT_QUERY_MODEL, testDb);
|
||||
expect(queries.length).toBeGreaterThan(1);
|
||||
expect(queries[0]).toBe("api documentation");
|
||||
});
|
||||
|
||||
test("performs RRF fusion on multiple result lists", () => {
|
||||
const list1: RankedResult[] = [
|
||||
{ file: "/a", displayPath: "a.md", title: "A", body: "body", score: 1 },
|
||||
{ file: "/b", displayPath: "b.md", title: "B", body: "body", score: 0.8 },
|
||||
];
|
||||
const list2: RankedResult[] = [
|
||||
{ file: "/b", displayPath: "b.md", title: "B", body: "body", score: 1 },
|
||||
{ file: "/c", displayPath: "c.md", title: "C", body: "body", score: 0.9 },
|
||||
];
|
||||
|
||||
const fused = reciprocalRankFusion([list1, list2]);
|
||||
expect(fused.length).toBe(3);
|
||||
// B appears in both lists, should have higher score
|
||||
const bResult = fused.find(r => r.file === "/b");
|
||||
expect(bResult).toBeDefined();
|
||||
});
|
||||
|
||||
test("reranks documents with LLM", async () => {
|
||||
const docs = [
|
||||
{ file: "/test/docs/readme.md", text: "Project readme" },
|
||||
{ file: "/test/docs/api.md", text: "API documentation" },
|
||||
];
|
||||
const reranked = await rerank("readme", docs, DEFAULT_RERANK_MODEL, testDb);
|
||||
expect(reranked.length).toBe(2);
|
||||
expect(reranked[0].score).toBeGreaterThan(0);
|
||||
});
|
||||
|
||||
test("full hybrid search pipeline", async () => {
|
||||
// Simulate full qmd_query flow
|
||||
const query = "meeting notes";
|
||||
const queries = await expandQuery(query, DEFAULT_QUERY_MODEL, testDb);
|
||||
|
||||
const rankedLists: RankedResult[][] = [];
|
||||
for (const q of queries) {
|
||||
const ftsResults = searchFTS(testDb, q, 20);
|
||||
if (ftsResults.length > 0) {
|
||||
rankedLists.push(ftsResults.map(r => ({
|
||||
file: r.file,
|
||||
displayPath: r.displayPath,
|
||||
title: r.title,
|
||||
body: r.body,
|
||||
score: r.score,
|
||||
})));
|
||||
}
|
||||
}
|
||||
|
||||
expect(rankedLists.length).toBeGreaterThan(0);
|
||||
|
||||
const fused = reciprocalRankFusion(rankedLists);
|
||||
expect(fused.length).toBeGreaterThan(0);
|
||||
|
||||
const candidates = fused.slice(0, 10);
|
||||
const reranked = await rerank(
|
||||
query,
|
||||
candidates.map(c => ({ file: c.file, text: c.body })),
|
||||
DEFAULT_RERANK_MODEL,
|
||||
testDb
|
||||
);
|
||||
|
||||
expect(reranked.length).toBeGreaterThan(0);
|
||||
});
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// Tool: qmd_get (Get Document)
|
||||
// ===========================================================================
|
||||
|
||||
describe("qmd_get tool", () => {
|
||||
test("retrieves document by display_path", () => {
|
||||
const result = getDocument(testDb, "readme.md");
|
||||
expect("error" in result).toBe(false);
|
||||
if (!("error" in result)) {
|
||||
expect(result.displayPath).toBe("readme.md");
|
||||
expect(result.body).toContain("Project README");
|
||||
}
|
||||
});
|
||||
|
||||
test("retrieves document by filepath", () => {
|
||||
const result = getDocument(testDb, "/test/docs/api.md");
|
||||
expect("error" in result).toBe(false);
|
||||
if (!("error" in result)) {
|
||||
expect(result.title).toBe("API Documentation");
|
||||
}
|
||||
});
|
||||
|
||||
test("retrieves document by partial path", () => {
|
||||
const result = getDocument(testDb, "api.md");
|
||||
expect("error" in result).toBe(false);
|
||||
});
|
||||
|
||||
test("returns not found for missing document", () => {
|
||||
const result = getDocument(testDb, "nonexistent.md");
|
||||
expect("error" in result).toBe(true);
|
||||
if ("error" in result) {
|
||||
expect(result.error).toBe("not_found");
|
||||
}
|
||||
});
|
||||
|
||||
test("suggests similar files when not found", () => {
|
||||
const result = getDocument(testDb, "readm.md"); // typo
|
||||
expect("error" in result).toBe(true);
|
||||
if ("error" in result) {
|
||||
expect(result.similarFiles.length).toBeGreaterThanOrEqual(0);
|
||||
}
|
||||
});
|
||||
|
||||
test("supports line range with :line suffix", () => {
|
||||
const result = getDocument(testDb, "readme.md:2", undefined, 2);
|
||||
expect("error" in result).toBe(false);
|
||||
if (!("error" in result)) {
|
||||
const lines = result.body.split("\n");
|
||||
expect(lines.length).toBeLessThanOrEqual(2);
|
||||
}
|
||||
});
|
||||
|
||||
test("supports fromLine parameter", () => {
|
||||
const result = getDocument(testDb, "readme.md", 3);
|
||||
expect("error" in result).toBe(false);
|
||||
if (!("error" in result)) {
|
||||
expect(result.body).not.toContain("# Project README");
|
||||
}
|
||||
});
|
||||
|
||||
test("supports maxLines parameter", () => {
|
||||
const result = getDocument(testDb, "api.md", 1, 3);
|
||||
expect("error" in result).toBe(false);
|
||||
if (!("error" in result)) {
|
||||
const lines = result.body.split("\n");
|
||||
expect(lines.length).toBeLessThanOrEqual(3);
|
||||
}
|
||||
});
|
||||
|
||||
test("includes context for documents in context path", () => {
|
||||
const result = getDocument(testDb, "meetings/meeting-2024-01.md");
|
||||
expect("error" in result).toBe(false);
|
||||
if (!("error" in result)) {
|
||||
expect(result.context).toBe("Meeting notes and transcripts");
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// Tool: qmd_multi_get (Multi Get)
|
||||
// ===========================================================================
|
||||
|
||||
describe("qmd_multi_get tool", () => {
|
||||
test("retrieves multiple documents by glob pattern", () => {
|
||||
const { files, errors } = getMultipleDocuments(testDb, "meetings/*.md");
|
||||
expect(errors.length).toBe(0);
|
||||
expect(files.length).toBe(2);
|
||||
expect(files.some(f => f.displayPath === "meetings/meeting-2024-01.md")).toBe(true);
|
||||
expect(files.some(f => f.displayPath === "meetings/meeting-2024-02.md")).toBe(true);
|
||||
});
|
||||
|
||||
test("retrieves documents by comma-separated list", () => {
|
||||
const { files, errors } = getMultipleDocuments(testDb, "readme.md, api.md");
|
||||
expect(errors.length).toBe(0);
|
||||
expect(files.length).toBe(2);
|
||||
});
|
||||
|
||||
test("returns errors for missing files in comma list", () => {
|
||||
const { files, errors } = getMultipleDocuments(testDb, "readme.md, nonexistent.md");
|
||||
expect(files.length).toBe(1);
|
||||
expect(errors.length).toBe(1);
|
||||
expect(errors[0]).toContain("not found");
|
||||
});
|
||||
|
||||
test("skips files larger than maxBytes", () => {
|
||||
const { files } = getMultipleDocuments(testDb, "*.md", undefined, 1000); // 1KB limit
|
||||
const largeFile = files.find(f => f.displayPath === "large-file.md");
|
||||
expect(largeFile).toBeDefined();
|
||||
expect(largeFile?.skipped).toBe(true);
|
||||
if (largeFile?.skipped) {
|
||||
expect(largeFile.skipReason).toContain("too large");
|
||||
}
|
||||
});
|
||||
|
||||
test("respects maxLines parameter", () => {
|
||||
const { files } = getMultipleDocuments(testDb, "readme.md", 2);
|
||||
expect(files.length).toBe(1);
|
||||
if (!files[0].skipped) {
|
||||
const lines = files[0].body.split("\n");
|
||||
// maxLines + truncation message
|
||||
expect(lines.length).toBeLessThanOrEqual(4);
|
||||
}
|
||||
});
|
||||
|
||||
test("returns error for non-matching glob", () => {
|
||||
const { files, errors } = getMultipleDocuments(testDb, "nonexistent/*.md");
|
||||
expect(files.length).toBe(0);
|
||||
expect(errors.length).toBe(1);
|
||||
expect(errors[0]).toContain("No files matched");
|
||||
});
|
||||
|
||||
test("includes context in results", () => {
|
||||
const { files } = getMultipleDocuments(testDb, "meetings/meeting-2024-01.md");
|
||||
expect(files.length).toBe(1);
|
||||
if (!files[0].skipped) {
|
||||
expect(files[0].context).toBe("Meeting notes and transcripts");
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// Tool: qmd_status
|
||||
// ===========================================================================
|
||||
|
||||
describe("qmd_status tool", () => {
|
||||
test("returns index status", () => {
|
||||
const status = getStatus(testDb);
|
||||
expect(status.totalDocuments).toBe(5);
|
||||
expect(status.hasVectorIndex).toBe(true);
|
||||
expect(status.collections.length).toBe(1);
|
||||
expect(status.collections[0].path).toBe("/test/docs");
|
||||
});
|
||||
|
||||
test("shows documents needing embedding", () => {
|
||||
const status = getStatus(testDb);
|
||||
// large-file.md doesn't have embeddings
|
||||
expect(status.needsEmbedding).toBe(1);
|
||||
});
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// Resource: qmd://{path}
|
||||
// ===========================================================================
|
||||
|
||||
describe("qmd:// resource", () => {
|
||||
test("lists all documents", () => {
|
||||
const docs = testDb.prepare(`
|
||||
SELECT display_path, title
|
||||
FROM documents
|
||||
WHERE active = 1
|
||||
ORDER BY modified_at DESC
|
||||
LIMIT 1000
|
||||
`).all() as { display_path: string; title: string }[];
|
||||
|
||||
expect(docs.length).toBe(5);
|
||||
expect(docs.map(d => d.display_path)).toContain("readme.md");
|
||||
});
|
||||
|
||||
test("reads document by display_path", () => {
|
||||
const path = "readme.md";
|
||||
const doc = testDb.prepare(`
|
||||
SELECT filepath, display_path, body
|
||||
FROM documents
|
||||
WHERE display_path = ? AND active = 1
|
||||
`).get(path) as { filepath: string; display_path: string; body: string } | null;
|
||||
|
||||
expect(doc).not.toBeNull();
|
||||
expect(doc?.body).toContain("Project README");
|
||||
});
|
||||
|
||||
test("reads document by URL-encoded path", () => {
|
||||
// Simulate URL encoding that MCP clients may send
|
||||
const encodedPath = "meetings%2Fmeeting-2024-01.md";
|
||||
const decodedPath = decodeURIComponent(encodedPath);
|
||||
|
||||
const doc = testDb.prepare(`
|
||||
SELECT filepath, display_path, body
|
||||
FROM documents
|
||||
WHERE display_path = ? AND active = 1
|
||||
`).get(decodedPath) as { filepath: string; display_path: string; body: string } | null;
|
||||
|
||||
expect(doc).not.toBeNull();
|
||||
expect(doc?.display_path).toBe("meetings/meeting-2024-01.md");
|
||||
});
|
||||
|
||||
test("reads document by suffix match", () => {
|
||||
const path = "meeting-2024-01.md"; // without meetings/ prefix
|
||||
let doc = testDb.prepare(`
|
||||
SELECT filepath, display_path, body
|
||||
FROM documents
|
||||
WHERE display_path = ? AND active = 1
|
||||
`).get(path) as { filepath: string; display_path: string; body: string } | null;
|
||||
|
||||
if (!doc) {
|
||||
doc = testDb.prepare(`
|
||||
SELECT filepath, display_path, body
|
||||
FROM documents
|
||||
WHERE display_path LIKE ? AND active = 1
|
||||
LIMIT 1
|
||||
`).get(`%${path}`) as { filepath: string; display_path: string; body: string } | null;
|
||||
}
|
||||
|
||||
expect(doc).not.toBeNull();
|
||||
expect(doc?.display_path).toBe("meetings/meeting-2024-01.md");
|
||||
});
|
||||
|
||||
test("returns not found for missing document", () => {
|
||||
const path = "nonexistent.md";
|
||||
const doc = testDb.prepare(`
|
||||
SELECT filepath, display_path, body
|
||||
FROM documents
|
||||
WHERE display_path = ? AND active = 1
|
||||
`).get(path) as { filepath: string; display_path: string; body: string } | null;
|
||||
|
||||
expect(doc).toBeNull();
|
||||
});
|
||||
|
||||
test("includes context in document body", () => {
|
||||
const path = "meetings/meeting-2024-01.md";
|
||||
const doc = testDb.prepare(`
|
||||
SELECT filepath, display_path, body
|
||||
FROM documents
|
||||
WHERE display_path = ? AND active = 1
|
||||
`).get(path) as { filepath: string; display_path: string; body: string } | null;
|
||||
|
||||
expect(doc).not.toBeNull();
|
||||
const context = getContextForFile(testDb, doc!.filepath);
|
||||
expect(context).toBe("Meeting notes and transcripts");
|
||||
|
||||
// Verify context would be prepended
|
||||
let text = doc!.body;
|
||||
if (context) {
|
||||
text = `<!-- Context: ${context} -->\n\n` + text;
|
||||
}
|
||||
expect(text).toContain("<!-- Context: Meeting notes and transcripts -->");
|
||||
});
|
||||
|
||||
test("handles URL-encoded special characters", () => {
|
||||
// Test various URL encodings
|
||||
const testCases = [
|
||||
{ encoded: "readme.md", decoded: "readme.md" },
|
||||
{ encoded: "meetings%2Fmeeting-2024-01.md", decoded: "meetings/meeting-2024-01.md" },
|
||||
{ encoded: "api.md%3A10", decoded: "api.md:10" }, // with line number
|
||||
];
|
||||
|
||||
for (const { encoded, decoded } of testCases) {
|
||||
expect(decodeURIComponent(encoded)).toBe(decoded);
|
||||
}
|
||||
});
|
||||
|
||||
test("handles double-encoded URLs", () => {
|
||||
// Some clients may double-encode
|
||||
const doubleEncoded = "meetings%252Fmeeting-2024-01.md";
|
||||
const singleDecoded = decodeURIComponent(doubleEncoded);
|
||||
expect(singleDecoded).toBe("meetings%2Fmeeting-2024-01.md");
|
||||
|
||||
const fullyDecoded = decodeURIComponent(singleDecoded);
|
||||
expect(fullyDecoded).toBe("meetings/meeting-2024-01.md");
|
||||
});
|
||||
|
||||
test("handles URL-encoded paths with spaces", () => {
|
||||
// Add a document with spaces in the path
|
||||
const now = new Date().toISOString();
|
||||
testDb.prepare(`
|
||||
INSERT INTO documents (collection_id, name, title, hash, filepath, display_path, body, created_at, modified_at, active)
|
||||
VALUES (1, ?, ?, ?, ?, ?, ?, ?, ?, 1)
|
||||
`).run(
|
||||
"podcast with spaces.md",
|
||||
"Podcast Episode",
|
||||
"hash_spaces",
|
||||
"/test/docs/External Podcast/2023 April - Interview.md",
|
||||
"External Podcast/2023 April - Interview.md",
|
||||
"# Podcast Episode\n\nInterview content here.",
|
||||
now,
|
||||
now
|
||||
);
|
||||
|
||||
// Simulate URL-encoded path from MCP client
|
||||
const encodedPath = "External%20Podcast%2F2023%20April%20-%20Interview.md";
|
||||
const decodedPath = decodeURIComponent(encodedPath);
|
||||
|
||||
expect(decodedPath).toBe("External Podcast/2023 April - Interview.md");
|
||||
|
||||
const doc = testDb.prepare(`
|
||||
SELECT filepath, display_path, body
|
||||
FROM documents
|
||||
WHERE display_path = ? AND active = 1
|
||||
`).get(decodedPath) as { filepath: string; display_path: string; body: string } | null;
|
||||
|
||||
expect(doc).not.toBeNull();
|
||||
expect(doc?.display_path).toBe("External Podcast/2023 April - Interview.md");
|
||||
expect(doc?.body).toContain("Podcast Episode");
|
||||
});
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// Prompt: query
|
||||
// ===========================================================================
|
||||
|
||||
describe("query prompt", () => {
|
||||
test("returns usage guide", () => {
|
||||
// The prompt content is static, just verify the structure
|
||||
const promptContent = `# QMD - Quick Markdown Search
|
||||
|
||||
QMD is your on-device search engine for markdown knowledge bases.`;
|
||||
|
||||
expect(promptContent).toContain("QMD");
|
||||
expect(promptContent).toContain("search");
|
||||
});
|
||||
|
||||
test("describes all available tools", () => {
|
||||
const toolNames = [
|
||||
"qmd_search",
|
||||
"qmd_vsearch",
|
||||
"qmd_query",
|
||||
"qmd_get",
|
||||
"qmd_multi_get",
|
||||
"qmd_status",
|
||||
];
|
||||
|
||||
// Verify these are documented in the prompt
|
||||
const promptGuide = `
|
||||
### 1. qmd_search (Fast keyword search)
|
||||
### 2. qmd_vsearch (Semantic search)
|
||||
### 3. qmd_query (Hybrid search - highest quality)
|
||||
### 4. qmd_get (Retrieve document)
|
||||
### 5. qmd_multi_get (Retrieve multiple documents)
|
||||
### 6. qmd_status (Index info)
|
||||
`;
|
||||
|
||||
for (const tool of toolNames) {
|
||||
expect(promptGuide).toContain(tool);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// Edge Cases
|
||||
// ===========================================================================
|
||||
|
||||
describe("edge cases", () => {
|
||||
test("handles empty query", () => {
|
||||
const results = searchFTS(testDb, "", 10);
|
||||
expect(results.length).toBe(0);
|
||||
});
|
||||
|
||||
test("handles special characters in query", () => {
|
||||
const results = searchFTS(testDb, "project's", 10);
|
||||
// Should not throw
|
||||
expect(Array.isArray(results)).toBe(true);
|
||||
});
|
||||
|
||||
test("handles unicode in query", () => {
|
||||
const results = searchFTS(testDb, "文档", 10);
|
||||
expect(Array.isArray(results)).toBe(true);
|
||||
});
|
||||
|
||||
test("handles very long query", () => {
|
||||
const longQuery = "documentation ".repeat(100);
|
||||
const results = searchFTS(testDb, longQuery, 10);
|
||||
expect(Array.isArray(results)).toBe(true);
|
||||
});
|
||||
|
||||
test("handles query with only stopwords", () => {
|
||||
const results = searchFTS(testDb, "the and or", 10);
|
||||
expect(Array.isArray(results)).toBe(true);
|
||||
});
|
||||
|
||||
test("extracts snippet around matching text", () => {
|
||||
const body = "Line 1\nLine 2\nThis is the important line with the keyword\nLine 4\nLine 5";
|
||||
const { line, snippet } = extractSnippet(body, "keyword", 200);
|
||||
expect(snippet).toContain("keyword");
|
||||
expect(line).toBe(3);
|
||||
});
|
||||
|
||||
test("handles snippet extraction with chunkPos", () => {
|
||||
const body = "A".repeat(1000) + "KEYWORD" + "B".repeat(1000);
|
||||
const chunkPos = 1000; // Position of KEYWORD
|
||||
const { snippet } = extractSnippet(body, "keyword", 200, chunkPos);
|
||||
expect(snippet).toContain("KEYWORD");
|
||||
});
|
||||
});
|
||||
|
||||
// ===========================================================================
|
||||
// CSV Formatting
|
||||
// ===========================================================================
|
||||
|
||||
describe("CSV formatting", () => {
|
||||
test("escapes quotes in CSV", () => {
|
||||
const results = [{
|
||||
file: 'test.md',
|
||||
title: 'Test "quoted" title',
|
||||
score: 0.9,
|
||||
context: null,
|
||||
snippet: 'Some "quoted" text',
|
||||
}];
|
||||
const csv = searchResultsToMcpCsv(results);
|
||||
expect(csv).toContain('""quoted""');
|
||||
});
|
||||
|
||||
test("escapes newlines in CSV", () => {
|
||||
const results = [{
|
||||
file: 'test.md',
|
||||
title: 'Test title',
|
||||
score: 0.9,
|
||||
context: null,
|
||||
snippet: 'Line 1\nLine 2',
|
||||
}];
|
||||
const csv = searchResultsToMcpCsv(results);
|
||||
expect(csv).not.toContain('\n\n'); // Should be escaped within quotes
|
||||
});
|
||||
|
||||
test("handles empty results", () => {
|
||||
const csv = searchResultsToMcpCsv([]);
|
||||
expect(csv).toBe("file,title,score,context,snippet");
|
||||
});
|
||||
});
|
||||
});
|
||||
503
mcp.ts
Normal file
503
mcp.ts
Normal file
@ -0,0 +1,503 @@
|
||||
#!/usr/bin/env bun
|
||||
/**
|
||||
* QMD MCP Server - Model Context Protocol server for QMD
|
||||
*
|
||||
* Exposes QMD search and document retrieval as MCP tools and resources.
|
||||
* Documents are accessible via qmd:// URIs.
|
||||
*/
|
||||
|
||||
import { McpServer, ResourceTemplate } from "@modelcontextprotocol/sdk/server/mcp.js";
|
||||
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
|
||||
import { z } from "zod";
|
||||
import {
|
||||
createStore,
|
||||
reciprocalRankFusion,
|
||||
extractSnippet,
|
||||
DEFAULT_EMBED_MODEL,
|
||||
DEFAULT_QUERY_MODEL,
|
||||
DEFAULT_RERANK_MODEL,
|
||||
DEFAULT_MULTI_GET_MAX_BYTES,
|
||||
} from "./store.js";
|
||||
import type { RankedResult } from "./store.js";
|
||||
import { searchResultsToMcpCsv } from "./formatter.js";
|
||||
|
||||
export async function startMcpServer(): Promise<void> {
|
||||
// Open database once at startup - keep it open for the lifetime of the server
|
||||
const store = createStore();
|
||||
|
||||
const server = new McpServer({
|
||||
name: "qmd",
|
||||
version: "1.0.0",
|
||||
});
|
||||
|
||||
// Register resource template for qmd:// URIs
|
||||
// This allows clients to list and read documents via the MCP resources API
|
||||
server.registerResource(
|
||||
"document",
|
||||
new ResourceTemplate("qmd://{path}", {
|
||||
list: async () => {
|
||||
// List all indexed documents
|
||||
const docs = store.db.prepare(`
|
||||
SELECT display_path, title
|
||||
FROM documents
|
||||
WHERE active = 1
|
||||
ORDER BY modified_at DESC
|
||||
LIMIT 1000
|
||||
`).all() as { display_path: string; title: string }[];
|
||||
|
||||
return {
|
||||
resources: docs.map(doc => ({
|
||||
uri: `qmd://${encodeURIComponent(doc.display_path)}`,
|
||||
name: doc.title || doc.display_path,
|
||||
mimeType: "text/markdown",
|
||||
})),
|
||||
};
|
||||
},
|
||||
}),
|
||||
{
|
||||
title: "QMD Document",
|
||||
description: "A markdown document from your QMD knowledge base",
|
||||
mimeType: "text/markdown",
|
||||
},
|
||||
async (uri, { path }) => {
|
||||
// Decode URL-encoded path (MCP clients send encoded URIs)
|
||||
const decodedPath = decodeURIComponent(path);
|
||||
|
||||
// Find document by display_path
|
||||
let doc = store.db.prepare(`SELECT filepath, display_path, body FROM documents WHERE display_path = ? AND active = 1`).get(decodedPath) as { filepath: string; display_path: string; body: string } | null;
|
||||
|
||||
// Try suffix match if exact match fails
|
||||
if (!doc) {
|
||||
doc = store.db.prepare(`SELECT filepath, display_path, body FROM documents WHERE display_path LIKE ? AND active = 1 LIMIT 1`).get(`%${decodedPath}`) as { filepath: string; display_path: string; body: string } | null;
|
||||
}
|
||||
|
||||
if (!doc) {
|
||||
return { contents: [{ uri: uri.href, text: `Document not found: ${decodedPath}` }] };
|
||||
}
|
||||
|
||||
const context = store.getContextForFile(doc.filepath);
|
||||
|
||||
let text = doc.body;
|
||||
if (context) {
|
||||
text = `<!-- Context: ${context} -->\n\n` + text;
|
||||
}
|
||||
|
||||
return {
|
||||
contents: [{
|
||||
uri: uri.href,
|
||||
mimeType: "text/markdown",
|
||||
text,
|
||||
}],
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
// Register the query prompt - describes ideal usage
|
||||
server.registerPrompt(
|
||||
"query",
|
||||
{
|
||||
title: "QMD Query Guide",
|
||||
description: "How to effectively search your knowledge base with QMD",
|
||||
},
|
||||
() => ({
|
||||
messages: [
|
||||
{
|
||||
role: "user",
|
||||
content: {
|
||||
type: "text",
|
||||
text: `# QMD - Quick Markdown Search
|
||||
|
||||
QMD is your on-device search engine for markdown knowledge bases. Use it to find information across your notes, documents, and meeting transcripts.
|
||||
|
||||
## Available Tools
|
||||
|
||||
### 1. qmd_search (Fast keyword search)
|
||||
Best for: Finding documents with specific keywords or phrases.
|
||||
- Uses BM25 full-text search
|
||||
- Fast, no LLM required
|
||||
- Good for exact matches
|
||||
- Use \`collection\` parameter to filter to a specific collection
|
||||
|
||||
### 2. qmd_vsearch (Semantic search)
|
||||
Best for: Finding conceptually related content even without exact keyword matches.
|
||||
- Uses vector embeddings
|
||||
- Understands meaning and context
|
||||
- Good for "how do I..." or conceptual queries
|
||||
- Use \`collection\` parameter to filter to a specific collection
|
||||
|
||||
### 3. qmd_query (Hybrid search - highest quality)
|
||||
Best for: Important searches where you want the best results.
|
||||
- Combines keyword + semantic search
|
||||
- Expands your query with variations
|
||||
- Re-ranks results with LLM
|
||||
- Slower but most accurate
|
||||
- Use \`collection\` parameter to filter to a specific collection
|
||||
|
||||
### 4. qmd_get (Retrieve document)
|
||||
Best for: Getting the full content of a single document you found.
|
||||
- Use the file path from search results
|
||||
- Supports line ranges: \`file.md:100\` or fromLine/maxLines parameters
|
||||
- Suggests similar files if not found
|
||||
|
||||
### 5. qmd_multi_get (Retrieve multiple documents)
|
||||
Best for: Getting content from multiple files at once.
|
||||
- Use glob patterns: \`journals/2025-05*.md\`
|
||||
- Or comma-separated: \`file1.md, file2.md\`
|
||||
- Skips files over maxBytes (default 10KB) - use qmd_get for large files
|
||||
|
||||
### 6. qmd_status (Index info)
|
||||
Shows collection info, document counts, and embedding status.
|
||||
|
||||
## Resources
|
||||
|
||||
You can also access documents directly via the \`qmd://\` URI scheme:
|
||||
- List all documents: \`resources/list\`
|
||||
- Read a document: \`resources/read\` with uri \`qmd://path/to/file.md\`
|
||||
|
||||
## Search Strategy
|
||||
|
||||
1. **Start with qmd_search** for quick keyword lookups
|
||||
2. **Use qmd_vsearch** when keywords aren't working or for conceptual queries
|
||||
3. **Use qmd_query** for important searches or when you need high confidence
|
||||
4. **Use qmd_get** to retrieve a single full document
|
||||
5. **Use qmd_multi_get** to batch retrieve multiple related files
|
||||
|
||||
## Tips
|
||||
|
||||
- Use \`minScore: 0.5\` to filter low-relevance results
|
||||
- Use \`collection: "notes"\` to search only in a specific collection
|
||||
- Check the "Context" field - it describes what kind of content the file contains
|
||||
- File paths are relative to their collection (e.g., \`pages/meeting.md\`)
|
||||
- For glob patterns, match on display_path (e.g., \`journals/2025-*.md\`)`,
|
||||
},
|
||||
},
|
||||
],
|
||||
})
|
||||
);
|
||||
|
||||
// Tool: search (BM25 full-text)
|
||||
server.registerTool(
|
||||
"qmd_search",
|
||||
{
|
||||
title: "Search (BM25)",
|
||||
description: "Fast keyword-based full-text search using BM25. Best for finding documents with specific words or phrases.",
|
||||
inputSchema: {
|
||||
query: z.string().describe("Search query - keywords or phrases to find"),
|
||||
limit: z.number().optional().default(10).describe("Maximum number of results (default: 10)"),
|
||||
minScore: z.number().optional().default(0).describe("Minimum relevance score 0-1 (default: 0)"),
|
||||
collection: z.string().optional().describe("Filter to a specific collection by name"),
|
||||
},
|
||||
},
|
||||
async ({ query, limit, minScore, collection }) => {
|
||||
// Resolve collection filter
|
||||
let collectionId: number | undefined;
|
||||
if (collection) {
|
||||
collectionId = store.getCollectionIdByName(collection) ?? undefined;
|
||||
if (collectionId === undefined) {
|
||||
return { content: [{ type: "text", text: `Error: Collection not found: ${collection}` }] };
|
||||
}
|
||||
}
|
||||
|
||||
const results = store.searchFTS(query, limit || 10, collectionId);
|
||||
const filtered = results
|
||||
.filter(r => r.score >= (minScore || 0))
|
||||
.map(r => ({
|
||||
file: r.displayPath,
|
||||
title: r.title,
|
||||
score: Math.round(r.score * 100) / 100,
|
||||
context: store.getContextForFile(r.file),
|
||||
snippet: extractSnippet(r.body, query, 300, r.chunkPos).snippet,
|
||||
}));
|
||||
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
mimeType: "text/csv",
|
||||
text: searchResultsToMcpCsv(filtered),
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
// Tool: vsearch (Vector semantic search)
|
||||
server.registerTool(
|
||||
"qmd_vsearch",
|
||||
{
|
||||
title: "Vector Search (Semantic)",
|
||||
description: "Semantic similarity search using vector embeddings. Finds conceptually related content even without exact keyword matches. Requires embeddings (run 'qmd embed' first).",
|
||||
inputSchema: {
|
||||
query: z.string().describe("Natural language query - describe what you're looking for"),
|
||||
limit: z.number().optional().default(10).describe("Maximum number of results (default: 10)"),
|
||||
minScore: z.number().optional().default(0.3).describe("Minimum relevance score 0-1 (default: 0.3)"),
|
||||
collection: z.string().optional().describe("Filter to a specific collection by name"),
|
||||
},
|
||||
},
|
||||
async ({ query, limit, minScore, collection }) => {
|
||||
// Resolve collection filter
|
||||
let collectionId: number | undefined;
|
||||
if (collection) {
|
||||
collectionId = store.getCollectionIdByName(collection) ?? undefined;
|
||||
if (collectionId === undefined) {
|
||||
return { content: [{ type: "text", text: `Error: Collection not found: ${collection}` }] };
|
||||
}
|
||||
}
|
||||
|
||||
const tableExists = store.db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get();
|
||||
if (!tableExists) {
|
||||
return {
|
||||
content: [{ type: "text", text: "Error: Vector index not found. Run 'qmd embed' first to create embeddings." }],
|
||||
};
|
||||
}
|
||||
|
||||
// Expand query
|
||||
const queries = await store.expandQuery(query, DEFAULT_QUERY_MODEL);
|
||||
|
||||
// Collect results
|
||||
const allResults = new Map<string, { file: string; displayPath: string; title: string; body: string; score: number }>();
|
||||
for (const q of queries) {
|
||||
const vecResults = await store.searchVec(q, DEFAULT_EMBED_MODEL, limit || 10, collectionId);
|
||||
for (const r of vecResults) {
|
||||
const existing = allResults.get(r.file);
|
||||
if (!existing || r.score > existing.score) {
|
||||
allResults.set(r.file, { file: r.file, displayPath: r.displayPath, title: r.title, body: r.body, score: r.score });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const filtered = Array.from(allResults.values())
|
||||
.sort((a, b) => b.score - a.score)
|
||||
.slice(0, limit || 10)
|
||||
.filter(r => r.score >= (minScore || 0.3))
|
||||
.map(r => ({
|
||||
file: r.displayPath,
|
||||
title: r.title,
|
||||
score: Math.round(r.score * 100) / 100,
|
||||
context: store.getContextForFile(r.file),
|
||||
snippet: extractSnippet(r.body, query, 300).snippet,
|
||||
}));
|
||||
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
mimeType: "text/csv",
|
||||
text: searchResultsToMcpCsv(filtered),
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
// Tool: query (Hybrid with reranking)
|
||||
server.registerTool(
|
||||
"qmd_query",
|
||||
{
|
||||
title: "Hybrid Query (Best Quality)",
|
||||
description: "Highest quality search combining BM25 + vector + query expansion + LLM reranking. Slower but most accurate. Use for important searches.",
|
||||
inputSchema: {
|
||||
query: z.string().describe("Natural language query - describe what you're looking for"),
|
||||
limit: z.number().optional().default(10).describe("Maximum number of results (default: 10)"),
|
||||
minScore: z.number().optional().default(0).describe("Minimum relevance score 0-1 (default: 0)"),
|
||||
collection: z.string().optional().describe("Filter to a specific collection by name"),
|
||||
},
|
||||
},
|
||||
async ({ query, limit, minScore, collection }) => {
|
||||
// Resolve collection filter
|
||||
let collectionId: number | undefined;
|
||||
if (collection) {
|
||||
collectionId = store.getCollectionIdByName(collection) ?? undefined;
|
||||
if (collectionId === undefined) {
|
||||
return { content: [{ type: "text", text: `Error: Collection not found: ${collection}` }] };
|
||||
}
|
||||
}
|
||||
|
||||
// Expand query
|
||||
const queries = await store.expandQuery(query, DEFAULT_QUERY_MODEL);
|
||||
|
||||
// Collect ranked lists
|
||||
const rankedLists: RankedResult[][] = [];
|
||||
const hasVectors = !!store.db.prepare(`SELECT name FROM sqlite_master WHERE type='table' AND name='vectors_vec'`).get();
|
||||
|
||||
for (const q of queries) {
|
||||
const ftsResults = store.searchFTS(q, 20, collectionId);
|
||||
if (ftsResults.length > 0) {
|
||||
rankedLists.push(ftsResults.map(r => ({ file: r.file, displayPath: r.displayPath, title: r.title, body: r.body, score: r.score })));
|
||||
}
|
||||
if (hasVectors) {
|
||||
const vecResults = await store.searchVec(q, DEFAULT_EMBED_MODEL, 20, collectionId);
|
||||
if (vecResults.length > 0) {
|
||||
rankedLists.push(vecResults.map(r => ({ file: r.file, displayPath: r.displayPath, title: r.title, body: r.body, score: r.score })));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// RRF fusion
|
||||
const weights = rankedLists.map((_, i) => i < 2 ? 2.0 : 1.0);
|
||||
const fused = reciprocalRankFusion(rankedLists, weights);
|
||||
const candidates = fused.slice(0, 30);
|
||||
|
||||
// Rerank
|
||||
const reranked = await store.rerank(
|
||||
query,
|
||||
candidates.map(c => ({ file: c.file, text: c.body })),
|
||||
DEFAULT_RERANK_MODEL
|
||||
);
|
||||
|
||||
// Blend scores
|
||||
const candidateMap = new Map(candidates.map(c => [c.file, { displayPath: c.displayPath, title: c.title, body: c.body }]));
|
||||
const rrfRankMap = new Map(candidates.map((c, i) => [c.file, i + 1]));
|
||||
|
||||
const finalResults = reranked.map(r => {
|
||||
const rrfRank = rrfRankMap.get(r.file) || candidates.length;
|
||||
let rrfWeight: number;
|
||||
if (rrfRank <= 3) rrfWeight = 0.75;
|
||||
else if (rrfRank <= 10) rrfWeight = 0.60;
|
||||
else rrfWeight = 0.40;
|
||||
const rrfScore = 1 / rrfRank;
|
||||
const blendedScore = rrfWeight * rrfScore + (1 - rrfWeight) * r.score;
|
||||
const candidate = candidateMap.get(r.file);
|
||||
return {
|
||||
file: candidate?.displayPath || "",
|
||||
title: candidate?.title || "",
|
||||
score: Math.round(blendedScore * 100) / 100,
|
||||
context: store.getContextForFile(r.file),
|
||||
snippet: extractSnippet(candidate?.body || "", query, 300).snippet,
|
||||
};
|
||||
}).filter(r => r.score >= (minScore || 0)).slice(0, limit || 10);
|
||||
|
||||
return {
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
mimeType: "text/csv",
|
||||
text: searchResultsToMcpCsv(finalResults),
|
||||
},
|
||||
],
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
// Tool: get (Retrieve document)
|
||||
server.registerTool(
|
||||
"qmd_get",
|
||||
{
|
||||
title: "Get Document",
|
||||
description: "Retrieve the full content of a document by its file path. Use paths from search results. Suggests similar files if not found.",
|
||||
inputSchema: {
|
||||
file: z.string().describe("File path from search results (e.g., 'pages/meeting.md' or 'pages/meeting.md:100' to start at line 100)"),
|
||||
fromLine: z.number().optional().describe("Start from this line number (1-indexed)"),
|
||||
maxLines: z.number().optional().describe("Maximum number of lines to return"),
|
||||
},
|
||||
},
|
||||
async ({ file, fromLine, maxLines }) => {
|
||||
const result = store.getDocument(file, fromLine, maxLines);
|
||||
|
||||
if ("error" in result) {
|
||||
let msg = `Error: Document not found: ${file}`;
|
||||
if (result.similarFiles.length > 0) {
|
||||
msg += `\n\nDid you mean one of these?\n${result.similarFiles.map(s => ` - ${s}`).join('\n')}`;
|
||||
}
|
||||
return { content: [{ type: "text", text: msg }] };
|
||||
}
|
||||
|
||||
let text = result.body;
|
||||
if (result.context) {
|
||||
text = `<!-- Context: ${result.context} -->\n\n` + text;
|
||||
}
|
||||
|
||||
return {
|
||||
content: [{
|
||||
type: "resource",
|
||||
resource: {
|
||||
uri: `qmd://${result.displayPath}`,
|
||||
mimeType: "text/markdown",
|
||||
text,
|
||||
},
|
||||
}],
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
// Tool: multi-get (Retrieve multiple documents)
|
||||
server.registerTool(
|
||||
"qmd_multi_get",
|
||||
{
|
||||
title: "Multi-Get Documents",
|
||||
description: "Retrieve multiple documents by glob pattern (e.g., 'journals/2025-05*.md') or comma-separated list. Skips files larger than maxBytes.",
|
||||
inputSchema: {
|
||||
pattern: z.string().describe("Glob pattern or comma-separated list of file paths"),
|
||||
maxLines: z.number().optional().describe("Maximum lines per file"),
|
||||
maxBytes: z.number().optional().default(10240).describe("Skip files larger than this (default: 10240 = 10KB)"),
|
||||
},
|
||||
},
|
||||
async ({ pattern, maxLines, maxBytes }) => {
|
||||
const { files, errors } = store.getMultipleDocuments(pattern, maxLines, maxBytes || DEFAULT_MULTI_GET_MAX_BYTES);
|
||||
|
||||
if (files.length === 0 && errors.length === 0) {
|
||||
return { content: [{ type: "text", text: `No files matched pattern: ${pattern}` }] };
|
||||
}
|
||||
|
||||
const content: ({ type: "text"; text: string } | { type: "resource"; resource: { uri: string; mimeType: string; text: string } })[] = [];
|
||||
|
||||
if (errors.length > 0) {
|
||||
content.push({ type: "text", text: `Errors:\n${errors.join('\n')}` });
|
||||
}
|
||||
|
||||
for (const file of files) {
|
||||
if (file.skipped) {
|
||||
content.push({
|
||||
type: "text",
|
||||
text: `[SKIPPED: ${file.displayPath} - ${file.skipReason}. Use 'qmd_get' with file="${file.displayPath}" to retrieve.]`,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
|
||||
let text = file.body;
|
||||
if (file.context) {
|
||||
text = `<!-- Context: ${file.context} -->\n\n` + text;
|
||||
}
|
||||
|
||||
content.push({
|
||||
type: "resource",
|
||||
resource: {
|
||||
uri: `qmd://${file.displayPath}`,
|
||||
mimeType: "text/markdown",
|
||||
text,
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
return { content };
|
||||
}
|
||||
);
|
||||
|
||||
// Tool: status (Index status)
|
||||
server.registerTool(
|
||||
"qmd_status",
|
||||
{
|
||||
title: "Index Status",
|
||||
description: "Show the status of the QMD index: collections, document counts, and health information.",
|
||||
inputSchema: {},
|
||||
},
|
||||
async () => {
|
||||
const status = store.getStatus();
|
||||
|
||||
return {
|
||||
content: [{ type: "text", text: JSON.stringify(status, null, 2) }],
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
// Connect via stdio
|
||||
const transport = new StdioServerTransport();
|
||||
await server.connect(transport);
|
||||
|
||||
// Note: Database stays open - it will be closed when the process exits
|
||||
}
|
||||
|
||||
// Run if this is the main module
|
||||
if (import.meta.main) {
|
||||
startMcpServer().catch(console.error);
|
||||
}
|
||||
@ -7,13 +7,15 @@
|
||||
"qmd": "./qmd"
|
||||
},
|
||||
"scripts": {
|
||||
"test": "bun test",
|
||||
"qmd": "bun qmd.ts",
|
||||
"index": "bun qmd.ts index",
|
||||
"vector": "bun qmd.ts vector",
|
||||
"search": "bun qmd.ts search",
|
||||
"vsearch": "bun qmd.ts vsearch",
|
||||
"rerank": "bun qmd.ts rerank",
|
||||
"link": "bun link"
|
||||
"link": "bun link",
|
||||
"inspector": "npx @modelcontextprotocol/inspector bun qmd.ts mcp"
|
||||
},
|
||||
"dependencies": {
|
||||
"@modelcontextprotocol/sdk": "^1.24.3",
|
||||
|
||||
1808
store.test.ts
Normal file
1808
store.test.ts
Normal file
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue
Block a user