@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
1 lines • 12.5 kB
Source Map (JSON)
{"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":[]}