fix: fixed prompt file descriptions and namespacing to unblock use as…

… tool (#1215)
firebase · Nov 11, 2024 · 4abb4b6 · 4abb4b6
1 parent c77d14e
commit 4abb4b6
Show file tree

Hide file tree

Showing 15 changed files with 511 additions and 103 deletions.
diff --git a/js/ai/package.json b/js/ai/package.json
@@ -32,7 +32,8 @@
     "colorette": "^2.0.20",
     "json5": "^2.2.3",
     "node-fetch": "^3.3.2",
-    "partial-json": "^0.1.7"
+    "partial-json": "^0.1.7",
+    "uuid": "^10.0.0"
   },
   "devDependencies": {
     "npm-run-all": "^4.1.5",
@@ -103,6 +104,18 @@
       "import": "./lib/reranker.mjs",
       "default": "./lib/reranker.js"
     },
+    "./chat": {
+      "types": "./lib/chat.d.ts",
+      "require": "./lib/chat.js",
+      "import": "./lib/chat.mjs",
+      "default": "./lib/chat.js"
+    },
+    "./session": {
+      "types": "./lib/session.d.ts",
+      "require": "./lib/session.js",
+      "import": "./lib/session.mjs",
+      "default": "./lib/session.js"
+    },
     "./formats": {
       "types": "./lib/formats/index.d.ts",
       "require": "./lib/formats/index.js",
@@ -135,6 +148,12 @@
       ],
       "reranker": [
         "lib/reranker"
+      ],
+      "chat": [
+        "lib/chat"
+      ],
+      "session": [
+        "lib/session"
       ]
     }
   }

diff --git a/js/genkit/src/chat.ts → js/ai/src/chat.ts b/js/genkit/src/chat.ts → js/ai/src/chat.ts
@@ -14,24 +14,25 @@
  * limitations under the License.
  */
 
+import { z } from '@genkit-ai/core';
+import { runInNewSpan } from '@genkit-ai/core/tracing';
 import {
+  generate,
   GenerateOptions,
   GenerateResponse,
+  generateStream,
   GenerateStreamOptions,
   GenerateStreamResponse,
   GenerationCommonConfigSchema,
   MessageData,
   Part,
-} from '@genkit-ai/ai';
-import { z } from '@genkit-ai/core';
-import { Genkit } from './genkit';
+} from './index.js';
 import {
   BaseGenerateOptions,
+  runWithSession,
   Session,
   SessionStore,
-  runWithSession,
 } from './session';
-import { runInNewSpan } from './tracing';
 
 export const MAIN_THREAD = 'main';
 
@@ -105,10 +106,21 @@ export class Chat {
         }
         requestBase.messages = [...(requestBase.messages ?? []), promptMessage];
       }
-      requestBase.messages = [
-        ...(options.messages ?? []),
-        ...(requestBase.messages ?? []),
-      ];
+      if (hasPreamble(requestBase.messages)) {
+        requestBase.messages = [
+          // if request base contains a preamble, always put it first
+          ...(getPreamble(requestBase.messages) ?? []),
+          // strip out the preamble from history
+          ...(stripPreamble(options.messages) ?? []),
+          // add whatever non-preamble remains from request
+          ...(stripPreamble(requestBase.messages) ?? []),
+        ];
+      } else {
+        requestBase.messages = [
+          ...(options.messages ?? []),
+          ...(requestBase.messages ?? []),
+        ];
+      }
       this._messages = requestBase.messages;
       return requestBase;
     });
@@ -145,7 +157,7 @@ export class Chat {
           messages: this.messages,
           ...resolvedOptions,
         };
-        let response = await this.genkit.generate({
+        let response = await generate(this.session.registry, {
           ...request,
           streamingCallback,
         });
@@ -185,11 +197,14 @@ export class Chat {
           resolvedOptions = options as GenerateStreamOptions<O, CustomOptions>;
         }
 
-        const { response, stream } = await this.genkit.generateStream({
-          ...(await this.requestBase),
-          messages: this.messages,
-          ...resolvedOptions,
-        });
+        const { response, stream } = await generateStream(
+          this.session.registry,
+          {
+            ...(await this.requestBase),
+            messages: this.messages,
+            ...resolvedOptions,
+          }
+        );
 
         return {
           response: response.finally(async () => {
@@ -208,10 +223,6 @@ export class Chat {
     );
   }
 
-  private get genkit(): Genkit {
-    return this.session.genkit;
-  }
-
   get messages(): MessageData[] {
     return this._messages ?? [];
   }
@@ -221,3 +232,15 @@ export class Chat {
     await this.session.updateMessages(this.threadName, messages);
   }
 }
+
+function hasPreamble(msgs?: MessageData[]) {
+  return !!msgs?.find((m) => m.metadata?.preamble);
+}
+
+function getPreamble(msgs?: MessageData[]) {
+  return msgs?.filter((m) => m.metadata?.preamble);
+}
+
+function stripPreamble(msgs?: MessageData[]) {
+  return msgs?.filter((m) => !m.metadata?.preamble);
+}
diff --git a/js/ai/src/model.ts b/js/ai/src/model.ts
@@ -496,6 +496,9 @@ export async function resolveModel<C extends z.ZodTypeAny = z.ZodTypeAny>(
   let out: ResolvedModel<C>;
   let modelId: string;
 
+  if (!model) {
+    model = await registry.lookupValue('defaultModel', 'defaultModel');
+  }
   if (!model) {
     throw new GenkitError({
       status: 'INVALID_ARGUMENT',

diff --git a/js/genkit/src/session.ts → js/ai/src/session.ts b/js/genkit/src/session.ts → js/ai/src/session.ts
@@ -14,19 +14,19 @@
  * limitations under the License.
  */
 
+import { z } from '@genkit-ai/core';
+import { Registry } from '@genkit-ai/core/registry';
+import { AsyncLocalStorage } from 'node:async_hooks';
+import { v4 as uuidv4 } from 'uuid';
+import { Chat, ChatOptions, MAIN_THREAD, PromptRenderOptions } from './chat';
 import {
   ExecutablePrompt,
   GenerateOptions,
   Message,
   MessageData,
   isExecutablePrompt,
   tagAsPreamble,
-} from '@genkit-ai/ai';
-import { z } from '@genkit-ai/core';
-import { AsyncLocalStorage } from 'node:async_hooks';
-import { v4 as uuidv4 } from 'uuid';
-import { Chat, ChatOptions, MAIN_THREAD, PromptRenderOptions } from './chat';
-import { Genkit } from './genkit';
+} from './index.js';
 
 export type BaseGenerateOptions<
   O extends z.ZodTypeAny = z.ZodTypeAny,
@@ -60,7 +60,7 @@ export class Session<S = any> {
   private store: SessionStore<S>;
 
   constructor(
-    readonly genkit: Genkit,
+    readonly registry: Registry,
     options?: {
       id?: string;
       stateSchema?: S;
@@ -82,10 +82,6 @@ export class Session<S = any> {
   }
 
   get state(): S | undefined {
-    // We always get state from the parent. Parent session is the source of truth.
-    if (this.genkit instanceof Session) {
-      return this.genkit.state;
-    }
     return this.sessionData!.state;
   }
 
@@ -228,6 +224,8 @@ export class Session<S = any> {
         requestBase = preamble
           .render({
             input: renderOptions?.input,
+            model: (renderOptions as BaseGenerateOptions)?.model,
+            config: (renderOptions as BaseGenerateOptions)?.config,
           })
           .then((rb) => {
             return {

diff --git a/js/core/src/action.ts b/js/core/src/action.ts
@@ -28,8 +28,6 @@ import {
 export { Status, StatusCodes, StatusSchema } from './statusTypes.js';
 export { JSONSchema7 };
 
-export const GENKIT_SESSION_STATE_INPUT_KEY = '__genkit__sessionState';
-
 export interface ActionMetadata<
   I extends z.ZodTypeAny,
   O extends z.ZodTypeAny,
@@ -122,12 +120,6 @@ export function action<
       ? config.name
       : `${config.name.pluginId}/${config.name.actionId}`;
   const actionFn = async (input: I) => {
-    let sessionStateData: Record<string, any> | undefined = undefined;
-    if (input?.hasOwnProperty(GENKIT_SESSION_STATE_INPUT_KEY)) {
-      sessionStateData = input[GENKIT_SESSION_STATE_INPUT_KEY];
-      input = { ...input };
-      delete input[GENKIT_SESSION_STATE_INPUT_KEY];
-    }
     input = parseSchema(input, {
       schema: config.inputSchema,
       jsonSchema: config.inputJsonSchema,
@@ -143,9 +135,6 @@ export function action<
         metadata.name = actionName;
         metadata.input = input;
 
-        if (sessionStateData) {
-          input[GENKIT_SESSION_STATE_INPUT_KEY] = sessionStateData;
-        }
         const output = await fn(input);
 
         metadata.output = JSON.stringify(output);

diff --git a/js/genkit/package.json b/js/genkit/package.json
@@ -29,7 +29,8 @@
   "dependencies": {
     "@genkit-ai/core": "workspace:*",
     "@genkit-ai/ai": "workspace:*",
-    "@genkit-ai/dotprompt": "workspace:*"
+    "@genkit-ai/dotprompt": "workspace:*",
+    "uuid": "^10.0.0"
   },
   "devDependencies": {
     "@types/express": "^4.17.21",
@@ -39,8 +40,7 @@
     "tsup": "^8.0.2",
     "typescript": "^4.9.0",
     "tsx": "^4.7.1",
-    "@types/body-parser": "^1.19.5",
-    "uuid": "^10.0.0"
+    "@types/body-parser": "^1.19.5"
   },
   "files": [
     "genkit-ui",