@genkit-ai/ai
Version:
Genkit AI framework generative AI APIs.
1 lines • 11.3 kB
Source Map (JSON)
{"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 type GenerationCommonConfigSchema,\n type MessageData,\n type Part,\n} from './index.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,OASK;AACP;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"]}