distinguish between import and package name in error

cleanlab · Feb 7, 2025 · 5249f33 · 5249f33
1 parent 31709e3
commit 5249f33
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/src/cleanlab_codex/codex_tool.py b/src/cleanlab_codex/codex_tool.py
@@ -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/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