{"version":3,"sources":["../src/session.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 { getAsyncContext, type z } from '@genkit-ai/core';\nimport type { Registry } from '@genkit-ai/core/registry';\nimport { v4 as uuidv4 } from 'uuid';\nimport {\n  Chat,\n  MAIN_THREAD,\n  type ChatOptions,\n  type PromptRenderOptions,\n} from './chat.js';\nimport {\n  Message,\n  isExecutablePrompt,\n  tagAsPreamble,\n  type ExecutablePrompt,\n  type GenerateOptions,\n  type MessageData,\n  type PromptGenerateOptions,\n} from './index.js';\n\nexport type BaseGenerateOptions<\n  O extends z.ZodTypeAny = z.ZodTypeAny,\n  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,\n> = Omit<GenerateOptions<O, CustomOptions>, 'prompt'>;\n\nexport interface SessionOptions<S = any> {\n  /** Session store implementation for persisting the session state. */\n  store?: SessionStore<S>;\n  /** Initial state of the session.  */\n  initialState?: S;\n  /** Custom session Id. */\n  sessionId?: string;\n}\n\n/**\n * Session encapsulates a statful execution environment for chat.\n * Chat session executed within a session in this environment will have acesss to\n * session session convesation history.\n *\n * ```ts\n * const ai = genkit({...});\n * const chat = ai.chat(); // create a Session\n * let response = await chat.send('hi'); // session/history aware conversation\n * response = await chat.send('tell me a story');\n * ```\n */\nexport class Session<S = any> {\n  readonly id: string;\n  private sessionData?: SessionData<S>;\n  private store: SessionStore<S>;\n\n  constructor(\n    readonly registry: Registry,\n    options?: {\n      id?: string;\n      stateSchema?: S;\n      sessionData?: SessionData<S>;\n      store?: SessionStore<S>;\n    }\n  ) {\n    this.id = options?.id ?? uuidv4();\n    this.sessionData = options?.sessionData ?? {\n      id: this.id,\n    };\n    if (!this.sessionData) {\n      this.sessionData = { id: this.id };\n    }\n    if (!this.sessionData.threads) {\n      this.sessionData!.threads = {};\n    }\n    this.store = options?.store ?? new InMemorySessionStore<S>();\n  }\n\n  get state(): S | undefined {\n    return this.sessionData!.state;\n  }\n\n  /**\n   * Update session state data.\n   */\n  async updateState(data: S): Promise<void> {\n    let sessionData = this.sessionData;\n    if (!sessionData) {\n      sessionData = {} as SessionData<S>;\n    }\n    sessionData.state = data;\n    this.sessionData = sessionData;\n\n    await this.store.save(this.id, sessionData);\n  }\n\n  /**\n   * Update messages for a given thread.\n   */\n  async updateMessages(thread: string, messages: MessageData[]): Promise<void> {\n    let sessionData = this.sessionData;\n    if (!sessionData) {\n      sessionData = {} as SessionData<S>;\n    }\n    if (!sessionData.threads) {\n      sessionData.threads = {};\n    }\n    sessionData.threads[thread] = messages.map((m: any) =>\n      m.toJSON ? m.toJSON() : m\n    );\n    this.sessionData = sessionData;\n\n    await this.store.save(this.id, sessionData);\n  }\n\n  /**\n   * Create a chat session with the provided options.\n   *\n   * ```ts\n   * const session = ai.createSession({});\n   * const chat = session.chat({\n   *   system: 'talk like a pirate',\n   * })\n   * let response = await chat.send('tell me a joke');\n   * response = await chat.send('another one');\n   * ```\n   */\n  chat<I>(options?: ChatOptions<I, S>): Chat;\n\n  /**\n   * Create a chat session with the provided preamble.\n   *\n   * ```ts\n   * const triageAgent = ai.definePrompt({\n   *   system: 'help the user triage a problem',\n   * })\n   * const session = ai.createSession({});\n   * const chat = session.chat(triageAgent);\n   * const { text } = await chat.send('my phone feels hot');\n   * ```\n   */\n  chat<I>(preamble: ExecutablePrompt<I>, options?: ChatOptions<I, S>): Chat;\n\n  /**\n   * Craete a separate chat conversation (\"thread\") within the given preamble.\n   *\n   * ```ts\n   * const session = ai.createSession({});\n   * const lawyerChat = session.chat('lawyerThread', {\n   *   system: 'talk like a lawyer',\n   * });\n   * const pirateChat = session.chat('pirateThread', {\n   *   system: 'talk like a pirate',\n   * });\n   * await lawyerChat.send('tell me a joke');\n   * await pirateChat.send('tell me a joke');\n   * ```\n   */\n  chat<I>(\n    threadName: string,\n    preamble: ExecutablePrompt<I>,\n    options?: ChatOptions<I, S>\n  ): Chat;\n\n  /**\n   * Craete a separate chat conversation (\"thread\").\n   *\n   * ```ts\n   * const session = ai.createSession({});\n   * const lawyerChat = session.chat('lawyerThread', {\n   *   system: 'talk like a lawyer',\n   * });\n   * const pirateChat = session.chat('pirateThread', {\n   *   system: 'talk like a pirate',\n   * });\n   * await lawyerChat.send('tell me a joke');\n   * await pirateChat.send('tell me a joke');\n   * ```\n   */\n  chat<I>(threadName: string, options?: ChatOptions<I, S>): Chat;\n\n  chat<I>(\n    optionsOrPreambleOrThreadName?:\n      | ChatOptions<I, S>\n      | string\n      | ExecutablePrompt<I>,\n    maybeOptionsOrPreamble?: ChatOptions<I, S> | ExecutablePrompt<I>,\n    maybeOptions?: ChatOptions<I, S>\n  ): Chat {\n    return runWithSession(this.registry, this, () => {\n      let options: ChatOptions<S> | undefined;\n      let threadName = MAIN_THREAD;\n      let preamble: ExecutablePrompt<I> | undefined;\n\n      if (optionsOrPreambleOrThreadName) {\n        if (typeof optionsOrPreambleOrThreadName === 'string') {\n          threadName = optionsOrPreambleOrThreadName as string;\n        } else if (isExecutablePrompt(optionsOrPreambleOrThreadName)) {\n          preamble = optionsOrPreambleOrThreadName as ExecutablePrompt<I>;\n        } else {\n          options = optionsOrPreambleOrThreadName as ChatOptions<I, S>;\n        }\n      }\n      if (maybeOptionsOrPreamble) {\n        if (isExecutablePrompt(maybeOptionsOrPreamble)) {\n          preamble = maybeOptionsOrPreamble as ExecutablePrompt<I>;\n        } else {\n          options = maybeOptionsOrPreamble as ChatOptions<I, S>;\n        }\n      }\n      if (maybeOptions) {\n        options = maybeOptions as ChatOptions<I, S>;\n      }\n\n      let requestBase: Promise<BaseGenerateOptions>;\n      if (preamble) {\n        const renderOptions = options as PromptRenderOptions<I>;\n        requestBase = preamble\n          .render(renderOptions?.input, renderOptions as PromptGenerateOptions)\n          .then((rb) => {\n            return {\n              ...rb,\n              messages: tagAsPreamble(rb?.messages),\n            };\n          });\n      } else {\n        const baseOptions = { ...(options as BaseGenerateOptions) };\n        const messages: MessageData[] = [];\n        if (baseOptions.system) {\n          messages.push({\n            role: 'system',\n            content: Message.parseContent(baseOptions.system),\n          });\n        }\n        delete baseOptions.system;\n        if (baseOptions.messages) {\n          messages.push(...baseOptions.messages);\n        }\n        baseOptions.messages = tagAsPreamble(messages);\n\n        requestBase = Promise.resolve(baseOptions);\n      }\n      return new Chat(this, requestBase, {\n        thread: threadName,\n        id: this.id,\n        messages:\n          (this.sessionData?.threads &&\n            this.sessionData?.threads[threadName]) ??\n          [],\n      });\n    });\n  }\n\n  /**\n   * Executes provided function within this session context allowing calling\n   * `ai.currentSession().state`\n   */\n  run<O>(fn: () => O) {\n    return runWithSession(this.registry, this, fn);\n  }\n\n  toJSON() {\n    return this.sessionData;\n  }\n}\n\nexport interface SessionData<S = any> {\n  id: string;\n  state?: S;\n  threads?: Record<string, MessageData[]>;\n}\n\nconst sessionAlsKey = 'ai.session';\n\n/**\n * Executes provided function within the provided session state.\n */\nexport function runWithSession<S = any, O = any>(\n  registry: Registry,\n  session: Session<S>,\n  fn: () => O\n): O {\n  return getAsyncContext().run(sessionAlsKey, session, fn);\n}\n\n/** Returns the current session. */\nexport function getCurrentSession<S = any>(\n  registry: Registry\n): Session<S> | undefined {\n  return getAsyncContext().getStore(sessionAlsKey);\n}\n\n/** Throw when session state errors occur, ex. missing state, etc. */\nexport class SessionError extends Error {\n  constructor(msg: string) {\n    super(msg);\n  }\n}\n\n/** Session store persists session data such as state and chat messages. */\nexport interface SessionStore<S = any> {\n  get(sessionId: string): Promise<SessionData<S> | undefined>;\n\n  save(sessionId: string, data: Omit<SessionData<S>, 'id'>): Promise<void>;\n}\n\nexport function inMemorySessionStore() {\n  return new InMemorySessionStore();\n}\n\nclass InMemorySessionStore<S = any> implements SessionStore<S> {\n  private data: Record<string, SessionData<S>> = {};\n\n  async get(sessionId: string): Promise<SessionData<S> | undefined> {\n    return this.data[sessionId];\n  }\n\n  async save(sessionId: string, sessionData: SessionData<S>): Promise<void> {\n    this.data[sessionId] = sessionData;\n  }\n}\n"],"mappings":"AAgBA,SAAS,uBAA+B;AAExC,SAAS,MAAM,cAAc;AAC7B;AAAA,EACE;AAAA,EACA;AAAA,OAGK;AACP;AAAA,EACE;AAAA,EACA;AAAA,EACA;AAAA,OAKK;AA4BA,MAAM,QAAiB;AAAA,EAK5B,YACW,UACT,SAMA;AAPS;AAQT,SAAK,KAAK,SAAS,MAAM,OAAO;AAChC,SAAK,cAAc,SAAS,eAAe;AAAA,MACzC,IAAI,KAAK;AAAA,IACX;AACA,QAAI,CAAC,KAAK,aAAa;AACrB,WAAK,cAAc,EAAE,IAAI,KAAK,GAAG;AAAA,IACnC;AACA,QAAI,CAAC,KAAK,YAAY,SAAS;AAC7B,WAAK,YAAa,UAAU,CAAC;AAAA,IAC/B;AACA,SAAK,QAAQ,SAAS,SAAS,IAAI,qBAAwB;AAAA,EAC7D;AAAA,EAxBS;AAAA,EACD;AAAA,EACA;AAAA,EAwBR,IAAI,QAAuB;AACzB,WAAO,KAAK,YAAa;AAAA,EAC3B;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,YAAY,MAAwB;AACxC,QAAI,cAAc,KAAK;AACvB,QAAI,CAAC,aAAa;AAChB,oBAAc,CAAC;AAAA,IACjB;AACA,gBAAY,QAAQ;AACpB,SAAK,cAAc;AAEnB,UAAM,KAAK,MAAM,KAAK,KAAK,IAAI,WAAW;AAAA,EAC5C;AAAA;AAAA;AAAA;AAAA,EAKA,MAAM,eAAe,QAAgB,UAAwC;AAC3E,QAAI,cAAc,KAAK;AACvB,QAAI,CAAC,aAAa;AAChB,oBAAc,CAAC;AAAA,IACjB;AACA,QAAI,CAAC,YAAY,SAAS;AACxB,kBAAY,UAAU,CAAC;AAAA,IACzB;AACA,gBAAY,QAAQ,MAAM,IAAI,SAAS;AAAA,MAAI,CAAC,MAC1C,EAAE,SAAS,EAAE,OAAO,IAAI;AAAA,IAC1B;AACA,SAAK,cAAc;AAEnB,UAAM,KAAK,MAAM,KAAK,KAAK,IAAI,WAAW;AAAA,EAC5C;AAAA,EAoEA,KACE,+BAIA,wBACA,cACM;AACN,WAAO,eAAe,KAAK,UAAU,MAAM,MAAM;AAC/C,UAAI;AACJ,UAAI,aAAa;AACjB,UAAI;AAEJ,UAAI,+BAA+B;AACjC,YAAI,OAAO,kCAAkC,UAAU;AACrD,uBAAa;AAAA,QACf,WAAW,mBAAmB,6BAA6B,GAAG;AAC5D,qBAAW;AAAA,QACb,OAAO;AACL,oBAAU;AAAA,QACZ;AAAA,MACF;AACA,UAAI,wBAAwB;AAC1B,YAAI,mBAAmB,sBAAsB,GAAG;AAC9C,qBAAW;AAAA,QACb,OAAO;AACL,oBAAU;AAAA,QACZ;AAAA,MACF;AACA,UAAI,cAAc;AAChB,kBAAU;AAAA,MACZ;AAEA,UAAI;AACJ,UAAI,UAAU;AACZ,cAAM,gBAAgB;AACtB,sBAAc,SACX,OAAO,eAAe,OAAO,aAAsC,EACnE,KAAK,CAAC,OAAO;AACZ,iBAAO;AAAA,YACL,GAAG;AAAA,YACH,UAAU,cAAc,IAAI,QAAQ;AAAA,UACtC;AAAA,QACF,CAAC;AAAA,MACL,OAAO;AACL,cAAM,cAAc,EAAE,GAAI,QAAgC;AAC1D,cAAM,WAA0B,CAAC;AACjC,YAAI,YAAY,QAAQ;AACtB,mBAAS,KAAK;AAAA,YACZ,MAAM;AAAA,YACN,SAAS,QAAQ,aAAa,YAAY,MAAM;AAAA,UAClD,CAAC;AAAA,QACH;AACA,eAAO,YAAY;AACnB,YAAI,YAAY,UAAU;AACxB,mBAAS,KAAK,GAAG,YAAY,QAAQ;AAAA,QACvC;AACA,oBAAY,WAAW,cAAc,QAAQ;AAE7C,sBAAc,QAAQ,QAAQ,WAAW;AAAA,MAC3C;AACA,aAAO,IAAI,KAAK,MAAM,aAAa;AAAA,QACjC,QAAQ;AAAA,QACR,IAAI,KAAK;AAAA,QACT,WACG,KAAK,aAAa,WACjB,KAAK,aAAa,QAAQ,UAAU,MACtC,CAAC;AAAA,MACL,CAAC;AAAA,IACH,CAAC;AAAA,EACH;AAAA;AAAA;AAAA;AAAA;AAAA,EAMA,IAAO,IAAa;AAClB,WAAO,eAAe,KAAK,UAAU,MAAM,EAAE;AAAA,EAC/C;AAAA,EAEA,SAAS;AACP,WAAO,KAAK;AAAA,EACd;AACF;AAQA,MAAM,gBAAgB;AAKf,SAAS,eACd,UACA,SACA,IACG;AACH,SAAO,gBAAgB,EAAE,IAAI,eAAe,SAAS,EAAE;AACzD;AAGO,SAAS,kBACd,UACwB;AACxB,SAAO,gBAAgB,EAAE,SAAS,aAAa;AACjD;AAGO,MAAM,qBAAqB,MAAM;AAAA,EACtC,YAAY,KAAa;AACvB,UAAM,GAAG;AAAA,EACX;AACF;AASO,SAAS,uBAAuB;AACrC,SAAO,IAAI,qBAAqB;AAClC;AAEA,MAAM,qBAAyD;AAAA,EACrD,OAAuC,CAAC;AAAA,EAEhD,MAAM,IAAI,WAAwD;AAChE,WAAO,KAAK,KAAK,SAAS;AAAA,EAC5B;AAAA,EAEA,MAAM,KAAK,WAAmB,aAA4C;AACxE,SAAK,KAAK,SAAS,IAAI;AAAA,EACzB;AACF;","names":[]}