feat: add cli for preparing env

pyproject.toml

@@ -19,16 +19,21 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry.dependencies]
 python = "^3.8"
 gitpython = "*"
-chromadb = "*"
+chromadb = "0.4.14"
 pydantic-settings = "*"
 pydantic = "*"
 tqdm = "*"
 loguru = "^0.7.2"
 pandas = "^2.0.3"
+click = "^8.1.7"
 
 # actually srctag still requires `sentence_transformers` here
 # but pytorch is a large dep which I don't want to manage it here
 sentence-transformers = { version = "^2.2.2", optional = true }
 
 [tool.poetry.extras]
 embedding = ["sentence-transformers"]
+
+
+[tool.poetry.scripts]
+srctag = 'srctag.cli:cli'

srctag/cli.py

@@ -0,0 +1,33 @@
+import chromadb
+import click
+
+from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
+from loguru import logger
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command()
+def prepare():
+    """ usually used for pre-downloading sentence-transformer models """
+
+    logger.info("Start checking env. It may takes a few minutes for downloading models ...")
+    # try to embed
+    chromadb_cli = chromadb.Client()
+    collection = chromadb_cli.get_or_create_collection(
+        "testonly",
+        embedding_function=SentenceTransformerEmbeddingFunction(),
+    )
+    collection.add(
+        documents=["doc"],
+        ids=["id"],
+    )
+    assert collection.count() == 1
+    click.echo("ok.")
+
+
+if __name__ == '__main__':
+    cli()