fix issues with optional dependencies (#204)

* fix issues with optional dependencies

* format document

* Tree: Format and comment

common/actions.py

@@ -3,13 +3,12 @@ from loguru import logger
 
 from common.tabby_config import config, generate_config_file
 from endpoints.server import export_openapi
-from endpoints.utils import do_export_openapi
 
 
 def branch_to_actions() -> bool:
     """Checks if a optional action needs to be run."""
 
-    if config.actions.export_openapi or do_export_openapi:
+    if config.actions.export_openapi:
         openapi_json = export_openapi()
 
         with open(config.actions.openapi_export_path, "w") as f:

common/model.py

@@ -13,22 +13,20 @@ from typing import Optional
 from common.logger import get_loading_progress_bar
 from common.networking import handle_request_error
 from common.tabby_config import config
-from endpoints.utils import do_export_openapi
+from common.optional_dependencies import dependencies
 
-if not do_export_openapi:
+if dependencies.exl2:
     from backends.exllamav2.model import ExllamaV2Container
 
     # Global model container
     container: Optional[ExllamaV2Container] = None
     embeddings_container = None
 
-    # Type hint the infinity emb container if it exists
-    from backends.infinity.model import has_infinity_emb
 
-    if has_infinity_emb:
-        from backends.infinity.model import InfinityContainer
+if dependencies.extras:
+    from backends.infinity.model import InfinityContainer
 
-        embeddings_container: Optional[InfinityContainer] = None
+    embeddings_container: Optional[InfinityContainer] = None
 
 
 class ModelType(Enum):
@@ -121,7 +119,7 @@ async def load_embedding_model(model_path: pathlib.Path, **kwargs):
     global embeddings_container
 
     # Break out if infinity isn't installed
-    if not has_infinity_emb:
+    if not dependencies.extras:
         raise ImportError(
             "Skipping embeddings because infinity-emb is not installed.\n"
             "Please run the following command in your environment "

common/optional_dependencies.py

@@ -0,0 +1,52 @@
+"""Construct a model of all optional dependencies"""
+
+import importlib.util
+from pydantic import BaseModel, computed_field
+
+
+# Declare the exported parts of this module
+__all__ = ["dependencies"]
+
+
+class DependenciesModel(BaseModel):
+    """Model of which optional dependencies are installed."""
+
+    torch: bool
+    exllamav2: bool
+    flash_attn: bool
+    outlines: bool
+    infinity_emb: bool
+    sentence_transformers: bool
+
+    @computed_field
+    @property
+    def extras(self) -> bool:
+        return self.outlines and self.infinity_emb and self.sentence_transformers
+
+    @computed_field
+    @property
+    def exl2(self) -> bool:
+        return self.torch and self.exllamav2 and self.flash_attn
+
+
+def is_installed(package_name: str) -> bool:
+    """Utility function to check if a package is installed."""
+
+    spec = importlib.util.find_spec(package_name)
+    return spec is not None
+
+
+def get_installed_deps() -> DependenciesModel:
+    """Check if optional dependencies are installed by looping over the fields."""
+
+    fields = DependenciesModel.model_fields
+
+    installed_deps = {}
+
+    for field_name in fields.keys():
+        installed_deps[field_name] = is_installed(field_name)
+
+    return DependenciesModel(**installed_deps)
+
+
+dependencies = get_installed_deps()