Add runtime type checking

This patch adds beartype for runtime type checking. This gives us the
best of both worlds: we do static type checking of our own library with
mypy, and we export our static types, but for clients who do not run
static type checking of their own code, runtime type checking in our
library can help them catch bugs earlier.
`tests/test_codex_tool.py::test_bad_argument_type` serves as an example:
this fails at initialization time of `CodexTool`, whereas without
runtime type checking, this would fail later (e.g., when the user calls
the `query` method on the object).

Because we're performing runtime type checking, some of the imports that
were behind `if TYPE_CHECKING` flags have to be moved to runtime. This
patch updates the linter config to allow imports that are only used for
type checking.

This patch also switches to consistent `from __future__ import
annotations` everywhere, stops using type hints deprecated by PEP 585,
and uses PEP 585 / PEP 604 syntax everywhere. This patch updates the
linter config to match this style.

beartype relies on `isinstance` for runtime type checks, which needs to
be taken into account when using mocks by overriding the `__class__`
attribute. This patch updates the tests accordingly.

Author: anishathalye

pyproject.toml

@@ -27,6 +27,7 @@ classifiers = [
 dependencies = [
   "codex-sdk==0.1.0a9",
   "pydantic>=1.9.0, <3",
+  "beartype>=0.17.0",
 ]
 
 [project.urls]
@@ -98,4 +99,8 @@ html = "coverage html"
 xml = "coverage xml"
 
 [tool.ruff.lint]
-ignore = ["FA100", "UP007", "UP006"]
+ignore = [
+  "TCH001", # this package does run-time type checking
+  "TCH002",
+  "TCH003"
+]

src/cleanlab_codex/__init__.py

@@ -1,4 +1,11 @@
 # SPDX-License-Identifier: MIT
+
+from beartype.claw import beartype_this_package
+
+# this must run before any other imports from the cleanlab_codex package
+beartype_this_package()
+
+# ruff: noqa: E402
 from cleanlab_codex.codex import Codex
 from cleanlab_codex.codex_tool import CodexTool
 

src/cleanlab_codex/codex.py

@@ -1,13 +1,9 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
-
 from cleanlab_codex.internal.project import create_project, query_project
 from cleanlab_codex.internal.utils import init_codex_client
-
-if TYPE_CHECKING:
-    from cleanlab_codex.types.entry import Entry, EntryCreate
-    from cleanlab_codex.types.organization import Organization
+from cleanlab_codex.types.entry import Entry, EntryCreate
+from cleanlab_codex.types.organization import Organization
 
 
 class Codex:
@@ -41,7 +37,7 @@ def list_organizations(self) -> list[Organization]:
         """
         return self._client.users.myself.organizations.list().organizations
 
-    def create_project(self, name: str, organization_id: str, description: Optional[str] = None) -> str:
+    def create_project(self, name: str, organization_id: str, description: str | None = None) -> str:
         """Create a new Codex project.
 
         Args:
@@ -77,7 +73,7 @@ def create_project_access_key(
         self,
         project_id: str,
         access_key_name: str,
-        access_key_description: Optional[str] = None,
+        access_key_description: str | None = None,
     ) -> str:
         """Create a new access key for a project.
 
@@ -99,15 +95,15 @@ def query(
         self,
         question: str,
         *,
-        project_id: Optional[str] = None,  # TODO: update to uuid once project IDs are changed to UUIDs
-        fallback_answer: Optional[str] = None,
+        project_id: str | None = None,  # TODO: update to uuid once project IDs are changed to UUIDs
+        fallback_answer: str | None = None,
         read_only: bool = False,
-    ) -> tuple[Optional[str], Optional[Entry]]:
+    ) -> tuple[str | None, Entry | None]:
         """Query Codex to check if the Codex project contains an answer to this question and add the question to the Codex project for SME review if it does not.
 
         Args:
             question (str): The question to ask the Codex API.
-            project_id (:obj:`int`, optional): The ID of the project to query.
+            project_id (:obj:`str`, optional): The ID of the project to query.
                 If the client is authenticated with a user-level API Key, this is required.
                 If the client is authenticated with a project-level Access Key, this is optional. The client will use the Access Key's project ID by default.
             fallback_answer (:obj:`str`, optional): Optional fallback answer to return if Codex is unable to answer the question.

src/cleanlab_codex/codex_tool.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, ClassVar, Optional
+from typing import Any, ClassVar
 
 from cleanlab_codex.codex import Codex
 
@@ -23,8 +23,8 @@ def __init__(
         self,
         codex_client: Codex,
         *,
-        project_id: Optional[str] = None,
-        fallback_answer: Optional[str] = DEFAULT_FALLBACK_ANSWER,
+        project_id: str | None = None,
+        fallback_answer: str | None = DEFAULT_FALLBACK_ANSWER,
     ):
         self._codex_client = codex_client
         self._project_id = project_id
@@ -35,8 +35,8 @@ def from_access_key(
         cls,
         access_key: str,
         *,
-        project_id: Optional[str] = None,
-        fallback_answer: Optional[str] = DEFAULT_FALLBACK_ANSWER,
+        project_id: str | None = None,
+        fallback_answer: str | None = DEFAULT_FALLBACK_ANSWER,
     ) -> CodexTool:
         """Creates a CodexTool from an access key. The project ID that the CodexTool will use is the one that is associated with the access key."""
         return cls(
@@ -50,8 +50,8 @@ def from_client(
         cls,
         codex_client: Codex,
         *,
-        project_id: Optional[str] = None,
-        fallback_answer: Optional[str] = DEFAULT_FALLBACK_ANSWER,
+        project_id: str | None = None,
+        fallback_answer: str | None = DEFAULT_FALLBACK_ANSWER,
     ) -> CodexTool:
         """Creates a CodexTool from a Codex client.
         If the Codex client is initialized with a project access key, the CodexTool will use the project ID that is associated with the access key.
@@ -74,16 +74,16 @@ def tool_description(self) -> str:
         return self._tool_description
 
     @property
-    def fallback_answer(self) -> Optional[str]:
+    def fallback_answer(self) -> str | None:
         """The fallback answer to use if the Codex project cannot answer the question."""
         return self._fallback_answer
 
     @fallback_answer.setter
-    def fallback_answer(self, value: Optional[str]) -> None:
+    def fallback_answer(self, value: str | None) -> None:
         """Sets the fallback answer to use if the Codex project cannot answer the question."""
         self._fallback_answer = value
 
-    def query(self, question: str) -> Optional[str]:
+    def query(self, question: str) -> str | None:
         """Asks an all-knowing advisor this question in cases where it cannot be answered from the provided Context. If the answer is not available, this returns a fallback answer or None.
 
         Args:

src/cleanlab_codex/internal/project.py

@@ -1,12 +1,8 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
-
-if TYPE_CHECKING:
-    from codex import Codex as _Codex
-
-    from cleanlab_codex.types.entry import Entry
+from codex import Codex as _Codex
 
+from cleanlab_codex.types.entry import Entry
 from cleanlab_codex.types.project import ProjectConfig
 
 
@@ -17,7 +13,7 @@ def __str__(self) -> str:
         return "project_id is required when authenticating with a user-level API Key"
 
 
-def create_project(client: _Codex, name: str, organization_id: str, description: Optional[str] = None) -> str:
+def create_project(client: _Codex, name: str, organization_id: str, description: str | None = None) -> str:
     project = client.projects.create(
         config=ProjectConfig(),
         organization_id=organization_id,
@@ -31,10 +27,10 @@ def query_project(
     client: _Codex,
     question: str,
     *,
-    project_id: Optional[str] = None,
-    fallback_answer: Optional[str] = None,
+    project_id: str | None = None,
+    fallback_answer: str | None = None,
     read_only: bool = False,
-) -> tuple[Optional[str], Optional[Entry]]:
+) -> tuple[str | None, Entry | None]:
     if client.access_key is not None:
         project_id = client.projects.access_keys.retrieve_project_id().project_id
     elif project_id is None:

src/cleanlab_codex/utils/llamaindex.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Callable
 from inspect import signature
-from typing import Any, Callable
+from typing import Any
 
 from llama_index.core.bridge.pydantic import BaseModel, FieldInfo, create_model
 

src/cleanlab_codex/utils/openai.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, Dict, List, Literal
+from typing import Any, Literal
 
 from pydantic import BaseModel
 
@@ -12,8 +12,8 @@ class Property(BaseModel):
 
 class FunctionParameters(BaseModel):
     type: Literal["object"] = "object"
-    properties: Dict[str, Property]
-    required: List[str]
+    properties: dict[str, Property]
+    required: list[str]
 
 
 class Function(BaseModel):
@@ -30,9 +30,9 @@ class Tool(BaseModel):
 def format_as_openai_tool(
     tool_name: str,
     tool_description: str,
-    tool_properties: Dict[str, Any],
-    required_properties: List[str],
-) -> Dict[str, Any]:
+    tool_properties: dict[str, Any],
+    required_properties: list[str],
+) -> dict[str, Any]:
     return Tool(
         function=Function(
             name=tool_name,

src/cleanlab_codex/utils/smolagents.py

@@ -1,15 +1,17 @@
-from typing import Callable, Dict, Optional
+from __future__ import annotations
+
+from collections.abc import Callable
 
 from smolagents import Tool  # type: ignore
 
 
 class CodexTool(Tool):  # type: ignore[misc]
     def __init__(
         self,
-        query: Callable[[str], Optional[str]],
+        query: Callable[[str], str | None],
         tool_name: str,
         tool_description: str,
-        inputs: Dict[str, Dict[str, str]],
+        inputs: dict[str, dict[str, str]],
     ):
         super().__init__()
         self._query = query
@@ -18,5 +20,5 @@ def __init__(
         self.inputs = inputs
         self.output_type = "string"
 
-    def forward(self, question: str) -> Optional[str]:
+    def forward(self, question: str) -> str | None:
         return self._query(question)

tests/fixtures/client.py

@@ -1,12 +1,15 @@
-from typing import Generator
+from collections.abc import Generator
 from unittest.mock import MagicMock, patch
 
 import pytest
 
+from cleanlab_codex.internal.utils import _Codex
+
 
 @pytest.fixture
 def mock_client() -> Generator[MagicMock, None, None]:
     with patch("cleanlab_codex.codex.init_codex_client") as mock_init:
         mock_client = MagicMock()
+        mock_client.__class__ = _Codex
         mock_init.return_value = mock_client
         yield mock_client

tests/internal/test_utils.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from cleanlab_codex.internal.utils import MissingAuthKeyError, init_codex_client, is_access_key
+from cleanlab_codex.internal.utils import MissingAuthKeyError, _Codex, init_codex_client, is_access_key
 
 DUMMY_ACCESS_KEY = "sk-1-EMOh6UrRo7exTEbEi8_azzACAEdtNiib2LLa1IGo6kA"
 DUMMY_API_KEY = "GP0FzPfA7wYy5L64luII2YaRT2JoSXkae7WEo7dH6Bw"
@@ -16,6 +16,7 @@ def test_is_access_key() -> None:
 
 def test_init_codex_client_access_key() -> None:
     mock_client = MagicMock()
+    mock_client.__class__ = _Codex
     with patch("cleanlab_codex.internal.utils._Codex", autospec=True, return_value=mock_client) as mock_init:
         mock_client.projects.access_keys.retrieve_project_id.return_value = "test_project_id"
         client = init_codex_client(DUMMY_ACCESS_KEY)
@@ -25,6 +26,7 @@ def test_init_codex_client_access_key() -> None:
 
 def test_init_codex_client_api_key() -> None:
     mock_client = MagicMock()
+    mock_client.__class__ = _Codex
     with patch("cleanlab_codex.internal.utils._Codex", autospec=True, return_value=mock_client) as mock_init:
         mock_client.users.myself.api_key.retrieve.return_value = "test_project_id"
         client = init_codex_client(DUMMY_API_KEY)
@@ -40,6 +42,7 @@ def test_init_codex_client_no_key() -> None:
 def test_init_codex_client_access_key_env_var() -> None:
     with patch.dict(os.environ, {"CODEX_ACCESS_KEY": DUMMY_ACCESS_KEY}):
         mock_client = MagicMock()
+        mock_client.__class__ = _Codex
         with patch(
             "cleanlab_codex.internal.utils._Codex",
             autospec=True,
@@ -54,6 +57,7 @@ def test_init_codex_client_access_key_env_var() -> None:
 def test_init_codex_client_api_key_env_var() -> None:
     with patch.dict(os.environ, {"CODEX_API_KEY": DUMMY_API_KEY}):
         mock_client = MagicMock()
+        mock_client.__class__ = _Codex
         with patch(
             "cleanlab_codex.internal.utils._Codex",
             autospec=True,

tests/test_codex.py

@@ -86,6 +86,9 @@ def test_create_project_access_key(mock_client: MagicMock) -> None:
     codex = Codex("")
     access_key_name = "Test Access Key"
     access_key_description = "Test Access Key Description"
+    access_key = MagicMock()
+    access_key.token.__class__ = str
+    mock_client.projects.access_keys.create.return_value = access_key
     codex.create_project_access_key(FAKE_PROJECT_ID, access_key_name, access_key_description)
     mock_client.projects.access_keys.create.assert_called_once_with(
         project_id=FAKE_PROJECT_ID,

tests/test_codex_tool.py

@@ -34,3 +34,10 @@ def test_to_smolagents_tool(mock_client: MagicMock) -> None:  # noqa: ARG001
     assert isinstance(smolagents_tool, Tool)
     assert smolagents_tool.name == tool.tool_name
     assert smolagents_tool.description == tool.tool_description
+
+
+def test_bad_argument_type() -> None:
+    from beartype.roar import BeartypeException
+
+    with pytest.raises(BeartypeException):
+        CodexTool("asdf")