add delete collection, refactor

main.py

@@ -1,7 +1,12 @@
 from __future__ import annotations
 
 from chromadb.errors import NotFoundError, InternalError
-from chroma_functions import create_collection, delete_collection, list_collections
+from chroma_functions import (
+    count_collection,
+    create_collection,
+    delete_collection,
+    list_collections,
+)
 
 from cli_parser import build_parser
 
@@ -17,16 +22,18 @@ def main() -> int:
 
         for name in collections:
             print(name)
+
         return 0
 
     if args.command in {"create-collection", "cc"}:
         try:
-            create_collection(args.name)
+            collection = create_collection(args.name)
         except InternalError:
             print(f"Collection '{args.name}' already exists.")
             return 1
 
-        print(f"Created collection '{args.name}'.")
+        print(f"Created collection '{collection}'.")
+
         return 0
 
     if args.command in {"delete-collection", "dc"}:
@@ -37,11 +44,22 @@ def main() -> int:
             return 1
 
         print(f"Deleted collection '{args.name}'.")
+
         return 0
 
-    # TODO Implement the count command.
+    if args.command in {"count", "co"}:
+        try:
+            count = count_collection(args.name)
+        except NotFoundError:
+            print(f"Collection '{args.name}' does not exist.")
+            return 1
+
+        print(count)
+
+        return 0
 
     print("Nothing to do. Use -h to see available commands.")
+
     return 0