Add py.typed stub

Author: pamelafox

.github/workflows/python.yaml

@@ -31,3 +31,5 @@ jobs:
         - name: Run unit tests
           run: |
             python3 -m pytest -s -vv --cov --cov-fail-under=99
+        - name: Run type checks
+          run: mypy .

.vscode/settings.json

@@ -3,5 +3,12 @@
         "tests"
     ],
     "python.testing.unittestEnabled": false,
-    "python.testing.pytestEnabled": true
+    "python.testing.pytestEnabled": true,
+    "files.exclude": {
+        ".coverage": true,
+        ".pytest_cache": true,
+        "__pycache__": true,
+        ".ruff_cache": true,
+        ".mypy_cache": true,
+    }
 }

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 All notable changes to this project will be documented in this file.
 
+## [0.1.1] - May 2, 2024
+
+- Add `py.typed` file so that mypy can find the type hints in this package.
+
 ## [0.1.0] - May 2, 2024
 
 - Add `count_tokens_for_system_and_tools` to count tokens for system message and tools. You should count the tokens for both together, since the token count for tools varies based off whether a system message is provided.

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "openai-messages-token-helper"
 description = "A helper library for estimating tokens used by messages sent through OpenAI Chat Completions API."
-version = "0.1.0"
+version = "0.1.1"
 authors = [{name = "Pamela Fox"}]
 requires-python = ">=3.9"
 readme = "README.md"
@@ -33,7 +33,8 @@ dev = [
     "black",
     "flit",
     "azure-identity",
-    "python-dotenv"
+    "python-dotenv",
+    "mypy"
 ]
 
 [build-system]

src/openai_messages_token_helper/model_helper.py

@@ -113,7 +113,7 @@ def count_tokens_for_message(model: str, message: Mapping[str, object], default_
 
 def count_tokens_for_system_and_tools(
     model: str,
-    system_message: dict | None = None,
+    system_message: Mapping[str, object] | None = None,
     tools: list[dict[str, dict]] | None = None,
     tool_choice: str | dict | None = None,
     default_to_cl100k: bool = False,

src/py.typed

@@ -3,29 +3,44 @@
 import azure.identity
 import openai
 from dotenv import load_dotenv
-from functions import FUNCTION_COUNTS
+from functions import FUNCTION_COUNTS  # type: ignore[import-not-found]
 
 # Setup the OpenAI client to use either Azure OpenAI or OpenAI API
 load_dotenv()
 API_HOST = os.getenv("API_HOST")
 
+client: openai.OpenAI | openai.AzureOpenAI
+
 if API_HOST == "azure":
+
+    if (azure_openai_version := os.getenv("AZURE_OPENAI_VERSION")) is None:
+        raise ValueError("Missing Azure OpenAI version")
+    if (azure_openai_endpoint := os.getenv("AZURE_OPENAI_ENDPOINT")) is None:
+        raise ValueError("Missing Azure OpenAI endpoint")
+    if (azure_openai_deployment := os.getenv("AZURE_OPENAI_DEPLOYMENT")) is None:
+        raise ValueError("Missing Azure OpenAI deployment")
+
     token_provider = azure.identity.get_bearer_token_provider(
         azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
     )
     client = openai.AzureOpenAI(
-        api_version=os.getenv("AZURE_OPENAI_VERSION"),
-        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+        api_version=azure_openai_version,
+        azure_endpoint=azure_openai_endpoint,
         azure_ad_token_provider=token_provider,
     )
-    MODEL_NAME = os.getenv("AZURE_OPENAI_DEPLOYMENT")
+    MODEL_NAME = azure_openai_deployment
 else:
-    client = openai.OpenAI(api_key=os.getenv("OPENAI_KEY"))
-    MODEL_NAME = os.getenv("OPENAI_MODEL")
+    if (openai_key := os.getenv("OPENAI_KEY")) is None:
+        raise ValueError("Missing OpenAI API key")
+    if (openai_model := os.getenv("OPENAI_MODEL")) is None:
+        raise ValueError("Missing OpenAI model")
+    client = openai.OpenAI(api_key=openai_key)
+    MODEL_NAME = openai_model
+
 
 # Test the token count for each message
 for function_count_pair in FUNCTION_COUNTS:
-    response = client.chat.completions.create(
+    response = client.chat.completions.create(  # type: ignore[call-overload]
         model=MODEL_NAME,
         temperature=0.7,
         n=1,
@@ -35,6 +50,7 @@
     )
 
     print(function_count_pair["tools"])
+    assert response.usage is not None, "Expected usage to be present"
     assert (
         response.usage.prompt_tokens == function_count_pair["count"]
     ), f"Expected {function_count_pair['count']} tokens, got {response.usage.prompt_tokens}"

tests/verify_functions.py

@@ -3,37 +3,51 @@
 import azure.identity
 import openai
 from dotenv import load_dotenv
-from messages import MESSAGE_COUNTS
+from messages import MESSAGE_COUNTS  # type: ignore[import-not-found]
 
 # Setup the OpenAI client to use either Azure OpenAI or OpenAI API
 load_dotenv()
 API_HOST = os.getenv("API_HOST")
 
+client: openai.OpenAI | openai.AzureOpenAI
+
 if API_HOST == "azure":
+    if (azure_openai_version := os.getenv("AZURE_OPENAI_VERSION")) is None:
+        raise ValueError("Missing Azure OpenAI version")
+    if (azure_openai_endpoint := os.getenv("AZURE_OPENAI_ENDPOINT")) is None:
+        raise ValueError("Missing Azure OpenAI endpoint")
+    if (azure_openai_deployment := os.getenv("AZURE_OPENAI_DEPLOYMENT")) is None:
+        raise ValueError("Missing Azure OpenAI deployment")
+
     token_provider = azure.identity.get_bearer_token_provider(
         azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
     )
     client = openai.AzureOpenAI(
-        api_version=os.getenv("AZURE_OPENAI_VERSION"),
-        azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+        api_version=azure_openai_version,
+        azure_endpoint=azure_openai_endpoint,
         azure_ad_token_provider=token_provider,
     )
-    MODEL_NAME = os.getenv("AZURE_OPENAI_DEPLOYMENT")
+    MODEL_NAME = azure_openai_deployment
 else:
-    client = openai.OpenAI(api_key=os.getenv("OPENAI_KEY"))
-    MODEL_NAME = os.getenv("OPENAI_MODEL")
+    if (openai_key := os.getenv("OPENAI_KEY")) is None:
+        raise ValueError("Missing OpenAI API key")
+    if (openai_model := os.getenv("OPENAI_MODEL")) is None:
+        raise ValueError("Missing OpenAI model")
+    client = openai.OpenAI(api_key=openai_key)
+    MODEL_NAME = openai_model
 
 # Test the token count for each message
 for message_count_pair in MESSAGE_COUNTS:
     response = client.chat.completions.create(
         model=MODEL_NAME,
         temperature=0.7,
         n=1,
-        messages=[message_count_pair["message"]],
+        messages=[message_count_pair["message"]],  # type: ignore[list-item]
     )
 
     print(message_count_pair["message"])
     expected_tokens = message_count_pair["count"]
+    assert response.usage is not None, "Expected usage to be present"
     assert (
         response.usage.prompt_tokens == expected_tokens
     ), f"Expected {expected_tokens} tokens, got {response.usage.prompt_tokens}"