add colors!

chromy/chroma_functions.py

@@ -54,10 +54,11 @@ def delete_data(collection_name: str, where: dict[str, str]) -> int:
     return int(result.get("deleted", 0))
 
 
-def count_collection(collection_name: str) -> int:
+def count_collection(collection_name: str) -> str:
     _, collection = _get_client_and_collection(collection_name)
+    count = collection.count()
 
-    return collection.count()
+    return f"The '{collection_name}' collection contains [bold green]{count}[/] records."
 
 
 def add_data(

chromy/cli.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import Annotated, Callable
 
 import typer
+from rich import print
 from chromadb.errors import InternalError, NotFoundError
 
 from chromy.handlers.import_data import handle_import
@@ -27,7 +28,7 @@ def _run(handler: ExitCodeHandler) -> None:
 
 
 def _fail(message: str) -> None:
-    typer.echo(message)
+    print("[bold red]Error[/]:", message)
     raise typer.Exit(1)
 
 

chromy/handlers/count_collection.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from rich import print
 from chromy.chroma_functions import count_collection
 
 

chromy/handlers/create_collection.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
+from rich import print
 from chromy.chroma_functions import create_collection
 
 
 def handle_create_collection(collection: str) -> int:
     collection_name = create_collection(collection)
-    print(f"Created collection '{collection_name}'.")
+    print(f"[bold green]Created[/]: collection '{collection_name}'.")
     return 0

chromy/handlers/delete_collection.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from rich import print
 from chromy.chroma_functions import delete_collection, delete_data
 
 
@@ -22,7 +23,7 @@ def _parse_where_clause(where_clause: str) -> dict[str, str]:
 
 def handle_delete_collection(collection: str) -> int:
     delete_collection(collection)
-    print(f"Deleted collection '{collection}'.")
+    print(f"[bold green]Deleted[/] collection '{collection}'.")
     return 0
 
 
@@ -31,7 +32,7 @@ def handle_delete_records(collection: str, where_clause: str) -> int:
     deleted = delete_data(collection, where)
     condition, value = next(iter(where.items()))
     print(
-        f"Deleted {deleted} record(s) from collection '{collection}' "
+        f"[bold green]Deleted[/] {deleted} record(s) from collection '{collection}' "
         f"where {condition}={value}."
     )
     return 0

chromy/handlers/import_data.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 import os
 from pathlib import Path
 
+from rich import print
 from chromy.utilities import ingest_file
 
 
@@ -25,5 +26,6 @@ def _get_absolute_path(file: str) -> str:
 
 def handle_import(collection: str, file: str) -> int:
     records_added = ingest_file(collection, _get_absolute_path(file))
-    print(f"Added {records_added} records to collection '{collection}'.")
+    print(
+        f"[bold green]Added[/] {records_added} records to collection '{collection}'.")
     return 0

pyproject.toml

@@ -13,6 +13,7 @@ dependencies = [
     "openai>=2.32.0",
     "pymupdf4llm>=1.27.2.2",
     "python-dotenv>=1.2.2",
+    "rich>=15.0.0",
     "semchunk>=4.0.0",
     "tiktoken>=0.12.0",
     "transformers>=5.5.4",

tests/test_cli.py

@@ -46,7 +46,7 @@ class CliTests(unittest.TestCase):
 
         create_collection.assert_called_once_with("notes")
         self.assertEqual(result.exit_code, 0)
-        self.assertEqual(result.stdout, "Created collection 'notes'.\n")
+        self.assertEqual(result.stdout, "Created: collection 'notes'.\n")
 
     def test_create_collection_with_same_name(self) -> None:
         with patch(
@@ -58,7 +58,8 @@ class CliTests(unittest.TestCase):
 
         create_collection.assert_called_once_with("notes")
         self.assertEqual(result.exit_code, 1)
-        self.assertEqual(result.stdout, "Collection 'notes' already exists.\n")
+        self.assertEqual(
+            result.stdout, "Error: Collection 'notes' already exists.\n")
 
     def test_delete_collection(self) -> None:
         with patch(
@@ -79,7 +80,8 @@ class CliTests(unittest.TestCase):
 
         delete_collection.assert_called_once_with("notes")
         self.assertEqual(result.exit_code, 1)
-        self.assertEqual(result.stdout, "Collection 'notes' does not exist.\n")
+        self.assertEqual(
+            result.stdout, "Error: Collection 'notes' does not exist.\n")
 
     def test_count(self) -> None:
         with patch(
@@ -147,7 +149,7 @@ class CliTests(unittest.TestCase):
         self.assertEqual(result.exit_code, 1)
         self.assertEqual(
             result.stdout,
-            "Invalid --where value. Expected <condition>=<value>.\n",
+            "Error: Invalid --where value. Expected <condition>=<value>.\n",
         )
 
     def test_delete_requires_where_option(self) -> None:

tests/test_handlers.py

@@ -61,7 +61,7 @@ class HandlerTests(unittest.TestCase):
 
         create_collection.assert_called_once_with("notes")
         self.assertEqual(exit_code, 0)
-        self.assertEqual(output, "Created collection 'notes'.\n")
+        self.assertEqual(output, "Created: collection 'notes'.\n")
 
     def test_delete_collection_uses_typed_input(self) -> None:
         with patch("chromy.handlers.delete_collection.delete_collection") as delete:

uv.lock

@@ -261,6 +261,7 @@ dependencies = [
     { name = "openai" },
     { name = "pymupdf4llm" },
     { name = "python-dotenv" },
+    { name = "rich" },
     { name = "semchunk" },
     { name = "tiktoken" },
     { name = "transformers" },
@@ -281,6 +282,7 @@ requires-dist = [
     { name = "openai", specifier = ">=2.32.0" },
     { name = "pymupdf4llm", specifier = ">=1.27.2.2" },
     { name = "python-dotenv", specifier = ">=1.2.2" },
+    { name = "rich", specifier = ">=15.0.0" },
     { name = "semchunk", specifier = ">=4.0.0" },
     { name = "tiktoken", specifier = ">=0.12.0" },
     { name = "transformers", specifier = ">=5.5.4" },