@aipmanager/search-mcp/dist/test-codebase-search.js

MCP server providing Cursor-like Search tools: read file, list dir, grep, search files, semantic code search, web search

#!/usr/bin/env ts-node
import 'dotenv/config';
import OpenAI from 'openai';
import { MilvusClient, DataType, MetricType } from '@zilliz/milvus2-sdk-node';
import { Pinecone } from '@pinecone-database/pinecone';
function parseArgs() {
    const argv = process.argv.slice(2);
    let query = '';
    let topK = 8;
    for (let i = 0; i < argv.length; i++) {
        const a = argv[i];
        if (a === '--query' || a === '-q')
            query = argv[++i] || '';
        else if (a.startsWith('--query='))
            query = a.split('=')[1];
        else if (a === '--topK' || a === '-k')
            topK = Number(argv[++i] || 8);
        else if (a.startsWith('--topK='))
            topK = Number(a.split('=')[1]);
    }
    if (!query) {
        console.error('Usage: ts-node src/test-codebase-search.ts --query "关键词" [--topK 8]');
        process.exit(1);
    }
    return { query, topK };
}
async function embedTexts(texts) {
    const provider = (process.env.SEARCH_EMBED_PROVIDER || process.env.INDEXER_EMBED_PROVIDER || (process.env.SILICONFLOW_API_KEY ? 'siliconflow' : 'openai')).toLowerCase();
    const model = process.env.SEARCH_EMBED_MODEL || process.env.INDEXER_EMBED_MODEL || 'text-embedding-3-small';
    let apiKey;
    let baseURL = process.env.SEARCH_EMBED_BASE_URL || process.env.INDEXER_EMBED_BASE_URL;
    if (provider === 'siliconflow') {
        apiKey = process.env.SILICONFLOW_API_KEY;
        baseURL = baseURL || 'https://api.siliconflow.cn/v1';
    }
    else if (provider === 'custom') {
        apiKey = process.env.OPENAI_API_KEY || process.env.SILICONFLOW_API_KEY;
        if (!baseURL)
            throw new Error('SEARCH_EMBED_BASE_URL is required for custom provider');
    }
    else {
        apiKey = process.env.OPENAI_API_KEY;
    }
    if (!apiKey)
        throw new Error('Embedding provider API key is missing');
    const client = new OpenAI({ apiKey, ...(baseURL ? { baseURL } : {}) });
    const res = await client.embeddings.create({ model, input: texts });
    return res.data.map((d) => d.embedding);
}
async function main() {
    const { query, topK } = parseArgs();
    const [vec] = await embedTexts([query]);
    const vectorProvider = (process.env.INDEXER_VECTOR_PROVIDER || process.env.SEARCH_VECTOR_PROVIDER || 'milvus').toLowerCase();
    if (vectorProvider === 'pinecone') {
        const apiKey = process.env.PINECONE_API_KEY;
        const indexName = process.env.PINECONE_INDEX;
        const namespace = process.env.PINECONE_NAMESPACE || 'default';
        if (!apiKey || !indexName)
            throw new Error('Pinecone requires PINECONE_API_KEY and PINECONE_INDEX');
        const pc = new Pinecone({ apiKey });
        const index = pc.index(indexName).namespace(namespace);
        const q = await index.query({ topK, vector: vec, includeMetadata: true });
        console.log(JSON.stringify(q.matches || [], null, 2));
        return;
    }
    // default milvus
    const address = process.env.MILVUS_ADDRESS || 'localhost:19530';
    const username = process.env.MILVUS_USERNAME || undefined;
    const password = process.env.MILVUS_PASSWORD || undefined;
    const collection = process.env.MILVUS_COLLECTION || 'code_chunks';
    const annsField = process.env.MILVUS_VECTOR_FIELD || 'vector';
    const metric = MetricType[process.env.MILVUS_METRIC || 'IP'] || MetricType.IP;
    const client = new MilvusClient({ address, username, password, ssl: false });
    try {
        await client.loadCollectionSync({ collection_name: collection });
    }
    catch { }
    const res = await client.search({
        collection_name: collection,
        vectors: [vec],
        vector_type: DataType.FloatVector,
        anns_field: annsField,
        topk: topK,
        metric_type: metric,
        params: { nprobe: 10 },
        output_fields: ['id', 'repo', 'path', 'url', 'commit', 'meta_json']
    });
    const out = [];
    for (const hit of (res.results || res.results_fields || [])) {
        const fields = hit.fields || hit;
        out.push({
            id: fields.id,
            score: hit.score || fields.distance || 0,
            repo: fields.repo,
            path: fields.path,
            url: fields.url,
            commit: fields.commit,
            meta: fields.meta_json ? JSON.parse(fields.meta_json) : undefined,
        });
    }
    console.log(JSON.stringify(out, null, 2));
}
main().catch((e) => {
    console.error('[test-codebase-search] failed:', e?.message || e);
    process.exit(1);
});