Make the llm library optional

CHANGELOG.md

@@ -1,3 +1,13 @@
+## Unreleased
+
+### Features
+
+- Make LLM support optional and installable via `litecli[ai]`.
+
+### Bug Fixes
+
+- Avoid completion refresh crashes when no database is connected.
+
 ## 1.18.0 
 
 ### Internal

litecli/packages/special/llm.py

@@ -13,8 +13,22 @@
 from typing import Any
 
 import click
-import llm
-from llm.cli import cli
+
+try:
+    import llm
+
+    LLM_IMPORTED = True
+except ImportError:
+    llm = None
+    LLM_IMPORTED = False
+
+try:
+    from llm.cli import cli
+
+    LLM_CLI_IMPORTED = True
+except ImportError:
+    cli = None
+    LLM_CLI_IMPORTED = False
 
 from . import export
 from .main import Verbosity, parse_special_command
@@ -23,10 +37,10 @@
 log = logging.getLogger(__name__)
 
 LLM_TEMPLATE_NAME = "litecli-llm-template"
-LLM_CLI_COMMANDS: list[str] = list(cli.commands.keys())
+LLM_CLI_COMMANDS: list[str] = list(cli.commands.keys()) if LLM_CLI_IMPORTED else []
 # Mapping of model_id to None used for completion tree leaves.
 # the file name is llm.py and module name is llm, hence ty is complaining that get_models is missing.
-MODELS: dict[str, None] = {x.model_id: None for x in llm.get_models()}  # type: ignore[attr-defined]
+MODELS: dict[str, None] = {x.model_id: None for x in llm.get_models()} if LLM_IMPORTED else {}  # type: ignore[attr-defined]
 
 
 def run_external_cmd(
@@ -110,7 +124,7 @@ def build_command_tree(cmd: click.Command) -> dict[str, Any] | None:
 
 
 # Generate the tree
-COMMAND_TREE: dict[str, Any] | None = build_command_tree(cli)
+COMMAND_TREE: dict[str, Any] | None = build_command_tree(cli) if LLM_CLI_IMPORTED else {}
 
 
 def get_completions(tokens: list[str], tree: dict[str, Any] | None = COMMAND_TREE) -> list[str]:
@@ -123,6 +137,8 @@ def get_completions(tokens: list[str], tree: dict[str, Any] | None = COMMAND_TRE
     Returns:
         list[str]: List of possible completions.
     """
+    if not LLM_CLI_IMPORTED:
+        return []
     for token in tokens:
         if token.startswith("-"):
             # Skip options (flags)
@@ -171,6 +187,18 @@ def __init__(self, results: Any | None = None) -> None:
 # https://llm.datasette.io/en/stable/plugins/directory.html
 """
 
+NEED_DEPENDENCIES = """
+To enable LLM features you need to install litecli with AI support:
+
+    pip install 'litecli[ai]'
+
+or install LLM libraries separately
+
+   pip install llm
+
+This is required to use the \\llm command.
+"""
+
 _SQL_CODE_FENCE = r"```sql\n(.*?)\n```"
 PROMPT = """
 You are a helpful assistant who is a SQLite expert. You are embedded in a SQLite
@@ -230,6 +258,10 @@ def handle_llm(text: str, cur: DBCursor) -> tuple[str, str | None, float]:
     is_verbose = mode is Verbosity.VERBOSE
     is_succinct = mode is Verbosity.SUCCINCT
 
+    if not LLM_IMPORTED:
+        output = [(None, None, None, NEED_DEPENDENCIES)]
+        raise FinishIteration(output)
+
     if not arg.strip():  # No question provided. Print usage and bail.
         output = [(None, None, None, USAGE)]
         raise FinishIteration(output)

litecli/packages/special/main.py

@@ -8,6 +8,13 @@
 
 log = logging.getLogger(__name__)
 
+try:
+    import llm  # noqa: F401
+
+    LLM_IMPORTED = True
+except ImportError:
+    LLM_IMPORTED = False
+
 NO_QUERY = 0
 PARSED_QUERY = 1
 RAW_QUERY = 2
@@ -176,13 +183,19 @@ def quit(*_args: Any) -> None:
     arg_type=NO_QUERY,
     case_sensitive=True,
 )
-@special_command(
-    "\\llm",
-    "\\ai",
-    "Use LLM to construct a SQL query.",
-    arg_type=NO_QUERY,
-    case_sensitive=False,
-    aliases=(".ai", ".llm"),
-)
 def stub() -> None:
     raise NotImplementedError
+
+
+if LLM_IMPORTED:
+
+    @special_command(
+        "\\llm",
+        "\\ai",
+        "Use LLM to construct a SQL query.",
+        arg_type=NO_QUERY,
+        case_sensitive=False,
+        aliases=(".ai", ".llm"),
+    )
+    def llm_stub() -> None:
+        raise NotImplementedError

litecli/sqlexecute.py

@@ -179,7 +179,8 @@ def get_result(self, cursor: Any) -> tuple[str | None, list | None, list | None,
 
     def tables(self) -> Generator[tuple[str], None, None]:
         """Yields table names"""
-        assert self.conn is not None
+        if not self.conn:
+            return
         with closing(self.conn.cursor()) as cur:
             _logger.debug("Tables Query. sql: %r", self.tables_query)
             cur.execute(self.tables_query)
@@ -188,7 +189,8 @@ def tables(self) -> Generator[tuple[str], None, None]:
 
     def table_columns(self) -> Generator[tuple[str, str], None, None]:
         """Yields column names"""
-        assert self.conn is not None
+        if not self.conn:
+            return
         with closing(self.conn.cursor()) as cur:
             _logger.debug("Columns Query. sql: %r", self.table_columns_query)
             cur.execute(self.table_columns_query)
@@ -206,7 +208,8 @@ def databases(self) -> Generator[str, None, None]:
 
     def functions(self) -> Iterable[tuple]:
         """Yields tuples of (schema_name, function_name)"""
-        assert self.conn is not None
+        if not self.conn:
+            return
         with closing(self.conn.cursor()) as cur:
             _logger.debug("Functions Query. sql: %r", self.functions_query)
             cur.execute(self.functions_query % self.dbname)

pyproject.toml

@@ -13,10 +13,7 @@ dependencies = [
     "configobj>=5.0.5",
     "prompt-toolkit>=3.0.3,<4.0.0",
     "pygments>=1.6",
-    "sqlparse>=0.4.4",
-    "setuptools", # Required by llm commands to install models
-    "pip",
-    "llm>=0.25.0"
+    "sqlparse>=0.4.4"
 ]
 
 [build-system]
@@ -33,7 +30,11 @@ build-backend = "setuptools.build_meta"
 litecli = "litecli.main:cli"
 
 [project.optional-dependencies]
-ai = ["llm"]
+ai = [
+    "llm>=0.25.0",
+    "setuptools", # Required by llm commands to install models
+    "pip",
+]
 sqlean = ["sqlean-py>=3.47.0", 
     "sqlean-stubs>=0.0.3"]
 
@@ -45,7 +46,9 @@ dev = [
     "pytest-cov>=4.1.0",
     "tox>=4.8.0",
     "pdbpp>=0.10.3",
-    "llm>=0.19.0",
+    "llm>=0.25.0",
+    "setuptools",
+    "pip",
     "ty>=0.0.4"
 ]
 

tests/test_llm_special.py

@@ -2,9 +2,16 @@
 
 import pytest
 
+import litecli.packages.special.llm as llm_module
 from litecli.packages.special.llm import USAGE, FinishIteration, handle_llm
 
 
+@pytest.fixture(autouse=True)
+def enable_llm(monkeypatch):
+    monkeypatch.setattr(llm_module, "LLM_IMPORTED", True)
+    monkeypatch.setattr(llm_module, "LLM_CLI_COMMANDS", ["models"])
+
+
 @patch("litecli.packages.special.llm.llm")
 def test_llm_command_without_args(mock_llm, executor):
     r"""