feat: updated syntax for creating preamble/agent chat (#1209)

js/ai/src/index.ts

@@ -73,6 +73,7 @@ export {
 } from './model.js';
 export {
   definePrompt,
+  isExecutablePrompt,
   renderPrompt,
   type ExecutablePrompt,
   type PromptAction,

js/ai/src/prompt.ts

@@ -211,3 +211,11 @@ export async function renderPrompt<
     tools: rendered.tools || [],
   } as GenerateOptions<O, CustomOptions>;
 }
+
+export function isExecutablePrompt(obj: any): boolean {
+  return (
+    !!(obj as ExecutablePrompt)?.render &&
+    !!(obj as ExecutablePrompt)?.asTool &&
+    !!(obj as ExecutablePrompt)?.stream
+  );
+}

js/genkit/src/chat.ts

@@ -15,7 +15,6 @@
  */
 
 import {
-  ExecutablePrompt,
   GenerateOptions,
   GenerateResponse,
   GenerateStreamOptions,
@@ -42,7 +41,6 @@ export type ChatGenerateOptions<
 > = GenerateOptions<O, CustomOptions>;
 
 export interface PromptRenderOptions<I> {
-  preamble: ExecutablePrompt<I>;
   input?: I;
 }
 

js/genkit/src/genkit.ts

@@ -902,8 +902,53 @@ export class Genkit {
    * response = await chat.send('another one')
    * ```
    */
-  chat<I>(options?: ChatOptions<I>): Chat {
+  chat<I>(options?: ChatOptions<I>): Chat;
+
+  /**
+   * Create a chat session with the provided preabmle.
+   *
+   * ```ts
+   * const triageAgent = ai.definePrompt({
+   *   system: 'help the user triage a problem',
+   * })
+   * const chat = ai.chat(triageAgent)
+   * const { text } = await chat.send('my phone feels hot');
+   * ```
+   */
+  chat<I>(preamble: ExecutablePrompt<I>, options?: ChatOptions<I>): Chat;
+
+  /**
+   * Create a chat session with the provided options.
+   *
+   * ```ts
+   * const chat = ai.chat({
+   *   system: 'talk like a pirate',
+   * })
+   * let response = await chat.send('tell me a joke')
+   * response = await chat.send('another one')
+   * ```
+   */
+  chat<I>(
+    preambleOrOptions?: ChatOptions<I> | ExecutablePrompt<I>,
+    maybeOptions?: ChatOptions<I>
+  ): Chat {
+    let options: ChatOptions<I> | undefined;
+    let preamble: ExecutablePrompt<I> | undefined;
+    if (maybeOptions) {
+      preamble = preambleOrOptions as ExecutablePrompt<I>;
+      options = maybeOptions;
+    } else if (preambleOrOptions) {
+      if ((preambleOrOptions as ExecutablePrompt<I>)?.render) {
+        preamble = preambleOrOptions as ExecutablePrompt<I>;
+      } else {
+        options = preambleOrOptions as ChatOptions<I>;
+      }
+    }
+
     const session = this.createSession();
+    if (preamble) {
+      return session.chat(preamble, options);
+    }
     return session.chat(options);
   }
 

js/genkit/src/session.ts

@@ -15,9 +15,11 @@
  */
 
 import {
+  ExecutablePrompt,
   GenerateOptions,
   Message,
   MessageData,
+  isExecutablePrompt,
   tagAsPreamble,
 } from '@genkit-ai/ai';
 import { z } from '@genkit-ai/core';
@@ -125,54 +127,107 @@ export class Session<S = any> {
    * Create a chat session with the provided options.
    *
    * ```ts
-   * const chat = ai.chat({
+   * const session = ai.createSession({});
+   * const chat = session.chat({
    *   system: 'talk like a pirate',
    * })
-   * let response = await chat.send('tell me a joke')
-   * response = await chat.send('another one')
+   * let response = await chat.send('tell me a joke');
+   * response = await chat.send('another one');
    * ```
    */
   chat<I>(options?: ChatOptions<I, S>): Chat;
 
   /**
-   * Craete a separaete chat conversation ("thread") within the same session state.
+   * Create a chat session with the provided preamble.
    *
    * ```ts
-   * const lawyerChat = ai.chat('lawyerThread', {
-   *   system: 'talk like a lawyer',
+   * const triageAgent = ai.definePrompt({
+   *   system: 'help the user triage a problem',
    * })
-   * const pirateChat = ai.chat('pirateThread', {
+   * const session = ai.createSession({});
+   * const chat = session.chat(triageAgent);
+   * const { text } = await chat.send('my phone feels hot');
+   * ```
+   */
+  chat<I>(preamble: ExecutablePrompt<I>, options?: ChatOptions<I, S>): Chat;
+
+  /**
+   * Craete a separate chat conversation ("thread") within the given preamble.
+   *
+   * ```ts
+   * const session = ai.createSession({});
+   * const lawyerChat = session.chat('lawyerThread', {
+   *   system: 'talk like a lawyer',
+   * });
+   * const pirateChat = session.chat('pirateThread', {
    *   system: 'talk like a pirate',
-   * })
-   * await lawyerChat.send('tell me a joke')
-   * await pirateChat.send('tell me a joke')
+   * });
+   * await lawyerChat.send('tell me a joke');
+   * await pirateChat.send('tell me a joke');
+   * ```
+   */
+  chat<I>(
+    threadName: string,
+    preamble: ExecutablePrompt<I>,
+    options?: ChatOptions<I, S>
+  ): Chat;
+
+  /**
+   * Craete a separate chat conversation ("thread").
+   *
+   * ```ts
+   * const session = ai.createSession({});
+   * const lawyerChat = session.chat('lawyerThread', {
+   *   system: 'talk like a lawyer',
+   * });
+   * const pirateChat = session.chat('pirateThread', {
+   *   system: 'talk like a pirate',
+   * });
+   * await lawyerChat.send('tell me a joke');
+   * await pirateChat.send('tell me a joke');
    * ```
    */
   chat<I>(threadName: string, options?: ChatOptions<I, S>): Chat;
 
   chat<I>(
-    optionsOrThreadName?: ChatOptions<I, S> | string,
+    optionsOrPreambleOrThreadName?:
+      | ChatOptions<I, S>
+      | string
+      | ExecutablePrompt<I>,
+    maybeOptionsOrPreamble?: ChatOptions<I, S> | ExecutablePrompt<I>,
     maybeOptions?: ChatOptions<I, S>
   ): Chat {
     return runWithSession(this, () => {
       let options: ChatOptions<S> | undefined;
       let threadName = MAIN_THREAD;
-      if (maybeOptions) {
-        threadName = optionsOrThreadName as string;
-        options = maybeOptions as ChatOptions<S>;
-      } else if (optionsOrThreadName) {
-        if (typeof optionsOrThreadName === 'string') {
-          threadName = optionsOrThreadName as string;
+      let preamble: ExecutablePrompt<I> | undefined;
+
+      if (optionsOrPreambleOrThreadName) {
+        if (typeof optionsOrPreambleOrThreadName === 'string') {
+          threadName = optionsOrPreambleOrThreadName as string;
+        } else if (isExecutablePrompt(optionsOrPreambleOrThreadName)) {
+          preamble = optionsOrPreambleOrThreadName as ExecutablePrompt<I>;
         } else {
-          options = optionsOrThreadName as ChatOptions<S>;
+          options = optionsOrPreambleOrThreadName as ChatOptions<I, S>;
         }
       }
+      if (maybeOptionsOrPreamble) {
+        if (isExecutablePrompt(maybeOptionsOrPreamble)) {
+          preamble = maybeOptionsOrPreamble as ExecutablePrompt<I>;
+        } else {
+          options = maybeOptionsOrPreamble as ChatOptions<I, S>;
+        }
+      }
+      if (maybeOptions) {
+        options = maybeOptions as ChatOptions<I, S>;
+      }
+
       let requestBase: Promise<BaseGenerateOptions>;
-      if (!!(options as PromptRenderOptions<I>)?.preamble?.render) {
+      if (preamble) {
         const renderOptions = options as PromptRenderOptions<I>;
-        requestBase = renderOptions.preamble
+        requestBase = preamble
           .render({
-            input: renderOptions.input,
+            input: renderOptions?.input,
           })
           .then((rb) => {
             return {

js/genkit/tests/chat_test.ts

@@ -120,12 +120,25 @@ describe('chat', () => {
   });
 
   it('can start chat from a prompt', async () => {
+    const preamble = ai.definePrompt(
+      { name: 'hi', config: { version: 'abc' } },
+      'hi from template'
+    );
+    const session = await ai.chat(preamble);
+    const response = await session.send('send it');
+
+    assert.strictEqual(
+      response.text,
+      'Echo: hi from template,send it; config: {"version":"abc"}'
+    );
+  });
+
+  it('can start chat from a prompt with input', async () => {
     const preamble = ai.definePrompt(
       { name: 'hi', config: { version: 'abc' } },
       'hi {{ name }} from template'
     );
-    const session = await ai.chat({
-      preamble,
+    const session = await ai.chat(preamble, {
       input: { name: 'Genkit' },
     });
     const response = await session.send('send it');
@@ -156,7 +169,7 @@ describe('chat', () => {
   });
 });
 
-describe('preabmle', () => {
+describe('preamble', () => {
   let ai: Genkit;
   let pm: ProgrammableModel;
 
@@ -207,9 +220,7 @@ describe('preabmle', () => {
       };
     };
 
-    const session = ai.chat({
-      preamble: agentA,
-    });
+    const session = ai.chat(agentA);
     let { text } = await session.send('hi');
     assert.strictEqual(text, 'hi from agent a');
     assert.deepStrictEqual(pm.lastRequest, {