From 4ecb123bd34f0b01d377045b97dace89676d5d16 Mon Sep 17 00:00:00 2001
From: David Miguel Lozano <me@davidmiguel.com>
Date: Fri, 15 Dec 2023 18:17:56 +0100
Subject: [PATCH] feat: Allow to pass options to countTokens method (#268)

---
 .../langchain/lib/src/model_io/language_models/base.dart   | 7 +++++--
 packages/langchain_openai/lib/src/chat_models/openai.dart  | 7 +++++--
 2 files changed, 10 insertions(+), 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
index 84b9ea4a..7a46cd09 100644
--- 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
index 9b7528d0..74115102 100644
--- 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();