Resolve default embeddings db in a function

llm/__init__.py

@@ -1,3 +1,4 @@
+import sqlite_utils
 from .hookspecs import hookimpl
 from .errors import (
     ModelError,
@@ -391,6 +392,9 @@ def user_dir():
     path.mkdir(exist_ok=True, parents=True)
     return path
 
+def get_default_embeddings_db() -> sqlite_utils.Database:
+    return sqlite_utils.Database(user_dir() / "embeddings.db")
+
 
 def set_alias(alias, model_id_or_alias):
     """

llm/cli.py

@@ -22,6 +22,7 @@
     KeyModel,
     encode,
     get_async_model,
+    get_default_embeddings_db,
     get_default_model,
     get_default_embedding_model,
     get_embedding_models_with_aliases,
@@ -2979,11 +2980,11 @@ def embed(
         raise click.ClickException("Must provide collection when using --store")
 
     # Lazy load this because we do not need it for -c or -i versions
-    def get_db():
+    def get_db() -> sqlite_utils.Database:
         if database:
             return sqlite_utils.Database(database)
         else:
-            return sqlite_utils.Database(user_dir() / "embeddings.db")
+            return get_default_embeddings_db()
 
     collection_obj = None
     model_obj = None
@@ -3168,7 +3169,7 @@ def embed_multi(
     if database:
         db = sqlite_utils.Database(database)
     else:
-        db = sqlite_utils.Database(user_dir() / "embeddings.db")
+        db = get_default_embeddings_db()
 
     for alias, attach_path in attach:
         db.attach(alias, attach_path)
@@ -3313,7 +3314,7 @@ def similar(collection, id, input, content, binary, number, plain, database, pre
     if database:
         db = sqlite_utils.Database(database)
     else:
-        db = sqlite_utils.Database(user_dir() / "embeddings.db")
+        db = get_default_embeddings_db()
 
     if not db["embeddings"].exists():
         raise click.ClickException("No embeddings table found in database")
@@ -3436,8 +3437,11 @@ def collections_path():
 @click.option("json_", "--json", is_flag=True, help="Output as JSON")
 def embed_db_collections(database, json_):
     "View a list of collections"
-    database = database or (user_dir() / "embeddings.db")
-    db = sqlite_utils.Database(str(database))
+    if database:
+        db = sqlite_utils.Database(str(database))
+    else:
+        db = get_default_embeddings_db()
+
     if not db["collections"].exists():
         raise click.ClickException("No collections table found in {}".format(database))
     rows = db.query(
@@ -3483,8 +3487,11 @@ def collections_delete(collection, database):
     \b
         llm collections delete my-collection
     """
-    database = database or (user_dir() / "embeddings.db")
-    db = sqlite_utils.Database(str(database))
+    if database:
+        db = sqlite_utils.Database(str(database))
+    else:
+        db = get_default_embeddings_db()
+
     try:
         collection_obj = Collection(collection, db, create=False)
     except Collection.DoesNotExist:

tests/test_embed_cli.py

@@ -1,6 +1,6 @@
 from click.testing import CliRunner
 from llm.cli import cli
-from llm import Collection
+from llm import Collection, get_default_embeddings_db
 import json
 import pathlib
 import pytest
@@ -153,7 +153,7 @@ def test_embed_store_binary(user_path):
     args = ["embed", "-m", "embed-demo", "items", "2", "--binary", "--store"]
     result = runner.invoke(cli, args, input=b"\x00\x01\x02")
     assert result.exit_code == 0
-    db = sqlite_utils.Database(str(user_path / "embeddings.db"))
+    db = get_default_embeddings_db()
     rows = list(db["embeddings"].rows)
     assert rows == [
         {
@@ -682,7 +682,7 @@ def test_default_embed_model_errors(user_path, default_is_set, command):
         result3 = runner.invoke(cli, args, input=input, catch_exceptions=False)
         assert result3.exit_code == 0
     # At the end of this, there should be 2 embeddings
-    db = sqlite_utils.Database(str(user_path / "embeddings.db"))
+    db = get_default_embeddings_db()
     assert db["embeddings"].count == 1