TalkingBot/chat-engine/tests/hybrid.test.js

const { mockClient } = require('aws-sdk-client-mock');
const { BedrockRuntimeClient, ConverseStreamCommand, InvokeModelCommand } = require('@aws-sdk/client-bedrock-runtime');
const { runHybrid } = require('../src/strategies/hybridAgent');
const { buildSearchRequests } = require('../src/tools/searchBook');
const { capToolResult } = require('../src/tools/execute');
const { getPages } = require('../src/tools/getPages');
const { chatConfig, sampleBook, textStream, usage } = require('./fixture');

const bedrock = mockClient(BedrockRuntimeClient);

function toolStream(calls) {
  return (async function* () {
    let index = 0;
    for (const call of calls) {
      yield {
        contentBlockStart: {
          contentBlockIndex: index,
          start: { toolUse: { toolUseId: call.id, name: call.name } },
        },
      };
      yield {
        contentBlockDelta: {
          contentBlockIndex: index,
          delta: { toolUse: { input: JSON.stringify(call.input) } },
        },
      };
      index += 1;
    }
    yield { messageStop: { stopReason: 'tool_use' } };
    yield { metadata: { usage: usage({ outputTokens: 4 }) } };
  })();
}

function searchClient() {
  return {
    async search() {
      return {
        body: {
          hits: {
            hits: [
              {
                _id: 'c1',
                _source: {
                  chunk_id: 'c1',
                  book_id: '3',
                  chapter_no: 1,
                  chapter_title: 'Opening',
                  page_start: 2,
                  page_end: 2,
                  text: 'Opening body',
                  token_count: 20,
                  prev_chunk_id: null,
                  next_chunk_id: null,
                },
              },
            ],
          },
        },
      };
    },
  };
}

async function collect(iterable) {
  const events = [];
  for await (const event of iterable) events.push(event);
  return events;
}

describe('hybrid agent', () => {
  let client;

  beforeEach(() => {
    bedrock.reset();
    client = new BedrockRuntimeClient({ region: 'ap-south-1' });
    bedrock.on(InvokeModelCommand).resolves({ body: Buffer.from(JSON.stringify({ embedding: [0.1, 0.2] })) });
  });

  function ctx(overrides = {}) {
    return {
      config: chatConfig(),
      logger: { info() {}, warn() {} },
      bedrock: client,
      opensearch: searchClient(),
      book: sampleBook('short'),
      session: { channel: 'chat', history: [], rollingSummary: '', lastCacheTouchAt: null, sessionId: 's' },
      userText: 'Where is the workshop?',
      ...overrides,
    };
  }

  it('runs one tool call and then answers', async () => {
    let calls = 0;
    bedrock.on(ConverseStreamCommand).callsFake(async (input) => {
      calls += 1;
      return {
        stream: input.toolConfig && calls === 1
          ? toolStream([{ id: 't1', name: 'search_book', input: { query: 'workshop', book_id: '999' } }])
          : textStream('It is on (p. 3).', usage({ outputTokens: 6 })),
      };
    });
    const events = await collect(runHybrid(ctx()));
    expect(events.filter((event) => event.type === 'tool_start')).toHaveLength(1);
    expect(events.find((event) => event.type === 'final').text).toContain('(p. 3)');
  });

  it('runs parallel tool calls', async () => {
    let calls = 0;
    bedrock.on(ConverseStreamCommand).callsFake(async (input) => {
      calls += 1;
      return {
        stream: input.toolConfig && calls === 1
          ? toolStream([
              { id: 't1', name: 'search_book', input: { query: 'morning' } },
              { id: 't2', name: 'search_book', input: { query: 'afternoon' } },
            ])
          : textStream('Both parts differ.', usage()),
      };
    });
    const events = await collect(runHybrid(ctx()));
    expect(events.filter((event) => event.type === 'tool_start').map((event) => event.name)).toEqual(['search_book', 'search_book']);
  });

  it('forces a tool-less answer on the last iteration', async () => {
    bedrock.on(ConverseStreamCommand).callsFake(async (input) => ({
      stream: input.toolConfig
        ? toolStream([{ id: 't1', name: 'get_chapter', input: { chapter_no: 1 } }])
        : textStream('Using what I have.', usage()),
    }));
    const events = await collect(runHybrid(ctx({ config: chatConfig({ agentMaxIterations: 2, agentMaxToolCalls: 6 }) })));
    expect(events.find((event) => event.type === 'final').text).toBe('Using what I have.');
    expect(events.filter((event) => event.type === 'tool_start').length).toBeLessThanOrEqual(6);
  });

  it('stops when the turn times out', async () => {
    const events = await collect(runHybrid(ctx({ config: chatConfig({ agentTurnTimeoutMs: 0 }) })));
    expect(events.find((event) => event.type === 'error').message).toMatch(/timed out/);
    expect(bedrock.commandCalls(ConverseStreamCommand)).toHaveLength(0);
  });

  it('continues after a tool error', async () => {
    bedrock.on(ConverseStreamCommand).callsFake(async (input) => ({
      stream: input.toolConfig
        ? toolStream([{ id: 't1', name: 'get_pages', input: { start: 5, end: 1 } }])
        : textStream('I could not read those pages.', usage()),
    }));
    const events = await collect(runHybrid(ctx()));
    expect(events.find((event) => event.type === 'error')).toBeUndefined();
    expect(events.find((event) => event.type === 'final').text).toContain('could not');
  });

  it('marks a citation verified only when that page came back from a tool', async () => {
    let calls = 0;
    bedrock.on(ConverseStreamCommand).callsFake(async (input) => {
      calls += 1;
      return {
        stream: input.toolConfig && calls === 1
          ? toolStream([{ id: 't1', name: 'get_pages', input: { start: 2, end: 2 } }])
          : textStream('See [[p:2]] and [[p:3]].', usage({ outputTokens: 8 })),
      };
    });
    const events = await collect(runHybrid(ctx()));
    const spoken = events.filter((event) => event.type === 'text_delta').map((event) => event.text).join('');
    expect(spoken).toBe('See and.');
    const items = events.find((event) => event.type === 'citations').items;
    expect(items.map((item) => [item.page, item.verified])).toEqual([[2, true], [3, false]]);
  });

  it('stops when the caller aborts', async () => {
    const abort = new AbortController();
    bedrock.on(ConverseStreamCommand).callsFake(async () => ({
      stream: (async function* () {
        yield { contentBlockDelta: { contentBlockIndex: 0, delta: { text: 'Hello' } } };
        abort.abort();
        yield { contentBlockDelta: { contentBlockIndex: 0, delta: { text: ' more' } } };
        yield { messageStop: { stopReason: 'end_turn' } };
        yield { metadata: { usage: usage() } };
      })(),
    }));
    const events = await collect(runHybrid(ctx({ signal: abort.signal, config: chatConfig({ agentMaxIterations: 1 }) })));
    const finalEvent = events.find((event) => event.type === 'final');
    expect(finalEvent.interrupted).toBe(true);
    expect(finalEvent.text).toBe('Hello');
  });
});

describe('tool guards', () => {
  it('does not let the model replace book_id', () => {
    const requests = buildSearchRequests('3', { query: 'river', book_id: '999', chapter_no: 1 }, [0.1, 0.2]);
    expect(JSON.stringify(requests.knn.post_filter)).toContain('"book_id":"3"');
    expect(JSON.stringify(requests.knn.post_filter)).not.toContain('999');
    expect(JSON.stringify(requests.bm25)).toContain('"book_id":"3"');
  });

  it('truncates oversized tool results and page ranges', () => {
    const capped = JSON.parse(capToolResult({ text: 'a'.repeat(500) }, 80));
    expect(capped.truncated).toBe(true);
    const pages = getPages(sampleBook(), 1, 9);
    expect(pages.page_end - pages.page_start).toBeLessThanOrEqual(4);
    expect(pages.truncated).toBe(true);
  });
});