Count tokens for OpenAI accurately with support for all parameters like name, functions.
"""Nox sessions."""
import os
import shlex
import sys
from pathlib import Path
from textwrap import dedent

import nox


try:
    from nox_poetry import Session, session
except ImportError:
    message = f"""\
    Nox failed to import the 'nox-poetry' package.

    Please install it using the following command:

    {sys.executable} -m pip install nox-poetry"""
    raise SystemExit(dedent(message)) from None


package = "openai_token_counter"
default_python_version = (
    "3.10"  # To be used in sessions where we don't run all python versions
)
python_versions = ["3.11", "3.10", "3.9"]
nox.needs_version = ">= 2021.6.6"
nox.options.sessions = (
    "pre-commit",
    "safety",
    "mypy",
    "tests",
    "typeguard",
)


def activate_virtualenv_in_precommit_hooks(session: Session) -> None:
    """Activate virtualenv in hooks installed by pre-commit.

    This function patches git hooks installed by pre-commit to activate the
    session's virtual environment. This allows pre-commit to locate hooks in
    that environment when invoked from git.

    Args:
        session: The Session object.
    """
    assert session.bin is not None  # noqa: S101

    # Only patch hooks containing a reference to this session's bindir. Support
    # quoting rules for Python and bash, but strip the outermost quotes so we
    # can detect paths within the bindir, like <bindir>/python.
    bindirs = [
        bindir[1:-1] if bindir[0] in "'\"" else bindir
        for bindir in (repr(session.bin), shlex.quote(session.bin))
    ]

    virtualenv = session.env.get("VIRTUAL_ENV")
    if virtualenv is None:
        return

    headers = {
        # pre-commit < 2.16.0
        "python": f"""\
            import os
            os.environ["VIRTUAL_ENV"] = {virtualenv!r}
            os.environ["PATH"] = os.pathsep.join((
                {session.bin!r},
                os.environ.get("PATH", ""),
            ))
            """,
        # pre-commit >= 2.16.0
        "bash": f"""\
            VIRTUAL_ENV={shlex.quote(virtualenv)}
            PATH={shlex.quote(session.bin)}"{os.pathsep}$PATH"
            """,
        # pre-commit >= 2.17.0 on Windows forces sh shebang
        "/bin/sh": f"""\
            VIRTUAL_ENV={shlex.quote(virtualenv)}
            PATH={shlex.quote(session.bin)}"{os.pathsep}$PATH"
            """,
    }

    hookdir = Path(".git") / "hooks"
    if not hookdir.is_dir():
        return

    for hook in hookdir.iterdir():
        if hook.name.endswith(".sample") or not hook.is_file():
            continue

        if not hook.read_bytes().startswith(b"#!"):
            continue

        text = hook.read_text()

        if not any(
            Path("A") == Path("a") and bindir.lower() in text.lower() or bindir in text
            for bindir in bindirs
        ):
            continue

        lines = text.splitlines()

        for executable, header in headers.items():
            if executable in lines[0].lower():
                lines.insert(1, dedent(header))
                hook.write_text("\n".join(lines))
                break


@session(name="pre-commit", python=default_python_version)
def precommit(session: Session) -> None:
    """Lint using pre-commit."""
    args = session.posargs or [
        "run",
        "--all-files",
        "--hook-stage=manual",
        "--show-diff-on-failure",
    ]
    session.install(
        "black",
        "darglint",
        "flake8",
        "flake8-bandit",
        "flake8-bugbear",
        "flake8-docstrings",
        "flake8-rst-docstrings",
        "isort",
        "pep8-naming",
        "pre-commit",
        "pre-commit-hooks",
        "pyupgrade",
    )
    session.run("pre-commit", *args)
    if args and args[0] == "install":
        activate_virtualenv_in_precommit_hooks(session)


@session(python=default_python_version)
def safety(session: Session) -> None:
    """Scan dependencies for insecure packages."""
    requirements = session.poetry.export_requirements()
    session.install("-r", str(requirements))

    session.run("safety", "check", "--full-report", f"--file={requirements}")


@session(python=python_versions)
def mypy(session: Session) -> None:
    """Type-check using mypy."""
    args = session.posargs or ["src", "tests"]
    session.install(".")

    requirements = session.poetry.export_requirements()
    session.install("-r", str(requirements))

    session.run("mypy", *args)
    if not session.posargs:
        session.run("mypy", f"--python-executable={sys.executable}", "noxfile.py")


@session(python=python_versions)
def tests(session: Session) -> None:
    """Run the test suite."""
    session.install(".")

    requirements = session.poetry.export_requirements()
    session.install("-r", str(requirements))

    try:
        session.run("coverage", "run", "--parallel", "-m", "pytest", *session.posargs)
    finally:
        if session.interactive:
            session.notify("coverage", posargs=[])


@session(python=default_python_version)
def coverage(session: Session) -> None:
    """Produce the coverage report."""
    args = session.posargs or ["report"]
    requirements = session.poetry.export_requirements()
    session.install("-r", str(requirements))

    if not session.posargs and any(Path().glob(".coverage.*")):
        session.run("coverage", "combine")

    session.run("coverage", *args)


@session(python=default_python_version)
def typeguard(session: Session) -> None:
    """Runtime type checking using Typeguard."""
    session.install(".")
    requirements = session.poetry.export_requirements()
    session.install("-r", str(requirements))

    session.run("pytest", f"--typeguard-packages={package}", *session.posargs)

noxfile.py
nyno-ai/openai-token-counter
from typing import Any, Literal, Optional, Union

from .models import OpenAIRequest
from .token_counter import TokenCounter


def openai_token_counter(
    messages: list[dict[str, Any]],
    model: Optional[str] = None,
    functions: Optional[list[dict[str, Any]]] = None,
    function_call: Optional[
        Union[dict[Literal["name"], Any], Literal["auto"], Literal["none"]]
    ] = None,
) -> int:
    """Token counter function.

    Args:
        messages (dict[str, Any]): The messages to count tokens for.
        model (Optional[str]): The model to use for token counting.
        functions (Optional[list[dict[str, Any]]]): The functions to count tokens for.
        function_call (Optional[dict[str, Any]]): The function call to count tokens for.

    Returns:
        int: The number of tokens the prompt will use.
    """
    token_counter = TokenCounter(model=model)
    return token_counter.estimate_token_count(
        OpenAIRequest.model_validate(
            {
                "messages": messages,
                "functions": functions,
                "function_call": function_call,
            }
        )
    )

src/openai_token_counter/__init__.py
nyno-ai/openai-token-counter
from openai_token_counter.models import (
    ArrayProp,
    BoolProp,
    NullProp,
    NumberProp,
    ObjectProp,
    OpenAIFunction,
    PropItem,
    StringProp,
)


def format_function_definitions(functions: list[OpenAIFunction]) -> str:
    """Format the function definitions to a string.

    Args:
        functions (list[OpenAIFunction]): The list of functions to format.

    Returns:
        str: The formatted string.
    """
    lines = ["namespace functions {", ""]

    for f in functions:
        if f.description:
            lines.append(f"// {f.description}")

        parameters = f.parameters
        properties = parameters.properties if parameters else {}

        if properties is None:
            lines.append(f"type {f.name} = () => any;")
        elif len(properties) == 0:
            lines.append(f"type {f.name} = () => any;")

        else:
            lines.append(f"type {f.name} = (_: " + "{")
            lines.append(format_object_properties(f.parameters, 0))
            lines.append("}) => any;")

        lines.append("")

    lines.append("} // namespace functions")
    return "\n".join(lines)


def format_object_properties(obj: ObjectProp, indent: int) -> str:
    """Format the object properties to a string.

    Args:
        obj (ObjectProp): The object to format.
        indent (int): The indentation level.

    Returns:
        str: The formatted string.
    """
    if obj.properties is None:
        return ""

    lines = []

    required_params = obj.required or []

    for name, param in obj.properties.items():
        if param.description and indent < 2:
            lines.append(f"// {param.description}")

        if name in required_params:
            lines.append(f"{name}: {format_type(param, indent)},")
        else:
            lines.append(f"{name}?: {format_type(param, indent)},")

    return "\n".join([(" " * indent + line) for line in lines])


def format_type(param: PropItem, indent: int) -> str:
    """Format the type to a string.

    Args:
        param (PropItem): The parameter to format.
        indent (int): The indentation level.

    Returns:
        str: The formatted string.
    """
    if isinstance(param, StringProp):
        if param.enum:
            return " | ".join(f'"{v}"' for v in param.enum)  # noqa: B907
        return "string"

    if isinstance(param, NumberProp):
        if param.enum:
            return " | ".join(str(v) for v in param.enum)
        return "number"

    if isinstance(param, BoolProp):
        return "boolean"

    if isinstance(param, NullProp):
        return "null"

    if isinstance(param, ArrayProp):
        if param.items:
            return f"{format_type(param.items, indent)}[]"
        return "any[]"

    if isinstance(param, ObjectProp):
        return "{\n" + format_object_properties(param, indent + 2) + "\n}"

src/openai_token_counter/format.py
nyno-ai/openai-token-counter
from __future__ import annotations

from typing import Literal, Optional, Union

from pydantic import BaseModel


class StringProp(BaseModel):
    """This is the property object for the OpenAI functions."""

    type: Literal["string"]
    description: Optional[str] = None
    enum: Optional[list[str]] = None


class NumberProp(BaseModel):
    """This is the property object for the OpenAI functions."""

    type: Literal["integer", "number"]
    description: Optional[str] = None
    minimum: Optional[int] = None
    maximum: Optional[int] = None
    enum: Optional[list[str]] = None


class BoolProp(BaseModel):
    """This is the property object for the OpenAI functions."""

    type: Literal["boolean"]
    description: Optional[str] = None


class NullProp(BaseModel):
    """This is the property object for the OpenAI functions."""

    type: Literal["null"]
    description: Optional[str] = None


class ArrayProp(BaseModel):
    """This is the property object for the OpenAI functions."""

    type: Literal["array"]
    description: Optional[str] = None
    items: PropItem


class ObjectProp(BaseModel):
    """This is the property object for the OpenAI functions."""

    type: Literal["object"]
    description: Optional[str] = None
    required: Optional[list[str]] = None
    properties: Optional[dict[str, PropItem]] = None


PropItem = Union[StringProp, NumberProp, BoolProp, NullProp, ArrayProp, ObjectProp]


class OpenAIFunctionParameters(ObjectProp):
    """This is the parameters object for the OpenAI functions."""


class OpenAIFunction(BaseModel):
    """This is the function object for the OpenAI request."""

    name: str
    description: Optional[str] = None
    parameters: OpenAIFunctionParameters


class OpenAIFunctionCall(BaseModel):
    """This is the function call object that is used in messages."""

    name: str
    arguments: str  # Json string


class OpenAIMessage(BaseModel):
    """This is the message object for the OpenAI API."""

    role: str
    content: Optional[str] = None  # Optional in case of] function response
    name: Optional[str] = None
    function_call: Optional[OpenAIFunctionCall] = None


class OpenAIRequest(BaseModel):
    """This is the request object for the OpenAI API."""

    messages: list[OpenAIMessage]
    functions: Optional[list[OpenAIFunction]] = None

    # Function call can be either:
    # None: defaults to "auto" behind the scenes
    # "auto": Call or no call whatever function
    # "none": Don't call any function
    # dict {"name": "my_func"}: Call the function named "my_func"
    function_call: Optional[
        Union[Literal["auto", "none"], dict[Literal["name"], str]]
    ] = None

src/openai_token_counter/models.py
nyno-ai/openai-token-counter
from copy import deepcopy
from dataclasses import dataclass, field
from typing import Optional

from tiktoken import encoding_for_model, get_encoding

from openai_token_counter.format import format_function_definitions

from .models import OpenAIFunction, OpenAIMessage, OpenAIRequest


@dataclass
class TokenCounter:
    """Token counter class.

    Attributes:
        model (Optional[str]): The model to use for token counting.
    """

    model: Optional[str] = field(default=None)

    def estimate_token_count(self, request: OpenAIRequest) -> int:
        """Estimate the number of tokens a prompt will use.

        Args:
            request (OpenAIRequest): The request to estimate the token count for.

        Returns:
            int: An estimate for the number of tokens the prompt will use.
        """
        messages = request.messages
        functions = request.functions
        function_call = request.function_call

        padded_system = False
        tokens = 0

        for message in messages:
            message_copy = deepcopy(message)
            if message_copy.role == "system" and functions and not padded_system:
                if message_copy.content:
                    message_copy.content += "\n"
                padded_system = True
            tokens += self.estimate_tokens_in_messages(message_copy)

        # Each completion (vs message) seems to carry a 3-token overhead
        tokens += 3

        # If there are functions, add the function definitions as they count towards token usage
        if functions:
            tokens += self.estimate_tokens_in_functions(functions)

        # If there's a system message _and_ functions are present, subtract four tokens
        if functions and any(message.role == "system" for message in messages):
            tokens -= 4

        # If function_call is 'none', add one token.
        # If it's a OpenAIFunctionCall object, add 4 + the number of tokens in the function name.
        # If it's undefined or 'auto', don't add anything.
        if function_call and function_call != "auto":
            if function_call == "none":
                tokens += 1

            elif isinstance(function_call, dict) and "name" in function_call:
                tokens += self.string_tokens(function_call["name"]) + 4

        return tokens

    def string_tokens(self, string: str) -> int:
        """Get the token count for a string.

        Args:
            string (str): The string to count.

        Returns:
            int: The token count.
        """
        if self.model:
            encoding = encoding_for_model(self.model)
        else:
            encoding = get_encoding("cl100k_base")

        return len(encoding.encode(string))

    def estimate_tokens_in_messages(self, message: OpenAIMessage) -> int:
        """Estimate token count for a single message.

        Args:
            message (OpenAIMessage): The message to estimate the token count for.

        Returns:
            int: The estimated token count.
        """
        tokens = 0

        if message.role:
            tokens += self.string_tokens(message.role)

        if message.content:
            tokens += self.string_tokens(message.content)

        if message.name:
            tokens += self.string_tokens(message.name) + 1  # +1 for the name

        if message.function_call:
            if message.function_call.name:
                tokens += self.string_tokens(message.function_call.name)

            if message.function_call.arguments:
                tokens += self.string_tokens(message.function_call.arguments)

            tokens += 3  # Additional tokens for function call

        tokens += 3  # Add three per message

        if message.role == "function":
            tokens -= 2  # Subtract 2 if role is "function"

        return tokens

    def estimate_tokens_in_functions(self, function: list[OpenAIFunction]) -> int:
        """Estimate token count for the functions.

        We take here a list of functions and not one function because we use the format_function_definitions function
        which has to takes a list of functions as input, since the formatting is done for all functions.

        Args:
            function (list[OpenAIFunction]): The functions to estimate the token count for.

        Returns:
            int: The estimated token count.
        """
        prompt_definition = format_function_definitions(function)
        tokens = self.string_tokens(prompt_definition)
        tokens += 9  # Additional tokens for function definition
        return tokens

src/openai_token_counter/token_counter.py
nyno-ai/openai-token-counter
import os
from typing import TypedDict

import pytest
from dotenv import load_dotenv


class ConfigTests(TypedDict):
    """Config for testing."""

    OPENAI_API_KEY: str


@pytest.fixture(autouse=True, scope="session")
def config() -> ConfigTests:
    """Load the test environment variables.

    Returns:
        ConfigTests: The test config.
    """
    load_dotenv(".env.test")

    config = ConfigTests(
        {
            "OPENAI_API_KEY": os.environ["OPENAI_API_KEY"],
        }
    )
    return config

tests/conftest.py
nyno-ai/openai-token-counter
from typing import Any, Literal, Optional, TypedDict, Union


class TestCase(TypedDict, total=False):
    """Test case dict type."""

    messages: list[dict[str, Any]]
    tokens: int
    functions: Optional[list[dict[str, Any]]]
    function_call: Optional[
        Union[dict[Literal["name"], Any], Literal["auto"], Literal["none"]]
    ]


test_cases_raw: list[TestCase] = [
    {
        "messages": [
            {
                "role": "system",
                "content": "You are a helpful, pattern-following assistant that "
                "translates corporate jargon into plain English.",
            },
        ],
        "tokens": 25,
    },
    {
        "messages": [
            {
                "role": "system",
                "name": "example_user",
                "content": "New synergies will help drive top-line growth.",
            },
        ],
        "tokens": 20,
    },
    {
        "messages": [
            {
                "role": "system",
                "name": "example_assistant",
                "content": "Things working well together will increase revenue.",
            },
        ],
        "tokens": 19,
    },
    {
        "messages": [
            {
                "role": "system",
                "name": "example_user",
                "content": "Let's circle back when we have more bandwidth "
                "to touch base on opportunities for increased leverage.",
            },
        ],
        "tokens": 28,
    },
    {
        "messages": [
            {
                "role": "system",
                "name": "example_assistant",
                "content": "Let's talk later when we're less busy about how to do better.",
            },
        ],
        "tokens": 26,
    },
    {
        "messages": [
            {
                "role": "user",
                "content": "This late pivot means we don't have time to boil the ocean for the client deliverable.",
            },
        ],
        "tokens": 26,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "tokens": 8,
    },
    {
        "messages": [{"role": "user", "content": "hello world"}],
        "tokens": 9,
    },
    {
        "messages": [{"role": "system", "content": "hello"}],
        "tokens": 8,
    },
    {
        "messages": [{"role": "system", "content": "hello:"}],
        "tokens": 9,
    },
    {
        "messages": [
            {"role": "system", "content": "# Important: you're the best robot"},
            {"role": "user", "content": "hello robot"},
            {"role": "assistant", "content": "hello world"},
        ],
        "tokens": 27,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "foo",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "tokens": 31,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "foo",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "function_call": "none",
        "tokens": 32,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "foo",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "function_call": "auto",
        "tokens": 31,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "foo",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "function_call": {"name": "foo"},
        "tokens": 36,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "foo",
                "description": "Do a foo",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "tokens": 36,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "bing_bong",
                "description": "Do a bing bong",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "foo": {"type": "string"},
                    },
                },
            },
        ],
        "tokens": 49,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "bing_bong",
                "description": "Do a bing bong",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "foo": {"type": "string"},
                        "bar": {"type": "number", "description": "A number"},
                    },
                },
            },
        ],
        "tokens": 57,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "bing_bong",
                "description": "Do a bing bong",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "foo": {
                            "type": "object",
                            "properties": {
                                "bar": {"type": "string", "enum": ["a", "b", "c"]},
                                "baz": {"type": "boolean"},
                            },
                        },
                    },
                },
            },
        ],
        "tokens": 68,
    },
    {
        "messages": [
            {"role": "user", "content": "hello world"},
            {"role": "function", "name": "do_stuff", "content": "{}"},
        ],
        "tokens": 15,
    },
    {
        "messages": [
            {"role": "user", "content": "hello world"},
            {
                "role": "function",
                "name": "do_stuff",
                "content": '{"foo": "bar", "baz": 1.5}',
            },
        ],
        "tokens": 28,
    },
    {
        "messages": [
            {
                "role": "function",
                "name": "dance_the_tango",
                "content": '{"a": { "b" : { "c": false}}}',
            },
        ],
        "tokens": 24,
    },
    {
        "messages": [
            {
                "role": "assistant",
                "content": "",
                "function_call": {
                    "name": "do_stuff",
                    "arguments": '{"foo": "bar", "baz": 1.5}',
                },
            },
        ],
        "tokens": 26,
    },
    {
        "messages": [
            {
                "role": "assistant",
                "content": "",
                "function_call": {
                    "name": "do_stuff",
                    "arguments": '{"foo":"bar", "baz":\n\n 1.5}',
                },
            },
        ],
        "tokens": 25,
    },
    {
        "messages": [
            {"role": "system", "content": "Hello"},
            {"role": "user", "content": "Hi there"},
        ],
        "functions": [
            {
                "name": "do_stuff",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "tokens": 35,
    },
    {
        "messages": [
            {"role": "system", "content": "Hello:"},
            {"role": "user", "content": "Hi there"},
        ],
        "functions": [
            {
                "name": "do_stuff",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "tokens": 35,
    },
    {
        "messages": [
            {"role": "system", "content": "Hello:"},
            {"role": "system", "content": "Hello"},
            {"role": "user", "content": "Hi there"},
        ],
        "functions": [
            {
                "name": "do_stuff",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "tokens": 40,
    },
    {
        "messages": [
            {"role": "system", "content": "Hello:"},
            {"role": "system", "content": "Hello"},
            {"role": "user", "content": "Hi there"},
        ],
        "functions": [
            {
                "name": "do_stuff",
                "parameters": {"type": "object", "properties": {}},
            },
            {
                "name": "do_other_stuff",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "tokens": 49,
    },
    {
        "messages": [
            {"role": "system", "content": "Hello:"},
            {"role": "system", "content": "Hello"},
            {"role": "user", "content": "Hi there"},
        ],
        "functions": [
            {
                "name": "do_stuff",
                "parameters": {"type": "object", "properties": {}},
            },
            {
                "name": "do_other_stuff",
                "parameters": {"type": "object", "properties": {}},
            },
        ],
        "function_call": {"name": "do_stuff"},
        "tokens": 55,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "get_recipe",
                "parameters": {
                    "type": "object",
                    "required": ["ingredients", "instructions", "time_to_cook"],
                    "properties": {
                        "ingredients": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "required": ["name", "unit", "amount"],
                                "properties": {
                                    "name": {
                                        "type": "string",
                                    },
                                    "unit": {
                                        "enum": [
                                            "grams",
                                            "ml",
                                            "cups",
                                            "pieces",
                                            "teaspoons",
                                        ],
                                        "type": "string",
                                    },
                                    "amount": {
                                        "type": "number",
                                    },
                                },
                            },
                        },
                        "instructions": {
                            "type": "array",
                            "items": {
                                "type": "string",
                            },
                            "description": "Steps to prepare the recipe (no numbering)",
                        },
                        "time_to_cook": {
                            "type": "number",
                            "description": "Total time to prepare the recipe in minutes",
                        },
                    },
                },
            },
        ],
        "tokens": 106,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "function",
                "description": "description",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "quality": {
                            "type": "object",
                            "properties": {
                                "pros": {
                                    "type": "array",
                                    "items": {
                                        "type": "string",
                                    },
                                    "description": "Write 3 points why this text is well written",
                                },
                            },
                        },
                    },
                },
            },
        ],
        "tokens": 46,
    },
    {
        "messages": [{"role": "user", "content": "hello"}],
        "functions": [
            {
                "name": "function",
                "description": "desctiption1",
                "parameters": {
                    "type": "object",
                    "description": "desctiption2",
                    "properties": {
                        "mainField": {
                            "type": "string",
                            "description": "description3",
                        },
                        "field number one": {
                            "type": "object",
                            "description": "description4",
                            "properties": {
                                "yesNoField": {
                                    "type": "string",
                                    "description": "description5",
                                    "enum": ["Yes", "No"],
                                },
                                "howIsInteresting": {
                                    "type": "string",
                                    "description": "description6",
                                },
                                "scoreInteresting": {
                                    "type": "number",
                                    "description": "description7",
                                },
                                "isInteresting": {
                                    "type": "string",
                                    "description": "description8",
                                    "enum": ["Yes", "No"],
                                },
                            },
                        },
                    },
                },
            },
        ],
        "tokens": 96,
    },
]

tests/counter/resources.py
nyno-ai/openai-token-counter
import time

import openai
from openai.error import APIConnectionError, APIError, ServiceUnavailableError

from openai_token_counter import openai_token_counter
from tests.conftest import ConfigTests
from tests.counter.resources import test_cases_raw


MAX_SERVICE_UNAVAILABLE_RETRY_ATTEMPTS = 5
SLEEP_INTERVAL_BETWEEN_ATTEMPTS = 5
MODEL = "gpt-3.5-turbo"
MODEL_PROMPT_TOKEN_COST_PER_TOKEN = 0.0015 / 1000
MODEL_COMPLETION_TOKEN_COST_PER_TOKEN = 0.002 / 1000
MAX_RESPONSE_TOKENS = 1  # The response tokens doesn't matter in this context, we only calculate request tokens


def test_token_counter(config: ConfigTests) -> None:
    """Test that the token counter works as expected."""
    token_usage = {"prompt": 0, "completion": 0}

    for test_case in test_cases_raw:
        for attempt in range(1, MAX_SERVICE_UNAVAILABLE_RETRY_ATTEMPTS + 1):
            try:
                optional_params = {
                    "functions": test_case.get("functions"),
                    "function_call": test_case.get("function_call"),
                }
                params = {
                    "api_key": config["OPENAI_API_KEY"],
                    "model": MODEL,
                    "max_tokens": MAX_RESPONSE_TOKENS,
                    "messages": test_case["messages"],
                    **{k: v for k, v in optional_params.items() if v is not None},
                }

                response = openai.ChatCompletion.create(**params)

                calculated_tokens = openai_token_counter(
                    messages=test_case["messages"],
                    model=MODEL,
                    functions=test_case.get("functions"),
                    function_call=test_case.get("function_call"),
                )

                token_usage["prompt"] += response["usage"]["prompt_tokens"]
                token_usage["completion"] += response["usage"]["completion_tokens"]

                assert response["usage"]["prompt_tokens"] == test_case["tokens"]
                assert calculated_tokens == test_case["tokens"]

            except (ServiceUnavailableError, APIConnectionError, APIError) as err:
                if attempt >= MAX_SERVICE_UNAVAILABLE_RETRY_ATTEMPTS:
                    raise Exception(
                        f"Failed to get response from OpenAI API after {attempt} attempts"
                    ) from err
                print(
                    f"Service unavailable, retrying in {SLEEP_INTERVAL_BETWEEN_ATTEMPTS} seconds..."
                )
                time.sleep(SLEEP_INTERVAL_BETWEEN_ATTEMPTS)

    print(f"Total token usage: {token_usage}")
    cost = (
        token_usage["prompt"] * MODEL_PROMPT_TOKEN_COST_PER_TOKEN
        + token_usage["completion"] * MODEL_COMPLETION_TOKEN_COST_PER_TOKEN
    )

    print(f"Test Costs: {cost}$")
    print("Thank you")

tests/counter/test_token_counter.py
nyno-ai/openai-token-counter
from openai_token_counter.format import format_function_definitions
from openai_token_counter.models import OpenAIFunction


def test_format_function_definitions() -> None:
    """Test that the format_function_definitions function works as expected."""
    functions: list[OpenAIFunction] = [
        OpenAIFunction.model_validate(
            {
                "name": "function1",
                "description": "This is function 1",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "param1": {"type": "string"},
                        "param2": {"type": "number"},
                    },
                    "required": ["param1"],
                },
            }
        ),
        OpenAIFunction.model_validate(
            {
                "name": "function2",
                "description": "This is function 2",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "param3": {"type": "boolean"},
                        "param4": {"type": "null"},
                    },
                },
            }
        ),
        OpenAIFunction.model_validate(
            {
                "name": "function3",
                "description": "This is function 3",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "param5": {"type": "boolean"},
                        "param6": {"type": "null"},
                        "param7": {
                            "type": "object",
                            "properties": {
                                "param8": {"type": "boolean"},
                                "param9": {"type": "null"},
                            },
                            "required": ["param8"],
                        },
                    },
                },
            }
        ),
    ]

    result = format_function_definitions(functions)
    assert (
        result
        == """
namespace functions {

// This is function 1
type function1 = (_: {
param1: string,
param2?: number,
}) => any;

// This is function 2
type function2 = (_: {
param3?: boolean,
param4?: null,
}) => any;

// This is function 3
type function3 = (_: {
param5?: boolean,
param6?: null,
param7?: {
  param8: boolean,
  param9?: null,
},
}) => any;

} // namespace functions
""".strip()
    )

tests/format/test_format.py
nyno-ai/openai-token-counter