Fix bugs related to combining tools and output conformance (#542)

js/ai/src/generate.ts

@@ -635,6 +635,9 @@ export async function generate<
  if (resolvedOptions.output?.schema || resolvedOptions.output?.jsonSchema) {
  // find a candidate with valid output schema
  const candidateErrors = response.candidates.map((c) => {
+ // don't validate messages that have no text or data
+ if (c.text() === '' && c.data() === null) return null;
+
  try {
  parseSchema(c.output(), {
  jsonSchema: resolvedOptions.output?.jsonSchema,

js/ai/src/model/middleware.ts

@@ -15,7 +15,7 @@
  */
 
 import { Document } from '../document.js';
-import { ModelInfo, ModelMiddleware, Part } from '../model.js';
+import { MessageData, ModelInfo, ModelMiddleware, Part } from '../model.js';
 
 /**
  * Preprocess a GenerateRequest to download referenced http(s) media URLs and
@@ -117,9 +117,18 @@ export function validateSupport(options: {
  };
 }
 
+function lastUserMessage(messages: MessageData[]) {
+ for (let i = messages.length - 1; i >= 0; i--) {
+ if (messages[i].role === 'user') {
+ return messages[i];
+ }
+ }
+}
+
 export function conformOutput(): ModelMiddleware {
  return async (req, next) => {
- const lastMessage = req.messages.at(-1)!;
+ const lastMessage = lastUserMessage(req.messages);
+ if (!lastMessage) return next(req);
  const outputPartIndex = lastMessage.content.findIndex(
  (p) => p.metadata?.purpose === 'output'
  );

js/ai/tests/model/middleware_test.ts

@@ -170,17 +170,21 @@ describe('conformOutput (default middleware)', () => {
  const response = await echoModel(req);
  const treq = response.candidates[0].message.content[0]
  .data as GenerateRequest;
+
+ const lastUserMessage = treq.messages
+ .reverse()
+ .find((m) => m.role === 'user');
  if (
- treq.messages
-  .at(-1)!
- ?.content.filter((p) => p.metadata?.purpose === 'output').length > 1
+ lastUserMessage &&
+ lastUserMessage.content.filter((p) => p.metadata?.purpose === 'output')
+ .length > 1
  ) {
  throw new Error('too many output parts');
  }
 
- return treq.messages
- .at(-1)
-  ?.content.find((p) => p.metadata?.purpose === 'output')!;
+ return lastUserMessage?.content.find(
+ (p) => p.metadata?.purpose === 'output'
+ )!;
  }
 
  it('adds output instructions to the last message', async () => {
@@ -199,6 +203,51 @@ describe('conformOutput (default middleware)', () => {
  );
  });
 
+ it('adds output to the last message with "user" role', async () => {
+ const part = await testRequest({
+ messages: [
+ {
+ content: [
+ {
+ text: 'First message.',
+ },
+ ],
+ role: 'user',
+ },
+ {
+ content: [
+ {
+ toolRequest: {
+ name: 'localRestaurant',
+ input: {
+ location: 'wtf',
+ },
+ },
+ },
+ ],
+ role: 'model',
+ },
+ {
+ content: [
+ {
+ toolResponse: {
+ name: 'localRestaurant',
+ output: 'McDonalds',
+ },
+ },
+ ],
+ role: 'tool',
+ },
+ ],
+ output: { format: 'json', schema },
+ });
+
+ assert(
+ part?.text?.includes(JSON.stringify(schema)),
+ "schema wasn't found in output part"
+ );
+ });
+
  it('does not add output instructions if already provided', async () => {
  const part = await testRequest({
  messages: [

js/testapps/flow-simple-ai/src/index.ts

@@ -215,15 +215,16 @@ export const jokeWithToolsFlow = defineFlow(
  modelName: z.enum([geminiPro.name, googleGeminiPro.name]),
  subject: z.string(),
  }),
- outputSchema: z.string(),
+ outputSchema: z.object({ model: z.string(), joke: z.string() }),
  },
  async (input) => {
  const llmResponse = await generate({
  model: input.modelName,
  tools,
+ output: { schema: z.object({ joke: z.string() }) },
  prompt: `Tell a joke about ${input.subject}.`,
  });
- return `From ${input.modelName}: ${llmResponse.text()}`;
+ return { ...llmResponse.output()!, model: input.modelName };
  }
 );