@langchain/community
Version:
Third-party integrations for LangChain.js
1 lines • 18.2 kB
Source Map (JSON)
{"version":3,"file":"premai.cjs","names":["ChatMessage","AIMessage","HumanMessageChunk","AIMessageChunk","ChatMessageChunk","BaseChatModel","Prem","ChatGenerationChunk"],"sources":["../../src/chat_models/premai.ts"],"sourcesContent":["import {\n AIMessage,\n AIMessageChunk,\n type BaseMessage,\n ChatMessage,\n ChatMessageChunk,\n HumanMessageChunk,\n} from \"@langchain/core/messages\";\nimport {\n type BaseLanguageModelCallOptions,\n TokenUsage,\n} from \"@langchain/core/language_models/base\";\n\nimport { CallbackManagerForLLMRun } from \"@langchain/core/callbacks/manager\";\nimport {\n type BaseChatModelParams,\n BaseChatModel,\n} from \"@langchain/core/language_models/chat_models\";\n\nimport Prem, {\n ChatCompletionStreamingCompletionData,\n CreateChatCompletionRequest,\n CreateChatCompletionResponse,\n} from \"@premai/prem-sdk\";\nimport { getEnvironmentVariable } from \"@langchain/core/utils/env\";\nimport {\n ChatGeneration,\n ChatGenerationChunk,\n ChatResult,\n} from \"@langchain/core/outputs\";\n\nimport { NewTokenIndices } from \"@langchain/core/callbacks/base\";\n\nexport type RoleEnum = \"user\" | \"assistant\";\n\n/**\n * Input to chat model class.\n */\nexport interface ChatPremInput extends BaseChatModelParams {\n project_id?: number | string;\n session_id?: string;\n messages?: {\n role: \"user\" | \"assistant\";\n content: string;\n [k: string]: unknown;\n }[];\n model?: string;\n system_prompt?: string;\n frequency_penalty?: number;\n logit_bias?: { [k: string]: unknown };\n max_tokens?: number;\n n?: number;\n presence_penalty?: number;\n response_format?: { [k: string]: unknown };\n seed?: number;\n stop?: string;\n temperature?: number;\n top_p?: number;\n tools?: { [k: string]: unknown }[];\n user?: string;\n /**\n * The Prem API key to use for requests.\n * @default process.env.PREM_API_KEY\n */\n apiKey?: string;\n streaming?: boolean;\n}\n\nexport interface ChatCompletionCreateParamsNonStreaming extends CreateChatCompletionRequest {\n stream?: false;\n}\n\nexport interface ChatCompletionCreateParamsStreaming extends CreateChatCompletionRequest {\n stream: true;\n}\n\nexport type ChatCompletionCreateParams =\n | ChatCompletionCreateParamsNonStreaming\n | ChatCompletionCreateParamsStreaming;\n\nfunction extractGenericMessageCustomRole(message: ChatMessage) {\n if (message.role !== \"assistant\" && message.role !== \"user\") {\n console.warn(`Unknown message role: ${message.role}`);\n }\n return message.role as RoleEnum;\n}\n\nexport function messageToPremRole(message: BaseMessage): RoleEnum {\n const type = message._getType();\n switch (type) {\n case \"ai\":\n return \"assistant\";\n case \"human\":\n return \"user\";\n case \"generic\": {\n if (!ChatMessage.isInstance(message))\n throw new Error(\"Invalid generic chat message\");\n return extractGenericMessageCustomRole(message);\n }\n default:\n throw new Error(`Unknown message type: ${type}`);\n }\n}\n\nfunction convertMessagesToPremParams(\n messages: BaseMessage[]\n): Array<CreateChatCompletionResponse[\"choices\"][\"0\"][\"message\"]> {\n return messages.map((message) => {\n if (typeof message.content !== \"string\") {\n throw new Error(\"Non string message content not supported\");\n }\n return {\n role: messageToPremRole(message),\n content: message.content,\n name: message.name,\n function_call: message.additional_kwargs.function_call,\n };\n });\n}\n\nfunction premResponseToChatMessage(\n message: CreateChatCompletionResponse[\"choices\"][\"0\"][\"message\"]\n): BaseMessage {\n switch (message.role) {\n case \"assistant\":\n return new AIMessage(message.content as string);\n default:\n return new ChatMessage(\n message.content as string,\n message.role ?? \"unknown\"\n );\n }\n}\n\nfunction _convertDeltaToMessageChunk(\n // oxlint-disable-next-line typescript/no-explicit-any\n delta: Record<string, any>\n) {\n const { role } = delta;\n const content = delta.content ?? \"\";\n let additional_kwargs;\n if (delta.function_call) {\n additional_kwargs = {\n function_call: delta.function_call,\n };\n } else {\n additional_kwargs = {};\n }\n if (role === \"user\") {\n return new HumanMessageChunk({ content });\n } else if (role === \"assistant\") {\n return new AIMessageChunk({ content, additional_kwargs });\n } else {\n return new ChatMessageChunk({ content, role });\n }\n}\n\n/**\n * Integration with a chat model.\n */\nexport class ChatPrem<\n CallOptions extends BaseLanguageModelCallOptions =\n BaseLanguageModelCallOptions,\n>\n extends BaseChatModel<CallOptions>\n implements ChatPremInput\n{\n client: Prem;\n\n apiKey?: string;\n\n project_id: number;\n\n session_id?: string;\n\n messages: {\n [k: string]: unknown;\n role: \"user\" | \"assistant\";\n content: string;\n }[];\n\n model?: string;\n\n system_prompt?: string;\n\n frequency_penalty?: number;\n\n logit_bias?: { [k: string]: unknown };\n\n max_tokens?: number;\n\n n?: number;\n\n presence_penalty?: number;\n\n response_format?: { [k: string]: unknown };\n\n seed?: number;\n\n stop?: string;\n\n temperature?: number;\n\n top_p?: number;\n\n tools?: { [k: string]: unknown }[];\n\n user?: string;\n\n streaming = false;\n\n [k: string]: unknown;\n\n // Used for tracing, replace with the same name as your class\n static lc_name() {\n return \"ChatPrem\";\n }\n\n lc_serializable = true;\n\n /**\n * Replace with any secrets this class passes to `super`.\n * See {@link ../../langchain-cohere/src/chat_model.ts} for\n * an example.\n */\n get lc_secrets(): { [key: string]: string } | undefined {\n return {\n apiKey: \"PREM_API_KEY\",\n };\n }\n\n get lc_aliases(): { [key: string]: string } | undefined {\n return {\n apiKey: \"PREM_API_KEY\",\n };\n }\n\n constructor(fields?: ChatPremInput) {\n super(fields ?? {});\n const apiKey = fields?.apiKey ?? getEnvironmentVariable(\"PREM_API_KEY\");\n if (!apiKey) {\n throw new Error(\n `Prem API key not found. Please set the PREM_API_KEY environment variable or provide the key into \"apiKey\"`\n );\n }\n\n const projectId =\n fields?.project_id ??\n parseInt(getEnvironmentVariable(\"PREM_PROJECT_ID\") ?? \"-1\", 10);\n if (!projectId || projectId === -1 || typeof projectId !== \"number\") {\n throw new Error(\n `Prem project ID not found. Please set the PREM_PROJECT_ID environment variable or provide the key into \"project_id\"`\n );\n }\n\n this.client = new Prem({\n apiKey,\n });\n\n this.project_id = projectId;\n this.session_id = fields?.session_id ?? this.session_id;\n this.messages = fields?.messages ?? this.messages;\n this.model = fields?.model ?? this.model;\n this.system_prompt = fields?.system_prompt ?? this.system_prompt;\n this.frequency_penalty =\n fields?.frequency_penalty ?? this.frequency_penalty;\n this.logit_bias = fields?.logit_bias ?? this.logit_bias;\n this.max_tokens = fields?.max_tokens ?? this.max_tokens;\n this.n = fields?.n ?? this.n;\n this.presence_penalty = fields?.presence_penalty ?? this.presence_penalty;\n this.response_format = fields?.response_format ?? this.response_format;\n this.seed = fields?.seed ?? this.seed;\n this.stop = fields?.stop ?? this.stop;\n this.temperature = fields?.temperature ?? this.temperature;\n this.top_p = fields?.top_p ?? this.top_p;\n this.tools = fields?.tools ?? this.tools;\n this.user = fields?.user ?? this.user;\n this.streaming = fields?.streaming ?? this.streaming;\n }\n\n // Replace\n _llmType() {\n return \"prem\";\n }\n\n async completionWithRetry(\n request: ChatCompletionCreateParamsStreaming,\n // oxlint-disable-next-line typescript/no-explicit-any\n options?: any\n ): Promise<AsyncIterable<ChatCompletionStreamingCompletionData>>;\n\n async completionWithRetry(\n request: ChatCompletionCreateParams,\n // oxlint-disable-next-line typescript/no-explicit-any\n options?: any\n ): Promise<CreateChatCompletionResponse>;\n\n async completionWithRetry(\n request: ChatCompletionCreateParamsStreaming,\n // oxlint-disable-next-line typescript/no-explicit-any\n options?: any\n ): Promise<\n | AsyncIterable<ChatCompletionStreamingCompletionData>\n | CreateChatCompletionResponse\n > {\n return this.caller.call(async () =>\n this.client.chat.completions.create(request, options)\n );\n }\n\n invocationParams(options: this[\"ParsedCallOptions\"]) {\n const params = super.invocationParams(options);\n return {\n ...params,\n project_id: this.project_id,\n session_id: this.session_id,\n messages: this.messages,\n model: this.model,\n system_prompt: this.system_prompt,\n frequency_penalty: this.frequency_penalty,\n logit_bias: this.logit_bias,\n max_tokens: this.max_tokens,\n n: this.n,\n presence_penalty: this.presence_penalty,\n response_format: this.response_format,\n seed: this.seed,\n stop: this.stop,\n temperature: this.temperature,\n top_p: this.top_p,\n tools: this.tools,\n user: this.user,\n streaming: this.streaming,\n stream: this.streaming,\n };\n }\n\n /**\n * Implement to support streaming.\n * Should yield chunks iteratively.\n */\n async *_streamResponseChunks(\n messages: BaseMessage[],\n options: this[\"ParsedCallOptions\"],\n runManager?: CallbackManagerForLLMRun\n ): AsyncGenerator<ChatGenerationChunk> {\n const params = this.invocationParams(options);\n const messagesMapped = convertMessagesToPremParams(messages);\n\n // All models have a built-in `this.caller` property for retries\n const stream = await this.caller.call(async () =>\n this.completionWithRetry(\n {\n ...params,\n messages: messagesMapped,\n stream: true,\n },\n params\n )\n );\n\n for await (const data of stream) {\n const choice = data?.choices[0];\n if (!choice) {\n continue;\n }\n const chunk = new ChatGenerationChunk({\n message: _convertDeltaToMessageChunk(choice.delta ?? {}),\n text: choice.delta.content ?? \"\",\n generationInfo: {\n finishReason: choice.finish_reason,\n },\n });\n yield chunk;\n // eslint-disable-next-line no-void\n void runManager?.handleLLMNewToken(chunk.text ?? \"\");\n }\n if (options.signal?.aborted) {\n throw new Error(\"AbortError\");\n }\n }\n\n /** @ignore */\n _combineLLMOutput() {\n return [];\n }\n\n async _generate(\n messages: BaseMessage[],\n options: this[\"ParsedCallOptions\"],\n runManager?: CallbackManagerForLLMRun\n ): Promise<ChatResult> {\n const tokenUsage: TokenUsage = {};\n const params = this.invocationParams(options);\n const messagesMapped = convertMessagesToPremParams(messages);\n\n if (params.streaming) {\n const stream = this._streamResponseChunks(messages, options, runManager);\n const finalChunks: Record<number, ChatGenerationChunk> = {};\n for await (const chunk of stream) {\n const index =\n (chunk.generationInfo as NewTokenIndices)?.completion ?? 0;\n if (finalChunks[index] === undefined) {\n finalChunks[index] = chunk;\n } else {\n finalChunks[index] = finalChunks[index].concat(chunk);\n }\n }\n const generations = Object.entries(finalChunks)\n .sort(([aKey], [bKey]) => parseInt(aKey, 10) - parseInt(bKey, 10))\n .map(([_, value]) => value);\n\n return { generations, llmOutput: { estimatedTokenUsage: tokenUsage } };\n } else {\n const data = await this.completionWithRetry(\n {\n ...params,\n stream: false,\n messages: messagesMapped,\n },\n {\n signal: options?.signal,\n }\n );\n\n if (\"usage\" in data && data.usage) {\n const {\n completion_tokens: completionTokens,\n prompt_tokens: promptTokens,\n total_tokens: totalTokens,\n } = data.usage as CreateChatCompletionResponse[\"usage\"];\n\n if (completionTokens) {\n tokenUsage.completionTokens =\n (tokenUsage.completionTokens ?? 0) + completionTokens;\n }\n\n if (promptTokens) {\n tokenUsage.promptTokens =\n (tokenUsage.promptTokens ?? 0) + promptTokens;\n }\n\n if (totalTokens) {\n tokenUsage.totalTokens = (tokenUsage.totalTokens ?? 0) + totalTokens;\n }\n }\n\n const generations: ChatGeneration[] = [];\n\n if (\"choices\" in data && data.choices) {\n for (const part of (data as unknown as CreateChatCompletionResponse)\n .choices) {\n const text = part.message?.content ?? \"\";\n const generation: ChatGeneration = {\n text: text as string,\n message: premResponseToChatMessage(\n part.message ?? { role: \"assistant\" }\n ),\n };\n generation.generationInfo = {\n ...(part.finish_reason\n ? { finish_reason: part.finish_reason }\n : {}),\n ...(part.logprobs ? { logprobs: part.logprobs } : {}),\n };\n generations.push(generation);\n }\n }\n\n return {\n generations,\n llmOutput: { tokenUsage },\n };\n }\n }\n}\n"],"mappings":";;;;;;;;;;;;;AAgFA,SAAS,gCAAgC,SAAsB;AAC7D,KAAI,QAAQ,SAAS,eAAe,QAAQ,SAAS,OACnD,SAAQ,KAAK,yBAAyB,QAAQ,OAAO;AAEvD,QAAO,QAAQ;;AAGjB,SAAgB,kBAAkB,SAAgC;CAChE,MAAM,OAAO,QAAQ,UAAU;AAC/B,SAAQ,MAAR;EACE,KAAK,KACH,QAAO;EACT,KAAK,QACH,QAAO;EACT,KAAK;AACH,OAAI,CAACA,yBAAAA,YAAY,WAAW,QAAQ,CAClC,OAAM,IAAI,MAAM,+BAA+B;AACjD,UAAO,gCAAgC,QAAQ;EAEjD,QACE,OAAM,IAAI,MAAM,yBAAyB,OAAO;;;AAItD,SAAS,4BACP,UACgE;AAChE,QAAO,SAAS,KAAK,YAAY;AAC/B,MAAI,OAAO,QAAQ,YAAY,SAC7B,OAAM,IAAI,MAAM,2CAA2C;AAE7D,SAAO;GACL,MAAM,kBAAkB,QAAQ;GAChC,SAAS,QAAQ;GACjB,MAAM,QAAQ;GACd,eAAe,QAAQ,kBAAkB;GAC1C;GACD;;AAGJ,SAAS,0BACP,SACa;AACb,SAAQ,QAAQ,MAAhB;EACE,KAAK,YACH,QAAO,IAAIC,yBAAAA,UAAU,QAAQ,QAAkB;EACjD,QACE,QAAO,IAAID,yBAAAA,YACT,QAAQ,SACR,QAAQ,QAAQ,UACjB;;;AAIP,SAAS,4BAEP,OACA;CACA,MAAM,EAAE,SAAS;CACjB,MAAM,UAAU,MAAM,WAAW;CACjC,IAAI;AACJ,KAAI,MAAM,cACR,qBAAoB,EAClB,eAAe,MAAM,eACtB;KAED,qBAAoB,EAAE;AAExB,KAAI,SAAS,OACX,QAAO,IAAIE,yBAAAA,kBAAkB,EAAE,SAAS,CAAC;UAChC,SAAS,YAClB,QAAO,IAAIC,yBAAAA,eAAe;EAAE;EAAS;EAAmB,CAAC;KAEzD,QAAO,IAAIC,yBAAAA,iBAAiB;EAAE;EAAS;EAAM,CAAC;;;;;AAOlD,IAAa,WAAb,cAIUC,4CAAAA,cAEV;CACE;CAEA;CAEA;CAEA;CAEA;CAMA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA;CAEA,YAAY;CAKZ,OAAO,UAAU;AACf,SAAO;;CAGT,kBAAkB;;;;;;CAOlB,IAAI,aAAoD;AACtD,SAAO,EACL,QAAQ,gBACT;;CAGH,IAAI,aAAoD;AACtD,SAAO,EACL,QAAQ,gBACT;;CAGH,YAAY,QAAwB;AAClC,QAAM,UAAU,EAAE,CAAC;EACnB,MAAM,SAAS,QAAQ,WAAA,GAAA,0BAAA,wBAAiC,eAAe;AACvE,MAAI,CAAC,OACH,OAAM,IAAI,MACR,4GACD;EAGH,MAAM,YACJ,QAAQ,cACR,UAAA,GAAA,0BAAA,wBAAgC,kBAAkB,IAAI,MAAM,GAAG;AACjE,MAAI,CAAC,aAAa,cAAc,MAAM,OAAO,cAAc,SACzD,OAAM,IAAI,MACR,sHACD;AAGH,OAAK,SAAS,IAAIC,iBAAAA,QAAK,EACrB,QACD,CAAC;AAEF,OAAK,aAAa;AAClB,OAAK,aAAa,QAAQ,cAAc,KAAK;AAC7C,OAAK,WAAW,QAAQ,YAAY,KAAK;AACzC,OAAK,QAAQ,QAAQ,SAAS,KAAK;AACnC,OAAK,gBAAgB,QAAQ,iBAAiB,KAAK;AACnD,OAAK,oBACH,QAAQ,qBAAqB,KAAK;AACpC,OAAK,aAAa,QAAQ,cAAc,KAAK;AAC7C,OAAK,aAAa,QAAQ,cAAc,KAAK;AAC7C,OAAK,IAAI,QAAQ,KAAK,KAAK;AAC3B,OAAK,mBAAmB,QAAQ,oBAAoB,KAAK;AACzD,OAAK,kBAAkB,QAAQ,mBAAmB,KAAK;AACvD,OAAK,OAAO,QAAQ,QAAQ,KAAK;AACjC,OAAK,OAAO,QAAQ,QAAQ,KAAK;AACjC,OAAK,cAAc,QAAQ,eAAe,KAAK;AAC/C,OAAK,QAAQ,QAAQ,SAAS,KAAK;AACnC,OAAK,QAAQ,QAAQ,SAAS,KAAK;AACnC,OAAK,OAAO,QAAQ,QAAQ,KAAK;AACjC,OAAK,YAAY,QAAQ,aAAa,KAAK;;CAI7C,WAAW;AACT,SAAO;;CAeT,MAAM,oBACJ,SAEA,SAIA;AACA,SAAO,KAAK,OAAO,KAAK,YACtB,KAAK,OAAO,KAAK,YAAY,OAAO,SAAS,QAAQ,CACtD;;CAGH,iBAAiB,SAAoC;AAEnD,SAAO;GACL,GAFa,MAAM,iBAAiB,QAAQ;GAG5C,YAAY,KAAK;GACjB,YAAY,KAAK;GACjB,UAAU,KAAK;GACf,OAAO,KAAK;GACZ,eAAe,KAAK;GACpB,mBAAmB,KAAK;GACxB,YAAY,KAAK;GACjB,YAAY,KAAK;GACjB,GAAG,KAAK;GACR,kBAAkB,KAAK;GACvB,iBAAiB,KAAK;GACtB,MAAM,KAAK;GACX,MAAM,KAAK;GACX,aAAa,KAAK;GAClB,OAAO,KAAK;GACZ,OAAO,KAAK;GACZ,MAAM,KAAK;GACX,WAAW,KAAK;GAChB,QAAQ,KAAK;GACd;;;;;;CAOH,OAAO,sBACL,UACA,SACA,YACqC;EACrC,MAAM,SAAS,KAAK,iBAAiB,QAAQ;EAC7C,MAAM,iBAAiB,4BAA4B,SAAS;EAG5D,MAAM,SAAS,MAAM,KAAK,OAAO,KAAK,YACpC,KAAK,oBACH;GACE,GAAG;GACH,UAAU;GACV,QAAQ;GACT,EACD,OACD,CACF;AAED,aAAW,MAAM,QAAQ,QAAQ;GAC/B,MAAM,SAAS,MAAM,QAAQ;AAC7B,OAAI,CAAC,OACH;GAEF,MAAM,QAAQ,IAAIC,wBAAAA,oBAAoB;IACpC,SAAS,4BAA4B,OAAO,SAAS,EAAE,CAAC;IACxD,MAAM,OAAO,MAAM,WAAW;IAC9B,gBAAgB,EACd,cAAc,OAAO,eACtB;IACF,CAAC;AACF,SAAM;AAED,eAAY,kBAAkB,MAAM,QAAQ,GAAG;;AAEtD,MAAI,QAAQ,QAAQ,QAClB,OAAM,IAAI,MAAM,aAAa;;;CAKjC,oBAAoB;AAClB,SAAO,EAAE;;CAGX,MAAM,UACJ,UACA,SACA,YACqB;EACrB,MAAM,aAAyB,EAAE;EACjC,MAAM,SAAS,KAAK,iBAAiB,QAAQ;EAC7C,MAAM,iBAAiB,4BAA4B,SAAS;AAE5D,MAAI,OAAO,WAAW;GACpB,MAAM,SAAS,KAAK,sBAAsB,UAAU,SAAS,WAAW;GACxE,MAAM,cAAmD,EAAE;AAC3D,cAAW,MAAM,SAAS,QAAQ;IAChC,MAAM,QACH,MAAM,gBAAoC,cAAc;AAC3D,QAAI,YAAY,WAAW,KAAA,EACzB,aAAY,SAAS;QAErB,aAAY,SAAS,YAAY,OAAO,OAAO,MAAM;;AAOzD,UAAO;IAAE,aAJW,OAAO,QAAQ,YAAY,CAC5C,MAAM,CAAC,OAAO,CAAC,UAAU,SAAS,MAAM,GAAG,GAAG,SAAS,MAAM,GAAG,CAAC,CACjE,KAAK,CAAC,GAAG,WAAW,MAAM;IAEP,WAAW,EAAE,qBAAqB,YAAY;IAAE;SACjE;GACL,MAAM,OAAO,MAAM,KAAK,oBACtB;IACE,GAAG;IACH,QAAQ;IACR,UAAU;IACX,EACD,EACE,QAAQ,SAAS,QAClB,CACF;AAED,OAAI,WAAW,QAAQ,KAAK,OAAO;IACjC,MAAM,EACJ,mBAAmB,kBACnB,eAAe,cACf,cAAc,gBACZ,KAAK;AAET,QAAI,iBACF,YAAW,oBACR,WAAW,oBAAoB,KAAK;AAGzC,QAAI,aACF,YAAW,gBACR,WAAW,gBAAgB,KAAK;AAGrC,QAAI,YACF,YAAW,eAAe,WAAW,eAAe,KAAK;;GAI7D,MAAM,cAAgC,EAAE;AAExC,OAAI,aAAa,QAAQ,KAAK,QAC5B,MAAK,MAAM,QAAS,KACjB,SAAS;IAEV,MAAM,aAA6B;KACjC,MAFW,KAAK,SAAS,WAAW;KAGpC,SAAS,0BACP,KAAK,WAAW,EAAE,MAAM,aAAa,CACtC;KACF;AACD,eAAW,iBAAiB;KAC1B,GAAI,KAAK,gBACL,EAAE,eAAe,KAAK,eAAe,GACrC,EAAE;KACN,GAAI,KAAK,WAAW,EAAE,UAAU,KAAK,UAAU,GAAG,EAAE;KACrD;AACD,gBAAY,KAAK,WAAW;;AAIhC,UAAO;IACL;IACA,WAAW,EAAE,YAAY;IAC1B"}