From ac5a4bca8cedf648324f1134c05ed2d6bcc11e55 Mon Sep 17 00:00:00 2001 From: Michael Bleigh Date: Mon, 13 May 2024 07:01:32 -0700 Subject: [PATCH] Fixes issue with context not passing through generate() (#139) --- js/ai/src/generate.ts | 1 + js/ai/src/model/middleware.ts | 2 +- js/ai/tests/generate/generate_test.ts | 23 ++++++++++ js/pnpm-lock.yaml | 3 ++ js/samples/flow-simple-ai/package.json | 1 + .../prompts/dotpromptContext.prompt | 17 +++++++ js/samples/flow-simple-ai/src/index.ts | 44 +++++++++++++++++++ 7 files changed, 90 insertions(+), 1 deletion(-) create mode 100644 js/samples/flow-simple-ai/prompts/dotpromptContext.prompt diff --git a/js/ai/src/generate.ts b/js/ai/src/generate.ts index 53523106a..065624bb8 100755 --- a/js/ai/src/generate.ts +++ b/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: diff --git a/js/ai/src/model/middleware.ts b/js/ai/src/model/middleware.ts index f15c79456..b1edb369a 100644 --- a/js/ai/src/model/middleware.ts +++ b/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); diff --git a/js/ai/tests/generate/generate_test.ts b/js/ai/tests/generate/generate_test.ts index 430a8c448..adb3edb6d 100644 --- a/js/ai/tests/generate/generate_test.ts +++ b/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' }, }, diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml index a8d204cfa..c9da7305c 100644 --- a/js/pnpm-lock.yaml +++ b/js/pnpm-lock.yaml @@ -989,6 +989,9 @@ importers: '@genkit-ai/core': specifier: workspace:^ version: link:../../core + '@genkit-ai/dotprompt': + specifier: workspace:* + version: link:../../plugins/dotprompt '@genkit-ai/firebase': specifier: workspace:* version: link:../../plugins/firebase diff --git a/js/samples/flow-simple-ai/package.json b/js/samples/flow-simple-ai/package.json index dda05338c..76780cf99 100644 --- a/js/samples/flow-simple-ai/package.json +++ b/js/samples/flow-simple-ai/package.json @@ -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" diff --git a/js/samples/flow-simple-ai/prompts/dotpromptContext.prompt b/js/samples/flow-simple-ai/prompts/dotpromptContext.prompt new file mode 100644 index 000000000..809541f19 --- /dev/null +++ b/js/samples/flow-simple-ai/prompts/dotpromptContext.prompt @@ -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}} \ No newline at end of file diff --git a/js/samples/flow-simple-ai/src/index.ts b/js/samples/flow-simple-ai/src/index.ts index 97461ce57..9f0bcea7e 100644 --- a/js/samples/flow-simple-ai/src/index.ts +++ b/js/samples/flow-simple-ai/src/index.ts @@ -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; + } +);