pediatric-ai-scribe-v3/test/generated-image-tools.test.js

const test = require('node:test');
const assert = require('node:assert/strict');
const fs = require('node:fs');
const vm = require('node:vm');
function loadAI(provider = 'litellm', chunks, content = 'Body [3].', backend) {
  const requests = [];
  class OpenAI { constructor() { this.chat = { completions: { create: async r => {
    requests.push(r);
    if (r.stream) return (async function*() { for (const chunk of chunks || []) yield chunk; })();
    return { choices: [{ message: { content, tool_calls: [{ id: 'one', type: 'function', function: { name: 'generate_image', arguments: '{"prompt":"diagram"}' } }] }, finish_reason: 'tool_calls' }] };
  } } }; } }
  const context = { module: { exports: {} }, console: { log() {}, error() {} }, process: { env: { AI_PROVIDER: provider, LITELLM_API_BASE: backend?.url || 'https://synthetic.invalid', OPENROUTER_API_KEY: 'synthetic', ...(provider === 'bedrock' ? { AWS_BEDROCK_REGION: 'synthetic-region' } : {}), ...(provider === 'azure' ? { AZURE_OPENAI_ENDPOINT: 'https://azure.synthetic.invalid', AZURE_OPENAI_API_KEY: 'synthetic', AZURE_DEPLOYMENT_NAME: 'test' } : {}) } }, require(name) {
    if (name === 'openai') return backend ? require('openai') : { OpenAI };
    if (name === '@aws-sdk/client-bedrock-runtime') return { BedrockRuntimeClient: class {} };
    if (name === './models') return { FALLBACK_MODEL: 'fallback', getEffectiveDefaultModel: async () => 'test', getAllowedModelIds: async () => new Set(['test']) };
    if (name === '../db/database') return { getSetting: async () => 'false' };
    if (name === './logger') return { apiCall() {}, error() {} };
    if (name === './generationOptions') return require('../src/utils/generationOptions');
    throw Error(name);
  } };
  vm.runInNewContext(fs.readFileSync('src/utils/ai.js', 'utf8'), context);
  return { ai: context.module.exports, requests };
}
const tools = [{ type: 'function', function: { name: 'generate_image', parameters: { type: 'object' } } }];
test('compatible call sends tools, returns tool_calls; ordinary request unchanged', async () => {
  const { ai, requests } = loadAI();
  const result = await ai.callAI([], { model: 'test', tools });
  assert.deepEqual(requests[0].tools, tools);
  assert.equal(result.toolCalls[0].function.name, 'generate_image');
  await ai.callAI([], { model: 'test' });
  assert.equal('tools' in requests[1], false);
});
test('fragmented streaming tool arguments are accumulated without losing body', async () => {
  const chunks = [
    { choices: [{ delta: { content: 'Body [3].', tool_calls: [{ index: 0, id: 'one', type: 'function', function: { name: 'generate_', arguments: '{"prompt":' } }] } }] },
    { choices: [{ delta: { tool_calls: [{ index: 0, function: { name: 'image', arguments: '"diagram"}' } }] }, finish_reason: 'tool_calls' }] }
  ];
  const { ai, requests } = loadAI('litellm', chunks);
  let text = '';
  const result = await ai.callAIStream([], { model: 'test', tools }, t => { text += t; });
  assert.deepEqual(requests[0].tools, tools);
  assert.equal(text, 'Body [3].');
  assert.equal(result.toolCalls[0].function.arguments, '{"prompt":"diagram"}');
  assert.equal(result.toolCalls[0].function.name, 'generate_image');
});
const imageTool = require('../src/utils/imageTool');
const express = require('express');
const realImages = require('../src/utils/generatedImages');
const id = '12345678-1234-1234-1234-123456789abc';
test('legacy direct provider rejects tool mode before calling any provider', async () => {
  const {ai,requests} = loadAI('bedrock');
  await assert.rejects(ai.callAI([], {model:'test',tools}), /Tools require/);
  await assert.rejects(ai.callAIStream([], {model:'test',tools},()=>{}), /Tools require/);
  assert.equal(requests.length,0);
});
test('tool validation caps invocations/fields; first-body continuation disables tools and never rewrites an existing body', async () => {
  const calls=[]; const continuations=[];
  const opts = { owner:101,workflow:'clinical_assistant',body:{message:'diagram',idempotencyKey:'request'},imageContext:{request:'diagram',history:[]},messages:[{role:'user',content:'diagram'}],options:{model:'test'},
    images:{enqueue:async(...v)=>{calls.push(v);return {jobId:id,status:'pending'};}},callAI:async(...v)=>{continuations.push(v);return {content:'First body [3].',finishReason:'stop'};} };
  const call = {id:'one',type:'function',function:{name:'generate_image',arguments:'{"prompt":"diagram","layout":"portrait"}'}};
  const original = 'Body [3].\n\nTrailing words with';
  const result = await imageTool.dispatch({content:original,toolCalls:[call],finishReason:'tool_calls'},opts);
  assert.equal(result.content,original); assert.equal(continuations.length,0); assert.equal(calls.length,1);
  assert.equal(calls[0][0],101); assert.equal(calls[0][1],'clinical_assistant'); assert.equal(calls[0][3],'tool:request');
  const first = await imageTool.dispatch({content:null,toolCalls:[call]},opts);
  assert.equal(first.content,'First body [3].'); assert.equal(continuations.length,1); assert.equal(continuations[0][1].toolChoice,'none');
  assert.equal(continuations[0][0].at(-1).role,'tool');
  for (const bad of [[call,call],[{...call,function:{name:'fetch_url',arguments:'{}'}}],[{...call,function:{name:'generate_image',arguments:'{"prompt":"x","model":"bad"}'}}],[{...call,function:{name:'generate_image',arguments:'{'}}]]) {
    await assert.rejects(imageTool.dispatch({content:original,toolCalls:bad},opts));
  }
  await assert.rejects(imageTool.dispatch({content:original,toolCalls:[call]},{...opts,imageContext:undefined}),/original image request/);
  assert.equal(calls.length,2);
});
function route(file, ai, jobs) {
  const mocks = {
    express, axios:{}, crypto:require('crypto'), multer:require('multer'), path:require('path'),
    '../utils/ai':ai, '../db/database':{getSetting:async key=>key.includes('model')?'test':null,all:async()=>[]},
    '../middleware/auth':{authMiddleware(){},moderatorMiddleware(){}},
    '../utils/crypto':{},'../utils/urlSafety':{},'../utils/policy':{requireFeature:()=>()=>{}},
    // Retrieval is opt-in and these cases do not ask for it; the stub proves
    // the route never reaches the corpus unless useCorpus was set.
    '../utils/learningRetrieval':{ retrieve: async () => { throw new Error('retrieval must not run unless requested'); } },
    '../utils/logger':{audit(){},error(){}},'../utils/redis':{},'../utils/clinicalPromptPool':{createClinicalPromptPool:()=>({})},
    '../utils/clinicalPrompts':require('../src/utils/clinicalPrompts'),
    '../utils/clinicalConversation':require('../src/utils/clinicalConversation'),
    '../utils/clinicalAnswer':require('../src/utils/clinicalAnswer'),
    // Citation quality tracking, required lazily by the streaming route.
    '../utils/citationAudit':require('../src/utils/citationAudit'),
    '../utils/clinicalTranslation':require('../src/utils/clinicalTranslation'),
    '../utils/patientTakehome':require('../src/utils/patientTakehome'),
    './auth':{__sendEmail:async()=>false},
    '../utils/generatedImages':realImages, '../utils/generatedImageLinks':require('../src/utils/generatedImageLinks'),
    '../utils/visionTool': require('../src/utils/visionTool'),
    '../utils/modelVision': { supportsVision: async () => null },
    '../utils/imageTool':{tools:imageTool.tools,dispatch:(value,options)=>imageTool.dispatch(value,{...options,images:{enqueue:async(...args)=>{jobs.push(args);return {jobId:id,status:'pending'};}}})},
    '../utils/clinicalMcpClient':{semanticSearch:async()=>({})},
    '../utils/clinicalRetrieval':{normalizeMcpSearchResponse:()=>[{number:3,title:'Synthetic source',page:17,excerpt:'Synthetic reference'}],dedupeSources:s=>s,
      normalizeMcpMultimodalResponse:()=>[],cleanSourceExcerpt:s=>s,isVisualSourceQuery:()=>false,classifyAndRerankMultimodalResults:async()=>[]}
  };
  const module = {exports:{}};
  vm.runInNewContext(fs.readFileSync(file,'utf8'),{module,Buffer,console:{info(){},warn(){},error(){}},process:{env:{CLINICAL_ASSISTANT_MCP_WARMUP:'false'}},setTimeout(){},require:n=>{assert.ok(n in mocks,n);return mocks[n];}});
  return async (path,body) => {
    const endpoint = module.exports.stack.find(l=>l.route?.path===path).route.stack.at(-1).handle;
    const response = {statusCode:200,events:'',status(s){this.statusCode=s;return this;},json(d){this.data=d;},setHeader(){},flushHeaders(){},write(t){this.events+=t;},end(){}};
    await endpoint({user:{id:101},body},response);return response;
  };
}
test('actual Clinical chat + fragmented stream send real tools to SDK, dispatch jobs, and preserve citation/body/source/page identity without regeneration', async () => {
  const body='Exact [3].\n\nSentence ending with'; // Old truncation heuristic would regenerate this.
  for (const streaming of [false,true]) {
    const {ai,requests}=loadAI('litellm',[
      {choices:[{delta:{content:body,tool_calls:[{index:0,id:'one',type:'function',function:{name:'generate_image',arguments:'{"prompt":'}}]}}]},
      {choices:[{delta:{tool_calls:[{index:0,function:{arguments:'"diagram"}'}}]},finish_reason:'tool_calls'}]}
    ],body);
    const jobs=[];const request=route('src/routes/clinicalAssistant.js',ai,jobs);
    const response=await request('/clinical-assistant/chat'+(streaming?'/stream':''),{message:'  Create a clinical diagram for the precise current clinical findings 😀\n',history:[{role:'assistant',content:'Prior table [3].'}],idempotencyKey:'same-request'});
    assert.equal(response.statusCode,200,JSON.stringify(response.data));
    const result=streaming?JSON.parse(response.events.match(/event: done\ndata: (.*)/)[1]):response.data;
    assert.equal(result.answer,body);assert.equal(result.sources[0].number,3);assert.equal(result.sources[0].page,17);assert.equal(result.sources[0].title,'Synthetic source');
    assert.equal(result.imageJobs[0].jobId,id);assert.equal(jobs.length,1);assert.equal(jobs[0][5].request,'  Create a clinical diagram for the precise current clinical findings 😀\n');assert.equal(jobs[0][5].history[0].content,'Prior table [3].');assert.equal(requests.length,1);assert.equal(requests[0].tools[0].function.name,'generate_image');
  }
});
test('actual Learning generate/refine use callable tools and bind their own workflow (not sidebar heuristics)', async()=>{
  for (const [path,content,input] of [
    ['/ai-generate','{"title":"Teaching","body":"<p>Exact body.</p>","questions":[]}',{topic:'Create a diagram',idempotencyKey:'gen'}],
    ['/ai-refine','<p>Exact refined body.</p>',{content:'<p>Prior body.</p>',instructions:'Include an image',idempotencyKey:'ref'}]
  ]) {
    const {ai,requests}=loadAI('litellm',undefined,content);const jobs=[];
    const request=route('src/routes/learningAI.js',ai,jobs);const response=await request(path,input);
    assert.equal(response.statusCode,200,JSON.stringify(response.data));assert.equal(response.data.success,true);assert.equal(response.data.imageJobs[0].jobId,id);
    assert.equal(requests.length,1);assert.equal(requests[0].tools[0].function.name,'generate_image');assert.equal(jobs[0][1],'learning_hub');
    assert.equal(path==='/ai-refine'?response.data.refined:response.data.content.body,path==='/ai-refine'?input.content:'<p>Exact body.</p>');
  }
});
test('Learning tool-only refinement retains every existing body/citation/page byte, including edge whitespace', async () => {
  const content = '\n  <p>Exact [3, 1].</p>\n<table><tr><td>5 mg</td><td>page 19 [3]</td></tr></table>  \n';
  const { ai, requests } = loadAI('litellm', undefined, null);
  const jobs = []; const request = route('src/routes/learningAI.js', ai, jobs);
  const response = await request('/ai-refine', { content, instructions: 'Create a matching diagram', idempotencyKey: 'exact-refine' });
  assert.equal(response.statusCode, 200); assert.equal(response.data.refined, content);
  assert.equal(requests.length, 1); assert.equal(jobs.length, 1);
});
test('all compatible providers capture tool_choice/parallel cap, fragmented IDs, and unchanged ordinary streams', async () => {
  const chunks = [
    { choices: [{ delta: { tool_calls: [{ index: 0, id: 'to', type: 'function', function: { name: 'generate_', arguments: '{"prompt":' } }] } }] },
    { choices: [{ delta: { tool_calls: [{ index: 0, id: 'ol', function: { name: 'image', arguments: '"diagram"}' } }], content: 'Body [3].' }, finish_reason: 'tool_calls' }] }
  ];
  for (const provider of ['litellm', 'openrouter', 'azure']) {
    const { ai, requests } = loadAI(provider, chunks);
    await ai.callAI([], { model: 'test', tools, toolChoice: 'none' });
    assert.equal(requests[0].tool_choice, 'none'); assert.equal(requests[0].parallel_tool_calls, false);
    const result = await ai.callAIStream([], { model: 'test', tools });
    assert.equal(result.provider, provider); assert.equal(result.toolCalls[0].id, 'tool');
    assert.equal(result.toolCalls[0].function.name, 'generate_image'); assert.equal(result.toolCalls[0].function.arguments, '{"prompt":"diagram"}');
    assert.equal(requests[1].tool_choice, 'auto'); assert.equal(requests[1].parallel_tool_calls, false);
    const ordinary = loadAI(provider, [{ choices: [{ delta: { content: 'No tools [3].' }, finish_reason: 'stop' }] }]);
    const text = await ordinary.ai.callAIStream([], { model: 'test' });
    assert.equal(text.content, 'No tools [3].'); assert.equal('toolCalls' in text, false);
    assert.equal('tools' in ordinary.requests[0], false); assert.equal('parallel_tool_calls' in ordinary.requests[0], false);
  }
});
test('malformed/oversized fragmented tools stop before job dispatch; no repeated continuation', async () => {
  for (const fragment of [
    { index: -1 }, { index: 8 }, { index: 0, type: 'unknown' },
    { index: 0, function: { arguments: 'x'.repeat(40001) } },
    { index: 0, function: { name: 'x'.repeat(101) } }, { index: 0, id: 'x'.repeat(201) }
  ]) {
    const { ai } = loadAI('litellm', [{ choices: [{ delta: { tool_calls: [fragment] } }] }]);
    await assert.rejects(ai.callAIStream([], { model: 'test', tools }));
  }
  let paid = 0, continuations = 0;
  const call = { id: 'one', type: 'function', function: { name: 'generate_image', arguments: '{"prompt":"diagram"}' } };
  await assert.rejects(imageTool.dispatch({ content: null, toolCalls: [call] }, {
    owner: 101, workflow: 'learning_hub', body: {}, imageContext:{request:'diagram',history:[]}, messages: [], options: {},
    images: { enqueue: async () => { paid++; return { jobId: id, status: 'pending' }; } },
    callAI: async () => { continuations++; return { content: '', toolCalls: [call] }; }
  }), /did not return educational content/);
  assert.equal(paid, 1); assert.equal(continuations, 1);
});
test('actual OpenAI SDK HTTP capture from Clinical routes carries tools and dispatches fragmented SSE once', async () => {
  const http = require('node:http'); const captured = [];
  const body = 'Exact transport body [3].\n| Dose | Page |\n| 5 mg | 17 [3] |\n';
  const call = { id: 'one', type: 'function', function: { name: 'generate_image', arguments: '{"prompt":"Diagram 😀","layout":"portrait"}' } };
  const server = http.createServer(async (req, res) => {
    const chunks = []; for await (const chunk of req) chunks.push(chunk);
    const payload = JSON.parse(Buffer.concat(chunks).toString()); captured.push({ path: req.url, payload });
    if (!payload.stream) {
      res.writeHead(200, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ id: 'synthetic', object: 'chat.completion', choices: [{ index: 0, message: { role: 'assistant', content: body, tool_calls: [call] }, finish_reason: 'tool_calls' }] }));
    } else {
      res.writeHead(200, { 'Content-Type': 'text/event-stream' });
      const fragments = [
        { choices: [{ index: 0, delta: { content: body, tool_calls: [{ index: 0, id: 'o', type: 'function', function: { name: 'generate_', arguments: '{"prompt":"Diagram ' } }] } }] },
        { choices: [{ index: 0, delta: { tool_calls: [{ index: 0, id: 'ne', function: { name: 'image', arguments: '😀","layout":"portrait"}' } }] }, finish_reason: 'tool_calls' }] }
      ];
      const bytes = Buffer.from(fragments.map(f => 'data: ' + JSON.stringify(f) + '\n\n').join('') + 'data: [DONE]\n\n');
      const split = bytes.indexOf(Buffer.from('😀')) + 2; // Split a UTF-8 character at the transport boundary too.
      res.write(bytes.subarray(0, split)); setImmediate(() => res.end(bytes.subarray(split)));
    }
  });
  server.listen(0, '127.0.0.1'); await new Promise(resolve => server.once('listening', resolve));
  try {
    const { ai } = loadAI('litellm', undefined, undefined, { url: 'http://127.0.0.1:' + server.address().port + '/v1' });
    for (const streaming of [false, true]) {
      const jobs = []; const request = route('src/routes/clinicalAssistant.js', ai, jobs);
      const response = await request('/clinical-assistant/chat' + (streaming ? '/stream' : ''), { message: 'Create a diagram', idempotencyKey: 'transport' });
      assert.equal(response.statusCode, 200);
      const result = streaming ? JSON.parse(response.events.match(/event: done\ndata: (.*)/)[1]) : response.data;
      assert.equal(result.answer, body); assert.equal(result.imageJobs[0].jobId, id); assert.equal(result.sources[0].number, 3); assert.equal(result.sources[0].page, 17);
      assert.equal(jobs.length, 1); assert.equal(jobs[0][2].prompt, 'Diagram 😀'); assert.equal(jobs[0][2].layout, 'portrait');
    }
    assert.equal(captured.length, 2);
    for (const request of captured) {
      assert.equal(request.path, '/v1/chat/completions'); assert.equal(request.payload.tools[0].function.name, 'generate_image');
      assert.equal(request.payload.parallel_tool_calls, false); assert.equal(request.payload.tool_choice, 'auto');
    }
  } finally { await new Promise(resolve => server.close(resolve)); }
});
test('Learning image refinement ignores accompanying rewritten HTML/citations/table; text-only refinement still works', async () => {
  const original='\n <p>Original [3, 1].</p><table><tr><td>5 mg</td><td>19 [3]</td></tr></table> \n';
  for (const withTool of [true,false]) {
    const jobs=[];const replacement='<p>ALTERED [9].</p>';
    const request=route('src/routes/learningAI.js',{callAI:async()=>({content:replacement,...(withTool?{toolCalls:[{id:'x',type:'function',function:{name:'generate_image',arguments:'{"prompt":"diagram"}'}}]}:{})})},jobs);
    const response=await request('/ai-refine',{content:original,instructions:withTool?'Include an image':'Shorten this text'});
    assert.equal(response.statusCode,200);assert.equal(response.data.refined,withTool?original:replacement);
    assert.equal(response.data.bodyPreserved,withTool);assert.equal(jobs.length,withTool?1:0);
  }
});