{"version":3,"sources":["../src/chat.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n *     http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport type { StreamingCallback, z } from '@genkit-ai/core';\nimport { Channel } from '@genkit-ai/core/async';\nimport {\n  ATTR_PREFIX,\n  SPAN_TYPE_ATTR,\n  runInNewSpan,\n} from '@genkit-ai/core/tracing';\nimport {\n  generate,\n  type GenerateOptions,\n  type GenerateResponse,\n  type GenerateResponseChunk,\n  type GenerateStreamOptions,\n  type GenerateStreamResponse,\n} from './generate.js';\nimport {\n  GenerationCommonConfigSchema,\n  type MessageData,\n} from './model-types.js';\nimport { type Part } from './parts.js';\nimport {\n  runWithSession,\n  type BaseGenerateOptions,\n  type Session,\n  type SessionStore,\n} from './session.js';\n\nexport const MAIN_THREAD = 'main';\nexport const SESSION_ID_ATTR = `${ATTR_PREFIX}:sessionId`;\nexport const THREAD_NAME_ATTR = `${ATTR_PREFIX}:threadName`;\n\nexport type ChatGenerateOptions<\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,\n> = GenerateOptions<O, CustomOptions>;\n\nexport interface PromptRenderOptions<I> {\n  input?: I;\n}\n\nexport type ChatOptions<I = undefined, S = any> = (\n  | PromptRenderOptions<I>\n  | BaseGenerateOptions\n) & {\n  store?: SessionStore<S>;\n  sessionId?: string;\n};\n\n/**\n * Chat encapsulates a statful execution environment for chat.\n * Chat session executed within a session in this environment will have acesss to\n * session convesation history.\n *\n * ```ts\n * const ai = genkit({...});\n * const chat = ai.chat(); // create a Chat\n * let response = await chat.send('hi, my name is Genkit');\n * response = await chat.send('what is my name?'); // chat history aware conversation\n * ```\n */\nexport class Chat {\n  private requestBase?: Promise<BaseGenerateOptions>;\n  readonly sessionId: string;\n  private _messages?: MessageData[];\n  private threadName: string;\n\n  constructor(\n    readonly session: Session,\n    requestBase: Promise<BaseGenerateOptions>,\n    options: {\n      id: string;\n      thread: string;\n      messages?: MessageData[];\n    }\n  ) {\n    this.sessionId = options.id;\n    this.threadName = options.thread;\n    this.requestBase = requestBase?.then((rb) => {\n      const requestBase = { ...rb };\n      // this is handling dotprompt render case\n      if (requestBase && requestBase['prompt']) {\n        const basePrompt = requestBase['prompt'] as string | Part | Part[];\n        let promptMessage: MessageData;\n        if (typeof basePrompt === 'string') {\n          promptMessage = {\n            role: 'user',\n            content: [{ text: basePrompt }],\n          };\n        } else if (Array.isArray(basePrompt)) {\n          promptMessage = {\n            role: 'user',\n            content: basePrompt,\n          };\n        } else {\n          promptMessage = {\n            role: 'user',\n            content: [basePrompt],\n          };\n        }\n        requestBase.messages = [...(requestBase.messages ?? []), promptMessage];\n      }\n      if (hasPreamble(requestBase.messages)) {\n        requestBase.messages = [\n          // if request base contains a preamble, always put it first\n          ...(getPreamble(requestBase.messages) ?? []),\n          // strip out the preamble from history\n          ...(stripPreamble(options.messages) ?? []),\n          // add whatever non-preamble remains from request\n          ...(stripPreamble(requestBase.messages) ?? []),\n        ];\n      } else {\n        requestBase.messages = [\n          ...(options.messages ?? []),\n          ...(requestBase.messages ?? []),\n        ];\n      }\n      this._messages = requestBase.messages;\n      return requestBase;\n    });\n    this._messages = options.messages;\n  }\n\n  async send<\n    O extends z.ZodTypeAny = z.ZodTypeAny,\n    CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,\n  >(\n    options: string | Part[] | ChatGenerateOptions<O, CustomOptions>\n  ): Promise<GenerateResponse<z.infer<O>>> {\n    return runWithSession(this.session.registry, this.session, () =>\n      runInNewSpan(\n        this.session.registry,\n        {\n          metadata: {\n            name: 'send',\n          },\n          labels: {\n            [SPAN_TYPE_ATTR]: 'helper',\n            [SESSION_ID_ATTR]: this.session.id,\n            [THREAD_NAME_ATTR]: this.threadName,\n          },\n        },\n        async (metadata) => {\n          const resolvedOptions = resolveSendOptions(options);\n          let streamingCallback:\n            | StreamingCallback<GenerateResponseChunk>\n            | undefined = undefined;\n\n          if (resolvedOptions.onChunk || resolvedOptions.streamingCallback) {\n            streamingCallback =\n              resolvedOptions.onChunk ?? resolvedOptions.streamingCallback;\n          }\n          const request: GenerateOptions = {\n            ...(await this.requestBase),\n            messages: this.messages,\n            ...resolvedOptions,\n          };\n          metadata.input = resolvedOptions;\n          const response = await generate(this.session.registry, {\n            ...request,\n            onChunk: streamingCallback,\n          });\n          this.requestBase = Promise.resolve({\n            ...(await this.requestBase),\n            // these things may get changed by tools calling within generate.\n            tools: response?.request?.tools?.map((td) => td.name),\n            toolChoice: response?.request?.toolChoice,\n            config: response?.request?.config,\n          });\n          await this.updateMessages(response.messages);\n          metadata.output = JSON.stringify(response);\n          return response;\n        }\n      )\n    );\n  }\n\n  sendStream<\n    O extends z.ZodTypeAny = z.ZodTypeAny,\n    CustomOptions extends z.ZodTypeAny = typeof GenerationCommonConfigSchema,\n  >(\n    options: string | Part[] | GenerateStreamOptions<O, CustomOptions>\n  ): GenerateStreamResponse<z.infer<O>> {\n    const channel = new Channel<GenerateResponseChunk>();\n    const resolvedOptions = resolveSendOptions(options);\n\n    const sent = this.send({\n      ...resolvedOptions,\n      onChunk: (chunk) => channel.send(chunk),\n    });\n    sent.then(\n      () => channel.close(),\n      (err) => channel.error(err)\n    );\n\n    return {\n      response: sent,\n      stream: channel,\n    };\n  }\n\n  get messages(): MessageData[] {\n    return this._messages ?? [];\n  }\n\n  private async updateMessages(messages: MessageData[]): Promise<void> {\n    this._messages = messages;\n    await this.session.updateMessages(this.threadName, messages);\n  }\n}\n\nfunction hasPreamble(msgs?: MessageData[]) {\n  return !!msgs?.find((m) => m.metadata?.preamble);\n}\n\nfunction getPreamble(msgs?: MessageData[]) {\n  return msgs?.filter((m) => m.metadata?.preamble);\n}\n\nfunction stripPreamble(msgs?: MessageData[]) {\n  return msgs?.filter((m) => !m.metadata?.preamble);\n}\n\nfunction resolveSendOptions<\n  O extends z.ZodTypeAny,\n  CustomOptions extends z.ZodTypeAny,\n>(\n  options: string | Part[] | ChatGenerateOptions<O, CustomOptions>\n): ChatGenerateOptions<O, CustomOptions> {\n  let resolvedOptions: ChatGenerateOptions<O, CustomOptions>;\n\n  // string\n  if (typeof options === 'string') {\n    resolvedOptions = {\n      prompt: options,\n    } as ChatGenerateOptions<O, CustomOptions>;\n  } else if (Array.isArray(options)) {\n    // Part[]\n    resolvedOptions = {\n      prompt: options,\n    } as ChatGenerateOptions<O, CustomOptions>;\n  } else {\n    resolvedOptions = options as ChatGenerateOptions<O, CustomOptions>;\n  }\n  return resolvedOptions;\n}\n"],"mappings":"AAiBA,SAAS,eAAe;AACxB;AAAA,EACE;AAAA,EACA;AAAA,EACA;AAAA,OACK;AACP;AAAA,EACE;AAAA,OAMK;AAMP;AAAA,EACE;AAAA,OAIK;AAEA,MAAM,cAAc;AACpB,MAAM,kBAAkB,GAAG,WAAW;AACtC,MAAM,mBAAmB,GAAG,WAAW;AA+BvC,MAAM,KAAK;AAAA,EAMhB,YACW,SACT,aACA,SAKA;AAPS;AAQT,SAAK,YAAY,QAAQ;AACzB,SAAK,aAAa,QAAQ;AAC1B,SAAK,cAAc,aAAa,KAAK,CAAC,OAAO;AAC3C,YAAMA,eAAc,EAAE,GAAG,GAAG;AAE5B,UAAIA,gBAAeA,aAAY,QAAQ,GAAG;AACxC,cAAM,aAAaA,aAAY,QAAQ;AACvC,YAAI;AACJ,YAAI,OAAO,eAAe,UAAU;AAClC,0BAAgB;AAAA,YACd,MAAM;AAAA,YACN,SAAS,CAAC,EAAE,MAAM,WAAW,CAAC;AAAA,UAChC;AAAA,QACF,WAAW,MAAM,QAAQ,UAAU,GAAG;AACpC,0BAAgB;AAAA,YACd,MAAM;AAAA,YACN,SAAS;AAAA,UACX;AAAA,QACF,OAAO;AACL,0BAAgB;AAAA,YACd,MAAM;AAAA,YACN,SAAS,CAAC,UAAU;AAAA,UACtB;AAAA,QACF;AACA,QAAAA,aAAY,WAAW,CAAC,GAAIA,aAAY,YAAY,CAAC,GAAI,aAAa;AAAA,MACxE;AACA,UAAI,YAAYA,aAAY,QAAQ,GAAG;AACrC,QAAAA,aAAY,WAAW;AAAA;AAAA,UAErB,GAAI,YAAYA,aAAY,QAAQ,KAAK,CAAC;AAAA;AAAA,UAE1C,GAAI,cAAc,QAAQ,QAAQ,KAAK,CAAC;AAAA;AAAA,UAExC,GAAI,cAAcA,aAAY,QAAQ,KAAK,CAAC;AAAA,QAC9C;AAAA,MACF,OAAO;AACL,QAAAA,aAAY,WAAW;AAAA,UACrB,GAAI,QAAQ,YAAY,CAAC;AAAA,UACzB,GAAIA,aAAY,YAAY,CAAC;AAAA,QAC/B;AAAA,MACF;AACA,WAAK,YAAYA,aAAY;AAC7B,aAAOA;AAAA,IACT,CAAC;AACD,SAAK,YAAY,QAAQ;AAAA,EAC3B;AAAA,EA3DQ;AAAA,EACC;AAAA,EACD;AAAA,EACA;AAAA,EA0DR,MAAM,KAIJ,SACuC;AACvC,WAAO;AAAA,MAAe,KAAK,QAAQ;AAAA,MAAU,KAAK;AAAA,MAAS,MACzD;AAAA,QACE,KAAK,QAAQ;AAAA,QACb;AAAA,UACE,UAAU;AAAA,YACR,MAAM;AAAA,UACR;AAAA,UACA,QAAQ;AAAA,YACN,CAAC,cAAc,GAAG;AAAA,YAClB,CAAC,eAAe,GAAG,KAAK,QAAQ;AAAA,YAChC,CAAC,gBAAgB,GAAG,KAAK;AAAA,UAC3B;AAAA,QACF;AAAA,QACA,OAAO,aAAa;AAClB,gBAAM,kBAAkB,mBAAmB,OAAO;AAClD,cAAI,oBAEY;AAEhB,cAAI,gBAAgB,WAAW,gBAAgB,mBAAmB;AAChE,gCACE,gBAAgB,WAAW,gBAAgB;AAAA,UAC/C;AACA,gBAAM,UAA2B;AAAA,YAC/B,GAAI,MAAM,KAAK;AAAA,YACf,UAAU,KAAK;AAAA,YACf,GAAG;AAAA,UACL;AACA,mBAAS,QAAQ;AACjB,gBAAM,WAAW,MAAM,SAAS,KAAK,QAAQ,UAAU;AAAA,YACrD,GAAG;AAAA,YACH,SAAS;AAAA,UACX,CAAC;AACD,eAAK,cAAc,QAAQ,QAAQ;AAAA,YACjC,GAAI,MAAM,KAAK;AAAA;AAAA,YAEf,OAAO,UAAU,SAAS,OAAO,IAAI,CAAC,OAAO,GAAG,IAAI;AAAA,YACpD,YAAY,UAAU,SAAS;AAAA,YAC/B,QAAQ,UAAU,SAAS;AAAA,UAC7B,CAAC;AACD,gBAAM,KAAK,eAAe,SAAS,QAAQ;AAC3C,mBAAS,SAAS,KAAK,UAAU,QAAQ;AACzC,iBAAO;AAAA,QACT;AAAA,MACF;AAAA,IACF;AAAA,EACF;AAAA,EAEA,WAIE,SACoC;AACpC,UAAM,UAAU,IAAI,QAA+B;AACnD,UAAM,kBAAkB,mBAAmB,OAAO;AAElD,UAAM,OAAO,KAAK,KAAK;AAAA,MACrB,GAAG;AAAA,MACH,SAAS,CAAC,UAAU,QAAQ,KAAK,KAAK;AAAA,IACxC,CAAC;AACD,SAAK;AAAA,MACH,MAAM,QAAQ,MAAM;AAAA,MACpB,CAAC,QAAQ,QAAQ,MAAM,GAAG;AAAA,IAC5B;AAEA,WAAO;AAAA,MACL,UAAU;AAAA,MACV,QAAQ;AAAA,IACV;AAAA,EACF;AAAA,EAEA,IAAI,WAA0B;AAC5B,WAAO,KAAK,aAAa,CAAC;AAAA,EAC5B;AAAA,EAEA,MAAc,eAAe,UAAwC;AACnE,SAAK,YAAY;AACjB,UAAM,KAAK,QAAQ,eAAe,KAAK,YAAY,QAAQ;AAAA,EAC7D;AACF;AAEA,SAAS,YAAY,MAAsB;AACzC,SAAO,CAAC,CAAC,MAAM,KAAK,CAAC,MAAM,EAAE,UAAU,QAAQ;AACjD;AAEA,SAAS,YAAY,MAAsB;AACzC,SAAO,MAAM,OAAO,CAAC,MAAM,EAAE,UAAU,QAAQ;AACjD;AAEA,SAAS,cAAc,MAAsB;AAC3C,SAAO,MAAM,OAAO,CAAC,MAAM,CAAC,EAAE,UAAU,QAAQ;AAClD;AAEA,SAAS,mBAIP,SACuC;AACvC,MAAI;AAGJ,MAAI,OAAO,YAAY,UAAU;AAC/B,sBAAkB;AAAA,MAChB,QAAQ;AAAA,IACV;AAAA,EACF,WAAW,MAAM,QAAQ,OAAO,GAAG;AAEjC,sBAAkB;AAAA,MAChB,QAAQ;AAAA,IACV;AAAA,EACF,OAAO;AACL,sBAAkB;AAAA,EACpB;AACA,SAAO;AACT;","names":["requestBase"]}