add gpt-4o support for num_tokens

oai_utils/openai_utils.py

@@ -96,7 +96,10 @@ def num_tokens(messages: list[dict[str, str]] | str, model: str = "gpt-4") -> in
             4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
         )
         tokens_per_name = -1  # if there's a name, the role is omitted
-    elif model == "gpt-4":
+    elif model in [
+        "gpt-4",
+        "gpt-4o",  # TODO: confirm gpt-4o rules are same as gpt-4, see https://github.com/related-sciences/oai_utils/pull/2
+    ]:
         tokens_per_message = 3
         tokens_per_name = 1
     else:

oai_utils/tests/openai_utils_test.py

@@ -6,10 +6,21 @@
     chat_completion_batch,
     embedding,
 )
+from oai_utils.openai_utils import num_tokens
 
 pytestmark = pytest.mark.openai_api_tests
 
 
+@pytest.mark.parametrize("model", ["gpt-3.5-turbo", "gpt-4", "gpt-4o"])
+def test_support_different_model_tokens(model):
+    assert (
+        num_tokens(
+            "What is the answer to life, the universe, and everything?", model=model
+        )
+        == 13
+    )
+
+
 def test_chat_completion__dummy():
     r = chat_completion("What is the answer to life, the universe, and everything?")
     assert r.request.messages == [