From 551d3635540e2a82b4fc328a152205f432e0136d Mon Sep 17 00:00:00 2001 From: purp Date: Fri, 28 Apr 2023 14:11:10 -0600 Subject: [PATCH] Add token count estimation to OpenAIChat LLM instance, to ensure parity with OpenAI LLM instance. --- langchain/src/llms/openai-chat.ts | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/langchain/src/llms/openai-chat.ts b/langchain/src/llms/openai-chat.ts index 1ef21d9d6458..172a98e345fb 100644 --- a/langchain/src/llms/openai-chat.ts +++ b/langchain/src/llms/openai-chat.ts @@ -10,6 +10,7 @@ import { import type { AxiosRequestConfig } from "axios"; import type { StreamingAxiosConfiguration } from "../util/axios-types.js"; import fetchAdapter from "../util/axios-fetch-adapter.js"; +import { calculateMaxTokens } from "../base_language/count_tokens.js"; import { BaseLLMCallOptions, BaseLLMParams, LLM } from "./base.js"; import { CallbackManagerForLLMRun } from "../callbacks/manager.js"; @@ -235,6 +236,20 @@ export class OpenAIChat extends LLM implements OpenAIChatInput { const params = this.invocationParams(); params.stop = stop ?? params.stop; + if (params.max_tokens === -1) { + // Include prefixes in the token count, as this will be enforced + // (if prefixes exist) by the OpenAI API token limits: + // https://platform.openai.com/docs/api-reference/completions/create#max_tokens + // Dump this all into one big blob of text with `.strigify` to + // ensure that we are pessimistic about the token count and don't + // leave anything out! + const promptAndPrefixes = JSON.stringify(this.formatMessages(prompt)); + params.max_tokens = await calculateMaxTokens({ + prompt: promptAndPrefixes, + // Cast here to allow for other models that may not fit the union + modelName: this.modelName, + }); + } const data = params.stream ? await new Promise((resolve, reject) => { let response: CreateChatCompletionResponse;