thread

firebase · Oct 9, 2024 · 81b1054 · 81b1054
1 parent f9a14a8
commit 81b1054
Show file tree

Hide file tree

Showing 10 changed files with 478 additions and 450 deletions.
diff --git a/js/ai/src/generate.ts b/js/ai/src/generate.ts
@@ -549,7 +549,7 @@ export async function generate<
     config: {
       ...resolvedModel.config,
       version: resolvedModel.version,
-      ...resolvedOptions.config
+      ...resolvedOptions.config,
     },
     output: resolvedOptions.output && {
       format: resolvedOptions.output.format,

diff --git a/js/ai/src/model.ts b/js/ai/src/model.ts
@@ -385,7 +385,7 @@ export interface ModelReference<CustomOptions extends z.ZodTypeAny> {
   configSchema?: CustomOptions;
   info?: ModelInfo;
   version?: string;
-  config?: z.infer<CustomOptions>,
+  config?: z.infer<CustomOptions>;
 }
 
 /** Cretes a model reference. */

diff --git a/js/genkit/src/genkit.ts b/js/genkit/src/genkit.ts
@@ -46,7 +46,7 @@ import {
   retrieve,
   RetrieverParams,
   ToolAction,
-  ToolConfig
+  ToolConfig,
 } from '@genkit-ai/ai';
 import {
   CallableFlow,
@@ -84,7 +84,6 @@ import {
 } from './model.js';
 import { lookupAction, Registry, runWithRegistry } from './registry.js';
 import {
-  BaseGenerateOptions,
   Environment,
   getCurrentSession,
   Session,
@@ -700,43 +699,42 @@ export class Genkit {
 
   chat<S extends z.ZodTypeAny = z.ZodTypeAny>(
     options?: SessionOptions<S>
-  ): Session<S>;
-
-  chat<S extends z.ZodTypeAny = z.ZodTypeAny>(
-    requestBase: BaseGenerateOptions,
-    options?: SessionOptions<S>
-  ): Session<S>;
-
-  chat<S extends z.ZodTypeAny = z.ZodTypeAny>(
-    requestBaseOrOpts?: SessionOptions<S> | BaseGenerateOptions,
-    maybeOptions?: SessionOptions<S>
   ): Session<S> {
-    // parse overloaded args
-    let baseGenerateOptions: BaseGenerateOptions | undefined = undefined;
-    let options: SessionOptions<S> | undefined = undefined;
-    if (maybeOptions !== undefined) {
-      options = maybeOptions;
-      baseGenerateOptions = requestBaseOrOpts as BaseGenerateOptions;
-    } else if (requestBaseOrOpts !== undefined) {
-      if (
-        (requestBaseOrOpts as SessionOptions<S>).state ||
-        (requestBaseOrOpts as SessionOptions<S>).store ||
-        (requestBaseOrOpts as SessionOptions<S>).schema
-      ) {
-        options = requestBaseOrOpts as SessionOptions<S>;
-      } else {
-        baseGenerateOptions = requestBaseOrOpts as BaseGenerateOptions;
+    return new Session(
+      this,
+      {
+        ...options,
+      },
+      {
+        sessionData: {
+          state: options?.state,
+        },
+        stateSchema: options?.stateSchema,
+        store: options?.store,
       }
-    }
+    );
+  }
 
+  async loadChat<S extends z.ZodTypeAny = z.ZodTypeAny>(
+    sessionId: string,
+    options: SessionOptions<S>
+  ): Promise<Session<S>> {
+    if (!options.store) {
+      throw new Error('options.store is required for loading chat sessions');
+    }
+    const sessionData = await options.store.get(sessionId);
+    if (!sessionData) {
+      throw new Error(`chat session ${sessionId} not found`);
+    }
     return new Session(
       this,
       {
-        ...baseGenerateOptions,
+        ...options,
       },
       {
-        state: options?.state,
-        schema: options?.schema,
+        id: sessionId,
+        sessionData,
+        stateSchema: options?.stateSchema,
         store: options?.store,
       }
     );