Merge pull request #58 from cequence-io/token_count_polishing

Token count polishing

Author: peterbanda

README.md

@@ -363,16 +363,62 @@ For this to work you need to use `OpenAIServiceStreamedFactory` from `openai-sca
   }
 ```
 
-- 🔥 **New**: Count expected used tokens before calling `createChatCompletions` or `createChatFunCompletions`, this help you select proper model ex. `gpt-3.5-turbo` or `gpt-3.5-turbo-16k` and reduce costs. This is an experimental feature and it may not work for all models. Requires `openai-scala-count-tokens` lib.
+- 🔥 **New**: Count expected used tokens before calling `createChatCompletions` or `createChatFunCompletions`, this helps you select proper model ex. `gpt-3.5-turbo` or `gpt-3.5-turbo-16k` and reduce costs. This is an experimental feature and it may not work for all models. Requires `openai-scala-count-tokens` lib.
 
+An example how to count message tokens:
+```scala
+import io.cequence.openaiscala.domain.{AssistantMessage, BaseMessage, FunctionSpec, ModelId, SystemMessage, UserMessage}
+
+class MyCompletionService extends OpenAICountTokensHelper {
+  def exec = {
+    val model = ModelId.gpt_4_turbo_2024_04_09
+
+    // messages to be sent to OpenAI
+    val messages: Seq[BaseMessage] = Seq(
+      SystemMessage("You are a helpful assistant."),
+      UserMessage("Who won the world series in 2020?"),
+      AssistantMessage("The Los Angeles Dodgers won the World Series in 2020."),
+      UserMessage("Where was it played?"),
+    )
+
+    val tokens = countMessageTokens(model, messages)
+  }
+}
+```
+
+An example how to count message tokens when a function is involved:
 ```scala
 import io.cequence.openaiscala.service.OpenAICountTokensHelper
 import io.cequence.openaiscala.domain.{ChatRole, FunMessageSpec, FunctionSpec}
 
+// TODO: simpler example
+import io.cequence.openaiscala.domain.{BaseMessage, FunctionSpec, ModelId, SystemMessage, UserMessage}
+
 class MyCompletionService extends OpenAICountTokensHelper {
   def exec = {
-    val messages: Seq[FunMessageSpec] = ??? // messages to be sent to OpenAI
-    val function: FunctionSpec = ??? // function to be called
+    val model = ModelId.gpt_4_turbo_2024_04_09
+    
+    // messages to be sent to OpenAI
+    val messages: Seq[BaseMessage] = 
+     Seq(
+       SystemMessage("You are a helpful assistant."),
+       UserMessage("What's the weather like in San Francisco, Tokyo, and Paris?")
+     )
+     
+    // function to be called
+    val function: FunctionSpec = FunctionSpec(
+      name = "getWeather",
+      parameters = Map(
+        "type" -> "object",
+        "properties" -> Map(
+          "location" -> Map(
+            "type" -> "string",
+            "description" -> "The city to get the weather for"
+          ),
+          "unit" -> Map("type" -> "string", "enum" -> List("celsius", "fahrenheit"))
+        )
+      )
+    )
 
     val tokens = countFunMessageTokens(model, messages, Seq(function), Some(function.name))
   }

openai-count-tokens/README.md

@@ -27,16 +27,25 @@ or to *pom.xml* (if you use maven)
 
 ## Usage
 
+An example how to count message tokens:
 ```scala
-import io.cequence.openaiscala.service.OpenAICountTokensHelper
-import io.cequence.openaiscala.domain.{ChatRole, FunMessageSpec, FunctionSpec}
-
-val messages: Seq[FunMessageSpec] = ??? // messages to be sent to OpenAI
-val function: FunctionSpec = ??? // function to be called
-
-val service = new OpenAICountTokensService()
-
-val tokens = service.countFunMessageTokens(messages, List(function), Some(function.name))
+import io.cequence.openaiscala.domain.{AssistantMessage, BaseMessage, FunctionSpec, ModelId, SystemMessage, UserMessage}
+
+class MyCompletionService extends OpenAICountTokensHelper {
+  def exec = {
+    val model = ModelId.gpt_4_turbo_2024_04_09
+
+    // messages to be sent to OpenAI
+    val messages: Seq[BaseMessage] = Seq(
+      SystemMessage("You are a helpful assistant."),
+      UserMessage("Who won the world series in 2020?"),
+      AssistantMessage("The Los Angeles Dodgers won the World Series in 2020."),
+      UserMessage("Where was it played?"),
+    )
+
+    val tokens = countMessageTokens(model, messages)
+  }
+}
 ```
 
 

openai-count-tokens/src/main/scala/io/cequence/openaiscala/service/OpenAICountTokensHelper.scala

@@ -49,16 +49,16 @@ trait OpenAICountTokensHelper {
 
   private def tokensPerMessageAndName(model: String): (Int, Int) =
     model match {
-      case "gpt-3.5-turbo-0301" =>
+      case ModelId.gpt_3_5_turbo_0301 =>
         // every message follows <|start|>{role/name}\n{content}<|end|>\n
         // if there's a name, the role is omitted
         (4, -1)
-      case "gpt-3.5-turbo-0613" | "gpt-3.5-turbo-16k-0613" | "gpt-4-0314" | "gpt-4-32k-0314" |
-          "gpt-4-0613" | "gpt-4-32k-0613" =>
+      case ModelId.gpt_3_5_turbo_0613 | ModelId.gpt_3_5_turbo_16k_0613 | ModelId.gpt_4_0613 |
+          ModelId.gpt_4_32k_0613 | ModelId.gpt_4_turbo_2024_04_09 =>
         (3, 1)
-      case "gpt-3.5-turbo" => tokensPerMessageAndName("gpt-3.5-turbo-0613")
-      case "gpt-4"         => tokensPerMessageAndName("gpt-4-0613")
-      case _               =>
+      case ModelId.gpt_3_5_turbo => tokensPerMessageAndName(ModelId.gpt_3_5_turbo_0613)
+      case ModelId.gpt_4         => tokensPerMessageAndName(ModelId.gpt_4_0613)
+      case _                     =>
         // failover to (3, 1)
         (3, 1)
     }

openai-count-tokens/src/test/scala/io/cequence/openaiscala/service/OpenAICountTokensServiceSpec.scala

@@ -16,7 +16,7 @@ import org.scalatestplus.mockito.MockitoSugar
 import org.scalatest.concurrent.ScalaFutures
 import org.scalatest.matchers.should.Matchers
 import org.scalatest.wordspec.AnyWordSpecLike
-import org.scalatest.{BeforeAndAfterAll, Ignore}
+import org.scalatest.BeforeAndAfterAll
 
 import scala.collection.immutable.ListMap
 import scala.concurrent.ExecutionContext.Implicits.global