diff --git a/chroma_functions.py b/chroma_functions.py new file mode 100644 index 0000000..84df5c5 --- /dev/null +++ b/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() diff --git a/cli_parser.py b/cli_parser.py new file mode 100644 index 0000000..377469c --- /dev/null +++ b/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 " command. + + return parser diff --git a/main.py b/main.py index e69de29..7999369 100644 --- a/main.py +++ b/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())