feat: Allow to pass options to countTokens method (#268)

davidmigloz · Dec 15, 2023 · 4ecb123 · 4ecb123
1 parent ccbdf85
commit 4ecb123
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/packages/langchain/lib/src/model_io/language_models/base.dart b/packages/langchain/lib/src/model_io/language_models/base.dart
@@ -73,8 +73,11 @@ abstract class BaseLanguageModel<Input extends Object,
   ///
   /// Note: subclasses can override this method to provide a more accurate
   /// implementation.
-  Future<int> countTokens(final PromptValue promptValue) async {
-    final tokens = await tokenize(promptValue);
+  Future<int> countTokens(
+    final PromptValue promptValue, {
+    final Options? options,
+  }) async {
+    final tokens = await tokenize(promptValue, options: options);
     return tokens.length;
   }
 

diff --git a/packages/langchain_openai/lib/src/chat_models/openai.dart b/packages/langchain_openai/lib/src/chat_models/openai.dart
@@ -321,8 +321,11 @@ class ChatOpenAI extends BaseChatModel<ChatOpenAIOptions> {
   }
 
   @override
-  Future<int> countTokens(final PromptValue promptValue) async {
-    final model = defaultOptions.model;
+  Future<int> countTokens(
+    final PromptValue promptValue, {
+    final ChatOpenAIOptions? options,
+  }) async {
+    final model = options?.model ?? defaultOptions.model;
     final tiktoken = _getTiktoken(model);
     final messages = promptValue.toChatMessages();