Fixes context middleware to append to user message only. (#591)

js/ai/src/model/middleware.ts

@@ -231,7 +231,7 @@ export function augmentWithContext(
  return (req, next) => {
  // if there is no context in the request, no-op
  if (!req.context?.length) return next(req);
- const userMessage = req.messages.at(-1);
+ const userMessage = lastUserMessage(req.messages);
  // if there are no messages, no-op
  if (!userMessage) return next(req);
  // if there is already a context part, no-op

js/ai/tests/model/middleware_test.ts

@@ -18,16 +18,16 @@ import assert from 'node:assert';
 import { beforeEach, describe, it } from 'node:test';
 import { DocumentData } from '../../src/document.js';
 import {
+ defineModel,
  GenerateRequest,
  GenerateResponseData,
  MessageData,
  Part,
- defineModel,
 } from '../../src/model.js';
 import {
+ augmentWithContext,
  AugmentWithContextOptions,
  CONTEXT_PREFACE,
- augmentWithContext,
  simulateSystemPrompt,
  validateSupport,
 } from '../../src/model/middleware.js';
@@ -414,6 +414,24 @@ describe('augmentWithContext', () => {
  });
  });
 
+ it('should append to the last user message', async () => {
+ const messages: MessageData[] = [
+ { role: 'user', content: [{ text: 'first part' }] },
+ {
+ role: 'tool',
+ content: [{ toolResponse: { name: 'testTool', output: { abc: 123 } } }],
+ },
+ ];
+ const result = await testRequest(messages, [
+ { content: [{ text: 'i am context' }] },
+ { content: [{ text: 'i am more context' }] },
+ ]);
+ assert.deepEqual(result[0].content.at(-1), {
+ text: `${CONTEXT_PREFACE}- [0]: i am context\n- [1]: i am more context\n\n`,
+ metadata: { purpose: 'context' },
+ });
+ });
+
  it('should use a custom preface', async () => {
  const messages: MessageData[] = [
  { role: 'user', content: [{ text: 'first part' }] },