add list, create, delete commands

chroma_functions.py

@@ -0,0 +1,30 @@
+import chromadb
+from typing import List
+
+
+def list_collections() -> List[str]:
+    client = chromadb.PersistentClient()
+    collections = client.list_collections()
+
+    if not collections:
+        return []
+
+    return [getattr(collection, "name", str(collection)) for collection in collections]
+
+
+def create_collection(name: str) -> str:
+    client = chromadb.PersistentClient()
+    collection = client.create_collection(name=name)
+    return getattr(collection, "name", name)
+
+
+def delete_collection(name: str) -> None:
+    client = chromadb.PersistentClient()
+    client.delete_collection(name=name)
+
+
+def count_collection(name: str) -> int:
+    # TODO Implement this method.
+    # The function must use the count method and return how many records are
+    # in the collection. It must handle non-existent collections.
+    raise NotImplemented()

cli_parser.py

@@ -0,0 +1,30 @@
+import argparse
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Inspect local Chroma collections.")
+    subparsers = parser.add_subparsers(dest="command")
+
+    subparsers.add_parser(
+        "list-collections",
+        aliases=["lc"],
+        help="List all collections stored in the local Chroma database.",
+    )
+
+    create_parser = subparsers.add_parser(
+        "create-collection",
+        aliases=["cc"],
+        help="Create a collection in the local Chroma database.",
+    )
+    create_parser.add_argument("name", help="Name of the collection to create.")
+
+    delete_parser = subparsers.add_parser(
+        "delete-collection",
+        aliases=["dc"],
+        help="Delete a collection from the local Chroma database.",
+    )
+    delete_parser.add_argument("name", help="Name of the collection to delete.")
+
+    # TODO add "count <name>" command.
+
+    return parser

main.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from chromadb.errors import NotFoundError, InternalError
+from chroma_functions import create_collection, delete_collection, list_collections
+
+from cli_parser import build_parser
+
+
+def main() -> int:
+    args = build_parser().parse_args()
+
+    if args.command in {"list-collections", "lc"}:
+        collections = list_collections()
+        if not collections:
+            print("No collections found.")
+            return 0
+
+        for name in collections:
+            print(name)
+        return 0
+
+    if args.command in {"create-collection", "cc"}:
+        try:
+            create_collection(args.name)
+        except InternalError:
+            print(f"Collection '{args.name}' already exists.")
+            return 1
+
+        print(f"Created collection '{args.name}'.")
+        return 0
+
+    if args.command in {"delete-collection", "dc"}:
+        try:
+            delete_collection(args.name)
+        except NotFoundError:
+            print(f"Collection '{args.name}' does not exist.")
+            return 1
+
+        print(f"Deleted collection '{args.name}'.")
+        return 0
+
+    # TODO Implement the count command.
+
+    print("Nothing to do. Use -h to see available commands.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())