diff --git a/.env.example b/.env.example index f95c57d0..541ae47a 100644 --- a/.env.example +++ b/.env.example @@ -260,6 +260,7 @@ DB_PASSWORD=pedscribe_secret_change_me # CLINICAL_ASSISTANT_MCP_CONCURRENCY=3 # library searches in flight at once; they used to run one at a time # CLINICAL_ASSISTANT_MCP_QUEUE_MAX= # callers allowed to wait for a slot (default 4 x concurrency); past it: "busy, try again" # CLINICAL_ASSISTANT_MCP_QUEUE_WAIT_MS=8000 # how long a caller waits in that line before being told the library is busy +# CLINICAL_ASSISTANT_COLLECTIONS_TTL_MS=600000 # how often the list of search collections is re-read (off the query path); with one collection nothing changes # CLINICAL_ASSISTANT_ASK_LIMIT_PER_MINUTE=30 # paid assistant questions per user per minute, counted in Redis (no Redis: no limit) # CLINICAL_ASSISTANT_RETRIEVAL_CACHE_TTL_S=60 # same library search within this window is answered from Redis; 0 disables # CLINICAL_ASSISTANT_MCP_WARMUP= # open a session at boot diff --git a/src/utils/clinicalMcpClient.js b/src/utils/clinicalMcpClient.js index 087b7566..2f7992d0 100644 --- a/src/utils/clinicalMcpClient.js +++ b/src/utils/clinicalMcpClient.js @@ -1,5 +1,6 @@ var axios = require('axios'); var sleep = require('node:timers/promises').setTimeout; +var fanout = require('./collectionFanout'); var MCP_URLS = buildMcpUrls(); var _lastGoodMcpUrl = MCP_URLS[0]; @@ -87,9 +88,8 @@ function closeMcpSession() { return _closePromise; } -async function semanticSearch(query, opts) { - opts = opts || {}; - return callMcpTool(SEARCH_TOOL_NAME, { +function searchArgs(query, opts, collection) { + var args = { query: query, limit: opts.limit, doc_types: ['file'], @@ -97,7 +97,65 @@ async function semanticSearch(query, opts) { fusion: 'rrf', include_context: opts.includeContext, context_chars: opts.contextChars + }; + // The default collection is asked for by omission, exactly as before the + // service could hold more than one. + if (collection) args.collection = collection; + return args; +} + +async function semanticSearch(query, opts) { + opts = opts || {}; + var extra = knownExtraCollections(); + if (extra.length === 0) return callMcpTool(SEARCH_TOOL_NAME, searchArgs(query, opts)); + // Every collection at once, the default included. A collection whose call + // fails is left out of the fused list rather than failing the answer; the + // default failing is the same error it always was. + var calls = [{ collection: null, promise: callMcpTool(SEARCH_TOOL_NAME, searchArgs(query, opts)) }] + .concat(extra.map(function(name) { + return { collection: name, promise: callMcpTool(SEARCH_TOOL_NAME, searchArgs(query, opts, name)) }; + })); + var settled = await Promise.allSettled(calls.map(function(c) { return c.promise; })); + if (settled[0].status === 'rejected') throw settled[0].reason; + var responses = []; + settled.forEach(function(s, i) { + if (s.status === 'fulfilled') responses.push({ collection: calls[i].collection, response: s.value }); + else console.warn('[clinical-assistant] collection search skipped:', calls[i].collection, s.reason && s.reason.message); }); + return fanout.mergeCollectionResults(responses, positiveInt(opts.limit, 0)); +} + +// The collections the service will search besides its default, learned off +// the query path: at warm-up and on the session timer, never while someone +// waits for an answer. Until the first listing arrives, or on a service that +// has no listing tool, this is empty and a search is the single call it +// always was. +var COLLECTIONS_TTL_MS = positiveInt(process.env.CLINICAL_ASSISTANT_COLLECTIONS_TTL_MS, 10 * 60 * 1000); +var _extraCollections = []; +var _collectionsFetchedAt = 0; +var _collectionsPromise = null; + +function knownExtraCollections() { + return _extraCollections; +} + +function refreshCollections() { + if (_collectionsPromise) return _collectionsPromise; + _collectionsPromise = callMcpTool('clinical_list_collections', {}).then(function(listing) { + _extraCollections = fanout.extraCollectionNames(listing); + _collectionsFetchedAt = Date.now(); + return _extraCollections; + }).catch(function(e) { + // No listing tool, or a transient error: keep whatever was known. + _collectionsFetchedAt = Date.now(); + console.info('[clinical-assistant] collection listing unavailable:', e && e.message); + return _extraCollections; + }).finally(function() { _collectionsPromise = null; }); + return _collectionsPromise; +} + +function collectionsStale() { + return Date.now() - _collectionsFetchedAt > COLLECTIONS_TTL_MS; } async function indexedTopicSuggestions(limit) { @@ -117,10 +175,14 @@ function warmMcpSession() { if (soon) { var old = _mcpSession; _mcpSession = null; endMcpSession(old); } getMcpSession().catch(function() {}); } + if (collectionsStale()) refreshCollections(); }, Math.max(30000, Math.floor(MCP_SESSION_TTL_MS / 3))); if (_warmTimer.unref) _warmTimer.unref(); } - return getMcpSession(); + return getMcpSession().then(function(session) { + if (collectionsStale()) refreshCollections(); + return session; + }); } function busyError() { @@ -402,6 +464,8 @@ function parseMcpResponse(body) { module.exports = { semanticSearch: semanticSearch, + knownExtraCollections: knownExtraCollections, + refreshCollections: refreshCollections, queueDepth: queueDepth, indexedTopicSuggestions: indexedTopicSuggestions, getMcpHealth: getMcpHealth, diff --git a/src/utils/collectionFanout.js b/src/utils/collectionFanout.js new file mode 100644 index 00000000..7b0a623d --- /dev/null +++ b/src/utils/collectionFanout.js @@ -0,0 +1,103 @@ +// ============================================================ +// COLLECTION FAN-OUT +// Searching more than one Milvus collection without slowing down the one. +// +// The search service can hold several collections, each embedded with its +// own model (MILVUS_COLLECTIONS on the service; clinical_list_collections +// reports them). A search names one collection, or none for the default. So +// a second collection is only useful to the assistant if the client asks +// both — and the one-collection case, which is every deployment today, must +// cost exactly what it did: one call, no listing on the query path. +// +// The list is therefore fetched off the query path (at session warm-up and +// on a slow timer) and the fan-out is decided from that cache. With one +// collection cached, or none yet, the search is the single call it always +// was. With more, one call per collection runs in parallel and the ranked +// lists are fused by reciprocal rank — the same fusion the service uses +// within a collection — so a strong hit in a small collection is not buried +// under a large one's scores, which are not comparable across embedders. +// ============================================================ + +var RRF_K = 60; + +// The service's result list, wherever the MCP envelope put it. +function extractResults(response) { + var data = response && (response.structuredContent || response.data || response); + if ((!data || !Array.isArray(data.results)) && response && Array.isArray(response.content)) { + for (var i = 0; i < response.content.length; i++) { + var c = response.content[i]; + if (c && c.type === 'text' && c.text) { + try { + var parsed = JSON.parse(c.text); + if (parsed && Array.isArray(parsed.results)) data = parsed; + } catch (e) {} + } + } + } + return data && Array.isArray(data.results) ? data : null; +} + +function resultKey(r) { + if (r && r.id != null) return String(r.id); + return [r && r.file_path, r && r.page_number, r && r.chunk_index].join('|'); +} + +// Fuse per-collection ranked lists into one list of at most `limit`. +// `responses` is [{ collection, response }]; a collection whose call failed +// is simply absent. Each merged result carries the collection it came from. +function mergeCollectionResults(responses, limit) { + var fused = new Map(); + var verified = 0, dropped = 0, searched = []; + responses.forEach(function(entry) { + var data = extractResults(entry.response); + if (!data) return; + searched.push(entry.collection); + verified += Number(data.verified_chunk_count || data.verifiedChunkCount || 0); + dropped += Number(data.dropped_document_count || data.droppedDocumentCount || 0); + data.results.forEach(function(r, rank) { + var key = resultKey(r); + var slot = fused.get(key); + var contribution = 1 / (RRF_K + rank + 1); + if (slot) { slot.fused += contribution; return; } + fused.set(key, { fused: contribution, result: Object.assign({}, r, { collection: entry.collection }) }); + }); + }); + var merged = Array.from(fused.values()) + .sort(function(a, b) { return b.fused - a.fused; }) + .map(function(s) { return s.result; }); + if (limit > 0) merged = merged.slice(0, limit); + return { + results: merged, + verified_chunk_count: verified, + dropped_document_count: dropped, + collections: searched + }; +} + +// The names worth a separate call: everything the service lists beyond its +// default. An empty or one-entry list means "search as before". +function extraCollectionNames(listing) { + var data = listing && (listing.structuredContent || listing.data || listing); + if ((!data || !Array.isArray(data.collections)) && listing && Array.isArray(listing.content)) { + for (var i = 0; i < listing.content.length; i++) { + var c = listing.content[i]; + if (c && c.type === 'text' && c.text) { + try { + var parsed = JSON.parse(c.text); + if (parsed && Array.isArray(parsed.collections)) data = parsed; + } catch (e) {} + } + } + } + if (!data || !Array.isArray(data.collections)) return []; + return data.collections + .filter(function(c) { return c && c.name && !c.default && !c.is_default; }) + .map(function(c) { return String(c.name); }); +} + +module.exports = { + extractResults: extractResults, + mergeCollectionResults: mergeCollectionResults, + extraCollectionNames: extraCollectionNames, + RRF_K: RRF_K +}; diff --git a/test/mcp-collection-fanout.test.js b/test/mcp-collection-fanout.test.js new file mode 100644 index 00000000..14c5d3e2 --- /dev/null +++ b/test/mcp-collection-fanout.test.js @@ -0,0 +1,60 @@ +// A second search collection is searched alongside the first and fused by +// rank; with one collection the search is the single call it always was. +const test = require('node:test'); +const assert = require('node:assert/strict'); +const fs = require('node:fs'); +const path = require('node:path'); +const fanout = require('../src/utils/collectionFanout'); +const src = fs.readFileSync(path.join(__dirname, '..', 'src/utils/clinicalMcpClient.js'), 'utf8'); + +function listing(names, defaultName) { + return { structuredContent: { collections: names.map(n => ({ name: n, default: n === defaultName, model: 'm' })) } }; +} + +test('the default collection alone means nothing extra to search', () => { + assert.deepEqual(fanout.extraCollectionNames(listing(['mcp_bge_m3_1024'], 'mcp_bge_m3_1024')), []); + assert.deepEqual(fanout.extraCollectionNames({}), []); + assert.deepEqual(fanout.extraCollectionNames(null), []); +}); + +test('collections beyond the default are named, from either envelope', () => { + assert.deepEqual(fanout.extraCollectionNames(listing(['a', 'b', 'c'], 'a')), ['b', 'c']); + const textEnvelope = { content: [{ type: 'text', text: JSON.stringify({ collections: [{ name: 'a', default: true }, { name: 'peds' }] }) }] }; + assert.deepEqual(fanout.extraCollectionNames(textEnvelope), ['peds']); +}); + +test('ranked lists from two collections fuse by reciprocal rank and keep their origin', () => { + const first = { structuredContent: { results: [{ id: 'x', score: 0.9 }, { id: 'y', score: 0.8 }], verified_chunk_count: 2 } }; + const second = { structuredContent: { results: [{ id: 'z', score: 0.1 }, { id: 'y', score: 0.05 }], verified_chunk_count: 2, dropped_document_count: 1 } }; + const merged = fanout.mergeCollectionResults([{ collection: null, response: first }, { collection: 'peds', response: second }], 10); + // y appears in both lists (ranks 2 and 2) and outranks either list's first. + assert.equal(merged.results[0].id, 'y'); + assert.deepEqual(merged.results.map(r => r.id).sort(), ['x', 'y', 'z']); + assert.equal(merged.results.find(r => r.id === 'z').collection, 'peds'); + assert.equal(merged.results.find(r => r.id === 'x').collection, null); + assert.equal(merged.verified_chunk_count, 4); + assert.equal(merged.dropped_document_count, 1); + assert.deepEqual(merged.collections, [null, 'peds']); +}); + +test('the fused list respects the limit and a failed collection is simply absent', () => { + const one = { structuredContent: { results: [{ id: 1 }, { id: 2 }, { id: 3 }] } }; + const merged = fanout.mergeCollectionResults([{ collection: null, response: one }], 2); + assert.equal(merged.results.length, 2); + assert.deepEqual(merged.collections, [null]); +}); + +test('the search itself stays one call when no extra collection is known', () => { + // The single-call path is the first thing semanticSearch does, before any + // fan-out machinery, and it names no collection. + const body = src.slice(src.indexOf('async function semanticSearch'), src.indexOf('// The collections the service will search')); + assert.match(body, /if \(extra\.length === 0\) return callMcpTool\(SEARCH_TOOL_NAME, searchArgs\(query, opts\)\)/); + assert.match(src, /if \(collection\) args\.collection = collection;/); +}); + +test('the collection list is learned off the query path, never inside a search', () => { + const search = src.slice(src.indexOf('async function semanticSearch'), src.indexOf('// The collections the service will search')); + assert.doesNotMatch(search, /refreshCollections|clinical_list_collections/); + const warm = src.slice(src.indexOf('function warmMcpSession'), src.indexOf('function busyError')); + assert.match(warm, /if \(collectionsStale\(\)\) refreshCollections\(\);/); +});