feat: add encoding detection for 4o and 4o-mini ( tests)

knuddelsgmbh · Jul 19, 2024 · accaf58 · accaf58
1 parent 681ad70
commit accaf58
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 0 deletions.
diff --git a/lib/src/main/java/com/knuddels/jtokkit/AbstractEncodingRegistry.java b/lib/src/main/java/com/knuddels/jtokkit/AbstractEncodingRegistry.java
@@ -32,6 32,10 @@ public Optional<Encoding> getEncodingForModel(final String modelName) {
  return Optional.of(getEncodingForModel(modelType.get()));
  }
 
  if (modelName.startsWith(ModelType.GPT_4O.getName())) {
  return Optional.of(getEncodingForModel(ModelType.GPT_4O));
  }
 
  if (modelName.startsWith(ModelType.GPT_4_32K.getName())) {
  return Optional.of(getEncodingForModel(ModelType.GPT_4_32K));
  }

diff --git a/lib/src/test/java/com/knuddels/jtokkit/BaseEncodingRegistryTest.java b/lib/src/test/java/com/knuddels/jtokkit/BaseEncodingRegistryTest.java
@@ -88,6 88,27 @@ void getEncodingForModelByPrefixReturnsCorrectEncodingForGpt4() {
  assertEquals(encoding.get().getName(), ModelType.GPT_4.getEncodingType().getName());
  }
 
  @Test
  void getEncodingForModelByPrefixReturnsCorrectEncodingForGpt4o() {
  var encoding = registry.getEncodingForModel("gpt-4o-123");
  assertTrue(encoding.isPresent());
  assertEquals(encoding.get().getName(), ModelType.GPT_4O.getEncodingType().getName());
  }
 
  @Test
  void getEncodingForModelByPrefixReturnsCorrectEncodingForGpt4oMini() {
  var encoding = registry.getEncodingForModel("gpt-4o-mini-123");
  assertTrue(encoding.isPresent());
  assertEquals(encoding.get().getName(), ModelType.GPT_4O_MINI.getEncodingType().getName());
  }
 
  @Test
  void getEncodingForModelByPrefixReturnsCorrectEncodingForGpt4Turbo() {
  var encoding = registry.getEncodingForModel("gpt-4-turbo-123");
  assertTrue(encoding.isPresent());
  assertEquals(encoding.get().getName(), ModelType.GPT_4_TURBO.getEncodingType().getName());
  }
 
  @Test
  void getEncodingForModelByPrefixReturnsCorrectEncodingForGpt3Turbo() {
  var encoding = registry.getEncodingForModel("gpt-3.5-turbo-0301");