Merge branch 'main' into codex-as-backup

cleanlab · Feb 7, 2025 · 3286674 · 3286674
2 parents 2630a2c + cd6f2f8
commit 3286674
Show file tree

Hide file tree

Showing 10 changed files with 120 additions and 24 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.0.1a4] - 2025-02-07
+
+### Fixed
+
+- Pydantic model validation error when querying Project and listing Organizations.
+
 ## [0.0.1a3] - 2025-02-06
 
 ### Added
@@ -27,7 +33,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Pre-release of the Cleanlab Codex Python client.
 
-[Unreleased]: https://github.com/cleanlab/cleanlab-codex/compare/v0.0.1a3...HEAD
+[Unreleased]: https://github.com/cleanlab/cleanlab-codex/compare/v0.0.1a4...HEAD
+[0.0.1a4]: https://github.com/cleanlab/cleanlab-codex/compare/v0.0.1a3...v0.0.1a4
 [0.0.1a3]: https://github.com/cleanlab/cleanlab-codex/compare/v0.0.1a2...v0.0.1a3
 [0.0.1a2]: https://github.com/cleanlab/cleanlab-codex/compare/v0.0.1a1...v0.0.1a2
 [0.0.1a1]: https://github.com/cleanlab/cleanlab-codex/compare/267a93300f77c94e215d7697223931e7926cad9e...v0.0.1a1
diff --git a/pyproject.toml b/pyproject.toml
@@ -26,7 +26,7 @@ classifiers = [
 ]
 dependencies = [
   "codex-sdk==0.1.0a9",
-  "pydantic>=1.9.0, <3",
+  "pydantic>=2.0.0, <3",
 ]
 
 [project.urls]

diff --git a/src/cleanlab_codex/__about__.py b/src/cleanlab_codex/__about__.py
@@ -1,2 +1,2 @@
 # SPDX-License-Identifier: MIT
-__version__ = "0.0.1a3"
+__version__ = "0.0.1a4"
diff --git a/src/cleanlab_codex/codex_tool.py b/src/cleanlab_codex/codex_tool.py
@@ -126,7 +126,7 @@ def to_smolagents_tool(self) -> Any:
         try:
             from cleanlab_codex.utils.smolagents import CodexTool as SmolagentsCodexTool
         except ImportError as e:
-            raise MissingDependencyError("smolagents", "https://github.com/huggingface/smolagents") from e
+            raise MissingDependencyError(e.name or "smolagents", "https://github.com/huggingface/smolagents") from e
 
         return SmolagentsCodexTool(
             query=self.query,
@@ -145,8 +145,9 @@ def to_llamaindex_tool(self) -> Any:
 
         except ImportError as e:
             raise MissingDependencyError(
-                "llama-index-core",
-                "https://docs.llamaindex.ai/en/stable/getting_started/installation/",
+                import_name=e.name or "llama_index",
+                package_name="llama-index-core",
+                package_url="https://docs.llamaindex.ai/en/stable/getting_started/installation/",
             ) from e
 
         return FunctionTool.from_defaults(
@@ -165,7 +166,11 @@ def to_langchain_tool(self) -> Any:
             from langchain_core.tools.structured import StructuredTool
 
         except ImportError as e:
-            raise MissingDependencyError("langchain", "https://pypi.org/project/langchain/") from e
+            raise MissingDependencyError(
+                import_name=e.name or "langchain",
+                package_name="langchain",
+                package_url="https://pypi.org/project/langchain/",
+            ) from e
 
         return StructuredTool.from_function(
             func=self.query,

diff --git a/src/cleanlab_codex/internal/organization.py b/src/cleanlab_codex/internal/organization.py
@@ -9,4 +9,6 @@
 
 
 def list_organizations(client: _Codex) -> list[Organization]:
-    return [Organization.model_validate(org) for org in client.users.myself.organizations.list().organizations]
+    return [
+        Organization.model_validate(org.model_dump()) for org in client.users.myself.organizations.list().organizations
+    ]
diff --git a/src/cleanlab_codex/internal/project.py b/src/cleanlab_codex/internal/project.py
@@ -25,14 +25,16 @@ def query_project(
 ) -> tuple[Optional[str], Optional[Entry]]:
     maybe_entry = client.projects.entries.query(project_id, question=question)
     if maybe_entry is not None:
-        entry = Entry.model_validate(maybe_entry)
+        entry = Entry.model_validate(maybe_entry.model_dump())
         if entry.answer is not None:
             return entry.answer, entry
 
         return fallback_answer, entry
 
     if not read_only:
-        created_entry = Entry.model_validate(client.projects.entries.add_question(project_id, question=question))
+        created_entry = Entry.model_validate(
+            client.projects.entries.add_question(project_id, question=question).model_dump()
+        )
         return fallback_answer, created_entry
 
     return fallback_answer, None
diff --git a/src/cleanlab_codex/utils/errors.py b/src/cleanlab_codex/utils/errors.py
@@ -4,12 +4,19 @@
 class MissingDependencyError(Exception):
     """Raised when a lazy import is missing."""
 
-    def __init__(self, import_name: str, package_url: str | None = None) -> None:
+    def __init__(self, import_name: str, package_name: str | None = None, package_url: str | None = None) -> None:
+        """
+        Args:
+            import_name: The name of the import that failed.
+            package_name: The name of the package to install.
+            package_url: The URL for more information about the package.
+        """
         self.import_name = import_name
+        self.package_name = package_name
         self.package_url = package_url
 
     def __str__(self) -> str:
-        message = f"Failed to import {self.import_name}. Please install the package using `pip install {self.import_name}` and try again."
+        message = f"Failed to import {self.import_name}. Please install the package using `pip install {self.package_name or self.import_name}` and try again."
         if self.package_url:
             message += f" For more information, see {self.package_url}."
         return message
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -8,11 +8,11 @@
 from codex import AuthenticationError
 from codex.types.project_return_schema import Config as ProjectReturnConfig
 from codex.types.project_return_schema import ProjectReturnSchema
+from codex.types.users.myself.user_organizations_schema import Organization as SDKOrganization
 from codex.types.users.myself.user_organizations_schema import UserOrganizationsSchema
 
 from cleanlab_codex.client import Client
 from cleanlab_codex.project import MissingProjectError
-from cleanlab_codex.types.organization import Organization
 from cleanlab_codex.types.project import ProjectConfig
 
 FAKE_PROJECT_ID = str(uuid.uuid4())
@@ -29,7 +29,7 @@ def test_client_uses_default_organization(mock_client_from_api_key: MagicMock) -
     default_org_id = "default-org-id"
     mock_client_from_api_key.users.myself.organizations.list.return_value = UserOrganizationsSchema(
         organizations=[
-            Organization(
+            SDKOrganization(
                 organization_id=default_org_id,
                 created_at=datetime.now(),
                 updated_at=datetime.now(),
@@ -98,7 +98,7 @@ def test_get_project_not_found(mock_client_from_api_key: MagicMock) -> None:
 def test_list_organizations(mock_client_from_api_key: MagicMock) -> None:
     mock_client_from_api_key.users.myself.organizations.list.return_value = UserOrganizationsSchema(
         organizations=[
-            Organization(
+            SDKOrganization(
                 organization_id=FAKE_ORGANIZATION_ID,
                 created_at=datetime.now(),
                 updated_at=datetime.now(),

diff --git a/tests/test_codex_tool.py b/tests/test_codex_tool.py
@@ -1,11 +1,24 @@
+import builtins
+import importlib
 import sys
+from typing import Any
 from unittest.mock import MagicMock, patch
 
 import pytest
 from langchain_core.tools.structured import StructuredTool
 from llama_index.core.tools import FunctionTool
 
 from cleanlab_codex.codex_tool import CodexTool
+from cleanlab_codex.utils.errors import MissingDependencyError
+
+
+def patch_import_with_import_error(missing_module: str) -> None:
+    def custom_import(name: str, *args: Any, **kwargs: Any) -> Any:
+        if name.startswith(missing_module):
+            raise ImportError("test", name=missing_module)
+        return importlib.__import__(name, *args, **kwargs)
+
+    builtins.__import__ = custom_import
 
 
 def test_to_openai_tool(mock_client_from_access_key: MagicMock) -> None:
@@ -32,6 +45,20 @@ def test_to_llamaindex_tool(mock_client_from_access_key: MagicMock) -> None:
         assert llama_index_tool.fn == tool.query
 
 
+def test_to_llamaindex_tool_import_error(
+    mock_client_from_access_key: MagicMock,
+) -> None:
+    with patch("cleanlab_codex.codex_tool.Project") as mock_project:
+        mock_project.from_access_key.return_value = MagicMock(client=mock_client_from_access_key, id="test_project_id")
+
+        tool = CodexTool.from_access_key("sk-test-123")
+        patch_import_with_import_error("llama_index")
+        with pytest.raises(MissingDependencyError) as exc_info:
+            tool.to_llamaindex_tool()
+
+        assert exc_info.value.import_name == "llama_index"
+
+
 def test_to_langchain_tool(mock_client_from_access_key: MagicMock) -> None:
     with patch("cleanlab_codex.codex_tool.Project") as mock_project:
         mock_project.from_access_key.return_value = MagicMock(client=mock_client_from_access_key, id="test_project_id")
@@ -50,6 +77,18 @@ def test_to_langchain_tool(mock_client_from_access_key: MagicMock) -> None:
         ), f"Expected description '{tool.tool_description}', got '{langchain_tool.description}'."
 
 
+def test_to_langchain_tool_import_error(mock_client_from_access_key: MagicMock) -> None:
+    with patch("cleanlab_codex.codex_tool.Project") as mock_project:
+        mock_project.from_access_key.return_value = MagicMock(client=mock_client_from_access_key, id="test_project_id")
+
+        tool = CodexTool.from_access_key("sk-test-123")
+        patch_import_with_import_error("langchain")
+        with pytest.raises(MissingDependencyError) as exc_info:
+            tool.to_langchain_tool()
+
+        assert exc_info.value.import_name == "langchain"
+
+
 def test_to_aws_converse_tool(mock_client_from_access_key: MagicMock) -> None:
     with patch("cleanlab_codex.codex_tool.Project") as mock_project:
         mock_project.from_access_key.return_value = MagicMock(client=mock_client_from_access_key, id="test_project_id")
@@ -95,3 +134,21 @@ def test_to_smolagents_tool(mock_client_from_access_key: MagicMock) -> None:
         assert isinstance(smolagents_tool, Tool)
         assert smolagents_tool.name == tool.tool_name
         assert smolagents_tool.description == tool.tool_description
+
+
+def test_to_smolagents_tool_import_error(
+    mock_client_from_access_key: MagicMock,
+) -> None:
+    with patch("cleanlab_codex.codex_tool.Project") as mock_project:
+        mock_project.from_access_key.return_value = MagicMock(client=mock_client_from_access_key, id="test_project_id")
+
+        tool = CodexTool.from_access_key("sk-test-123")
+        import_module_name = "smolagents"
+        if sys.version_info >= (3, 10):
+            import_module_name = "cleanlab_codex.utils.smolagents"
+            patch_import_with_import_error(import_module_name)
+
+        with pytest.raises(MissingDependencyError) as exc_info:
+            tool.to_smolagents_tool()
+
+        assert exc_info.value.import_name == import_module_name
diff --git a/tests/test_project.py b/tests/test_project.py
@@ -6,9 +6,10 @@
 from codex import AuthenticationError
 from codex.types.project_create_params import Config
 from codex.types.projects.access_key_retrieve_project_id_response import AccessKeyRetrieveProjectIDResponse
+from codex.types.projects.entry import Entry as SDKEntry
 
 from cleanlab_codex.project import MissingProjectError, Project
-from cleanlab_codex.types.entry import Entry, EntryCreate
+from cleanlab_codex.types.entry import EntryCreate
 
 FAKE_PROJECT_ID = str(uuid.uuid4())
 FAKE_USER_ID = "Test User"
@@ -138,11 +139,12 @@ def test_query_read_only(mock_client_from_access_key: MagicMock) -> None:
         FAKE_PROJECT_ID, question="What is the capital of France?"
     )
     mock_client_from_access_key.projects.entries.add_question.assert_not_called()
-    assert res == (None, None)
+    assert res[0] is None
+    assert res[1] is None
 
 
 def test_query_question_found_fallback_answer(mock_client_from_access_key: MagicMock) -> None:
-    unanswered_entry = Entry(
+    unanswered_entry = SDKEntry(
         id=str(uuid.uuid4()),
         created_at=datetime.now(tz=timezone.utc),
         question="What is the capital of France?",
@@ -151,22 +153,32 @@ def test_query_question_found_fallback_answer(mock_client_from_access_key: Magic
     mock_client_from_access_key.projects.entries.query.return_value = unanswered_entry
     project = Project(mock_client_from_access_key, FAKE_PROJECT_ID)
     res = project.query("What is the capital of France?")
-    assert res == (None, unanswered_entry)
+    assert res[0] is None
+    assert res[1] is not None
+    assert res[1].model_dump() == unanswered_entry.model_dump()
 
 
 def test_query_question_not_found_fallback_answer(mock_client_from_access_key: MagicMock) -> None:
     mock_client_from_access_key.projects.entries.query.return_value = None
-    mock_client_from_access_key.projects.entries.add_question.return_value = MagicMock(spec=Entry)
+    mock_entry = SDKEntry(
+        id="fake-id",
+        created_at=datetime.now(tz=timezone.utc),
+        question="What is the capital of France?",
+        answer=None,
+    )
+    mock_client_from_access_key.projects.entries.add_question.return_value = mock_entry
 
     project = Project(mock_client_from_access_key, FAKE_PROJECT_ID)
     res = project.query("What is the capital of France?", fallback_answer="Paris")
     assert res[0] == "Paris"
+    assert res[1] is not None
+    assert res[1].model_dump() == mock_entry.model_dump()
 
 
 def test_query_add_question_when_not_found(mock_client_from_access_key: MagicMock) -> None:
     """Test that query adds question when not found and not read_only"""
     mock_client_from_access_key.projects.entries.query.return_value = None
-    new_entry = Entry(
+    new_entry = SDKEntry(
         id=str(uuid.uuid4()),
         created_at=datetime.now(tz=timezone.utc),
         question="What is the capital of France?",
@@ -180,11 +192,13 @@ def test_query_add_question_when_not_found(mock_client_from_access_key: MagicMoc
     mock_client_from_access_key.projects.entries.add_question.assert_called_once_with(
         FAKE_PROJECT_ID, question="What is the capital of France?"
     )
-    assert res == (None, new_entry)
+    assert res[0] is None
+    assert res[1] is not None
+    assert res[1].model_dump() == new_entry.model_dump()
 
 
 def test_query_answer_found(mock_client_from_access_key: MagicMock) -> None:
-    answered_entry = Entry(
+    answered_entry = SDKEntry(
         id=str(uuid.uuid4()),
         created_at=datetime.now(tz=timezone.utc),
         question="What is the capital of France?",
@@ -193,7 +207,9 @@ def test_query_answer_found(mock_client_from_access_key: MagicMock) -> None:
     mock_client_from_access_key.projects.entries.query.return_value = answered_entry
     project = Project(mock_client_from_access_key, FAKE_PROJECT_ID)
     res = project.query("What is the capital of France?")
-    assert res == ("Paris", answered_entry)
+    assert res[0] == answered_entry.answer
+    assert res[1] is not None
+    assert res[1].model_dump() == answered_entry.model_dump()
 
 
 def test_add_entries_empty_list(mock_client_from_access_key: MagicMock) -> None: