Fixes issue with context not passing through generate() (#139)

Author: mbleigh

js/ai/src/generate.ts

@@ -410,6 +410,7 @@ export async function toGenerateRequest(
     messages,
     candidates: options.candidates,
     config: options.config,
+    context: options.context,
     tools: tools?.map((tool) => toToolDefinition(tool)) || [],
     output: {
       format:

js/ai/src/model/middleware.ts

@@ -198,13 +198,13 @@ const CONTEXT_ITEM_TEMPLATE = (
   out += d.text() + '\n';
   return out;
 };
+
 export function augmentWithContext(
   options?: AugmentWithContextOptions
 ): ModelMiddleware {
   const preface =
     typeof options?.preface === 'undefined' ? CONTEXT_PREFACE : options.preface;
   const itemTemplate = options?.itemTemplate || CONTEXT_ITEM_TEMPLATE;
-  const citationKey = options?.citationKey;
   return (req, next) => {
     // if there is no context in the request, no-op
     if (!req.context?.length) return next(req);

js/ai/tests/generate/generate_test.ts

@@ -260,6 +260,7 @@ describe('toGenerateRequest', () => {
         ],
         candidates: undefined,
         config: undefined,
+        context: undefined,
         tools: [],
         output: { format: 'text' },
       },
@@ -278,6 +279,7 @@ describe('toGenerateRequest', () => {
         ],
         candidates: undefined,
         config: undefined,
+        context: undefined,
         tools: [
           {
             name: 'tellAFunnyJoke',
@@ -313,6 +315,7 @@ describe('toGenerateRequest', () => {
         ],
         candidates: undefined,
         config: undefined,
+        context: undefined,
         tools: [
           {
             name: 'tellAFunnyJoke',
@@ -365,6 +368,7 @@ describe('toGenerateRequest', () => {
         ],
         candidates: undefined,
         config: undefined,
+        context: undefined,
         tools: [],
         output: { format: 'text' },
       },
@@ -387,6 +391,25 @@ describe('toGenerateRequest', () => {
         ],
         candidates: undefined,
         config: undefined,
+        context: undefined,
+        tools: [],
+        output: { format: 'text' },
+      },
+    },
+    {
+      should: 'pass context through to the model',
+      prompt: {
+        model: 'vertexai/gemini-1.0-pro',
+        prompt: 'Tell a joke with context.',
+        context: [{ content: [{ text: 'context here' }] }],
+      },
+      expectedOutput: {
+        messages: [
+          { content: [{ text: 'Tell a joke with context.' }], role: 'user' },
+        ],
+        candidates: undefined,
+        config: undefined,
+        context: [{ content: [{ text: 'context here' }] }],
         tools: [],
         output: { format: 'text' },
       },

js/pnpm-lock.yaml

@@ -22,6 +22,7 @@
     "@genkit-ai/google-cloud": "workspace:*",
     "@genkit-ai/googleai": "workspace:*",
     "@genkit-ai/vertexai": "workspace:*",
+    "@genkit-ai/dotprompt": "workspace:*",
     "@opentelemetry/sdk-trace-base": "^1.22.0",
     "firebase-admin": "^12.1.0",
     "zod": "^3.22.4"

js/samples/flow-simple-ai/package.json

@@ -0,0 +1,17 @@
+---
+model: vertexai/gemini-1.0-pro
+input:
+  schema:
+    question: string
+output:
+  format: json
+  schema:
+    answer: string, the answer to the question
+    id: string, the selected id of the saying
+    reasoning: string, why the saying applies to the question
+---
+
+You are a mystic wisdom bot designed to help people with their problems. Use the provided
+sayings to answer the question.
+
+Question: {{question}}

js/samples/flow-simple-ai/prompts/dotpromptContext.prompt

@@ -17,6 +17,7 @@
 import { generate, generateStream, retrieve } from '@genkit-ai/ai';
 import { defineTool } from '@genkit-ai/ai/tool';
 import { configureGenkit } from '@genkit-ai/core';
+import { dotprompt, prompt } from '@genkit-ai/dotprompt';
 import { defineFirestoreRetriever, firebase } from '@genkit-ai/firebase';
 import { defineFlow, run } from '@genkit-ai/flow';
 import { googleCloud } from '@genkit-ai/google-cloud';
@@ -48,6 +49,7 @@ configureGenkit({
         metricExportIntervalMillis: 5_000,
       },
     }),
+    dotprompt(),
   ],
   flowStateStore: 'firebase',
   traceStore: 'firebase',
@@ -230,3 +232,45 @@ Available Options:\n- ${docs.map((d) => `${d.metadata!.name}: ${d.text()}`).join
     return result.text();
   }
 );
+
+export const dotpromptContext = defineFlow(
+  {
+    name: 'dotpromptContext',
+    inputSchema: z.string(),
+    outputSchema: z.object({
+      answer: z.string(),
+      id: z.string(),
+      reasoning: z.string(),
+    }),
+  },
+  async (question: string) => {
+    const docs = [
+      {
+        content: [{ text: 'an apple a day keeps the doctor away' }],
+        metadata: { id: 'apple' },
+      },
+      {
+        content: [
+          { text: 'those who live in glass houses should not throw stones' },
+        ],
+        metadata: { id: 'stone' },
+      },
+      {
+        content: [
+          {
+            text: "if you don't have anything nice to say, don't say anything at all",
+          },
+        ],
+        metadata: { id: 'nice' },
+      },
+    ];
+
+    const result = await (
+      await prompt('dotpromptContext')
+    ).generate({
+      input: { question: question },
+      context: docs,
+    });
+    return result.output() as any;
+  }
+);