162 lines
4.9 KiB
JavaScript
162 lines
4.9 KiB
JavaScript
import _ from "lodash";
|
|
import { Portkey as _Portkey } from "portkey-ai";
|
|
import { GenerationChunk } from "@langchain/core/outputs";
|
|
import { getEnvironmentVariable } from "@langchain/core/utils/env";
|
|
import { BaseLLM } from "@langchain/core/language_models/llms";
|
|
const readEnv = (env, default_val) => getEnvironmentVariable(env) ?? default_val;
|
|
export class PortkeySession {
|
|
constructor(options = {}) {
|
|
Object.defineProperty(this, "portkey", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
if (!options.apiKey) {
|
|
/* eslint-disable no-param-reassign */
|
|
options.apiKey = readEnv("PORTKEY_API_KEY");
|
|
}
|
|
if (!options.baseURL) {
|
|
/* eslint-disable no-param-reassign */
|
|
options.baseURL = readEnv("PORTKEY_BASE_URL", "https://api.portkey.ai");
|
|
}
|
|
this.portkey = new _Portkey({});
|
|
this.portkey.llms = [{}];
|
|
if (!options.apiKey) {
|
|
throw new Error("Set Portkey ApiKey in PORTKEY_API_KEY env variable");
|
|
}
|
|
this.portkey = new _Portkey(options);
|
|
}
|
|
}
|
|
const defaultPortkeySession = [];
|
|
/**
|
|
* Get a session for the Portkey API. If one already exists with the same options,
|
|
* it will be returned. Otherwise, a new session will be created.
|
|
* @param options
|
|
* @returns
|
|
*/
|
|
export function getPortkeySession(options = {}) {
|
|
let session = defaultPortkeySession.find((session) => _.isEqual(session.options, options))?.session;
|
|
if (!session) {
|
|
session = new PortkeySession(options);
|
|
defaultPortkeySession.push({ session, options });
|
|
}
|
|
return session;
|
|
}
|
|
/**
|
|
* @example
|
|
* ```typescript
|
|
* const model = new Portkey({
|
|
* mode: "single",
|
|
* llms: [
|
|
* {
|
|
* provider: "openai",
|
|
* virtual_key: "open-ai-key-1234",
|
|
* model: "gpt-3.5-turbo-instruct",
|
|
* max_tokens: 2000,
|
|
* },
|
|
* ],
|
|
* });
|
|
*
|
|
* // Stream the output of the model and process it
|
|
* const res = await model.stream(
|
|
* "Question: Write a story about a king\nAnswer:"
|
|
* );
|
|
* for await (const i of res) {
|
|
* process.stdout.write(i);
|
|
* }
|
|
* ```
|
|
*/
|
|
export class Portkey extends BaseLLM {
|
|
constructor(init) {
|
|
super(init ?? {});
|
|
Object.defineProperty(this, "apiKey", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: undefined
|
|
});
|
|
Object.defineProperty(this, "baseURL", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: undefined
|
|
});
|
|
Object.defineProperty(this, "mode", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: undefined
|
|
});
|
|
Object.defineProperty(this, "llms", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: undefined
|
|
});
|
|
Object.defineProperty(this, "session", {
|
|
enumerable: true,
|
|
configurable: true,
|
|
writable: true,
|
|
value: void 0
|
|
});
|
|
this.apiKey = init?.apiKey;
|
|
this.baseURL = init?.baseURL;
|
|
this.mode = init?.mode;
|
|
this.llms = init?.llms;
|
|
this.session = getPortkeySession({
|
|
apiKey: this.apiKey,
|
|
baseURL: this.baseURL,
|
|
llms: this.llms,
|
|
mode: this.mode,
|
|
});
|
|
}
|
|
_llmType() {
|
|
return "portkey";
|
|
}
|
|
async _generate(prompts, options, _) {
|
|
const choices = [];
|
|
for (let i = 0; i < prompts.length; i += 1) {
|
|
const response = await this.session.portkey.completions.create({
|
|
prompt: prompts[i],
|
|
...options,
|
|
stream: false,
|
|
});
|
|
choices.push(response.choices);
|
|
}
|
|
const generations = choices.map((promptChoices) => promptChoices.map((choice) => ({
|
|
text: choice.text ?? "",
|
|
generationInfo: {
|
|
finishReason: choice.finish_reason,
|
|
logprobs: choice.logprobs,
|
|
},
|
|
})));
|
|
return {
|
|
generations,
|
|
};
|
|
}
|
|
async *_streamResponseChunks(input, options, runManager) {
|
|
const response = await this.session.portkey.completions.create({
|
|
prompt: input,
|
|
...options,
|
|
stream: true,
|
|
});
|
|
for await (const data of response) {
|
|
const choice = data?.choices[0];
|
|
if (!choice) {
|
|
continue;
|
|
}
|
|
const chunk = new GenerationChunk({
|
|
text: choice.text ?? "",
|
|
generationInfo: {
|
|
finishReason: choice.finish_reason,
|
|
},
|
|
});
|
|
yield chunk;
|
|
void runManager?.handleLLMNewToken(chunk.text ?? "");
|
|
}
|
|
if (options.signal?.aborted) {
|
|
throw new Error("AbortError");
|
|
}
|
|
}
|
|
}
|