simplify the app using typer
This commit is contained in:
Matteo Rosati
+159
@@ -0,0 +1,159 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Callable
|
||||
|
||||
import typer
|
||||
from chromadb.errors import InternalError, NotFoundError
|
||||
|
||||
from chromy.handlers.add_data import handle_add_data
|
||||
from chromy.handlers.count_collection import handle_count_collection
|
||||
from chromy.handlers.create_collection import handle_create_collection
|
||||
from chromy.handlers.delete_collection import (
|
||||
handle_delete_collection,
|
||||
handle_delete_records,
|
||||
)
|
||||
from chromy.handlers.list_collections import handle_list_collections
|
||||
from chromy.handlers.query import handle_query
|
||||
|
||||
app = typer.Typer(help="Inspect local Chroma collections.")
|
||||
|
||||
ExitCodeHandler = Callable[[], int]
|
||||
|
||||
|
||||
def _run(handler: ExitCodeHandler) -> None:
|
||||
exit_code = handler()
|
||||
if exit_code != 0:
|
||||
raise typer.Exit(exit_code)
|
||||
|
||||
|
||||
def _fail(message: str) -> None:
|
||||
typer.echo(message)
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
@app.command("lc", help="List all collections stored in the local Chroma database.")
|
||||
@app.command(
|
||||
"list-collections",
|
||||
help="List all collections stored in the local Chroma database.",
|
||||
)
|
||||
def list_collections() -> None:
|
||||
_run(handle_list_collections)
|
||||
|
||||
|
||||
@app.command("cc", help="Create a collection in the local Chroma database.")
|
||||
@app.command(
|
||||
"create-collection",
|
||||
help="Create a collection in the local Chroma database.",
|
||||
)
|
||||
def create_collection(
|
||||
collection: Annotated[
|
||||
str,
|
||||
typer.Argument(help="Name of the collection to create."),
|
||||
],
|
||||
) -> None:
|
||||
try:
|
||||
_run(lambda: handle_create_collection(collection))
|
||||
except InternalError:
|
||||
_fail(f"Collection '{collection}' already exists.")
|
||||
|
||||
|
||||
@app.command("dc", help="Delete a collection from the local Chroma database.")
|
||||
@app.command(
|
||||
"delete-collection",
|
||||
help="Delete a collection from the local Chroma database.",
|
||||
)
|
||||
def delete_collection(
|
||||
collection: Annotated[
|
||||
str,
|
||||
typer.Argument(help="Name of the collection to delete."),
|
||||
],
|
||||
) -> None:
|
||||
try:
|
||||
_run(lambda: handle_delete_collection(collection))
|
||||
except NotFoundError:
|
||||
_fail(f"Collection '{collection}' does not exist.")
|
||||
|
||||
|
||||
@app.command("co", help="Count records in a collection from the local Chroma database.")
|
||||
@app.command(
|
||||
"count",
|
||||
help="Count records in a collection from the local Chroma database.",
|
||||
)
|
||||
def count(
|
||||
collection: Annotated[
|
||||
str,
|
||||
typer.Argument(help="Name of the collection to count."),
|
||||
],
|
||||
) -> None:
|
||||
try:
|
||||
_run(lambda: handle_count_collection(collection))
|
||||
except NotFoundError:
|
||||
_fail(f"Collection '{collection}' does not exist.")
|
||||
|
||||
|
||||
@app.command(
|
||||
"ad",
|
||||
help="Chunk, embed, and add a file to a collection in the local Chroma database.",
|
||||
)
|
||||
@app.command(
|
||||
"add-data",
|
||||
help="Chunk, embed, and add a file to a collection in the local Chroma database.",
|
||||
)
|
||||
def add_data(
|
||||
collection: Annotated[
|
||||
str,
|
||||
typer.Argument(help="Name of the target collection."),
|
||||
],
|
||||
file: Annotated[
|
||||
str,
|
||||
typer.Argument(help="Path to the file to chunk and add to the collection."),
|
||||
],
|
||||
) -> None:
|
||||
try:
|
||||
_run(lambda: handle_add_data(collection, file))
|
||||
except NotFoundError:
|
||||
_fail(f"Collection '{collection}' does not exist.")
|
||||
except FileNotFoundError:
|
||||
_fail(f"The file {file} was not found.")
|
||||
|
||||
|
||||
@app.command("q", help="Query a collection with the provided text.")
|
||||
@app.command("query", help="Query a collection with the provided text.")
|
||||
def query(
|
||||
collection: Annotated[
|
||||
str,
|
||||
typer.Argument(help="Name of the target collection."),
|
||||
],
|
||||
query_text: Annotated[
|
||||
str,
|
||||
typer.Argument(help="The text to query."),
|
||||
],
|
||||
) -> None:
|
||||
try:
|
||||
_run(lambda: handle_query(collection, query_text))
|
||||
except NotFoundError:
|
||||
_fail(f"Collection '{collection}' does not exist.")
|
||||
|
||||
|
||||
@app.command("del", help="Delete records from a collection using a metadata filter.")
|
||||
@app.command("delete", help="Delete records from a collection using a metadata filter.")
|
||||
def delete_records(
|
||||
collection: Annotated[
|
||||
str,
|
||||
typer.Argument(help="Name of the target collection."),
|
||||
],
|
||||
where: Annotated[
|
||||
str,
|
||||
typer.Option(
|
||||
"--where",
|
||||
help="Metadata filter in the format <condition>=<value>.",
|
||||
metavar="CONDITION=VALUE",
|
||||
),
|
||||
],
|
||||
) -> None:
|
||||
try:
|
||||
_run(lambda: handle_delete_records(collection, where))
|
||||
except NotFoundError:
|
||||
_fail(f"Collection '{collection}' does not exist.")
|
||||
except ValueError as exc:
|
||||
_fail(str(exc))
|
||||
@@ -1,180 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from argparse import Namespace
|
||||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Generic, TypeVar, assert_never
|
||||
|
||||
from chromadb.errors import InternalError, NotFoundError
|
||||
|
||||
from chromy.command_inputs import (
|
||||
AddDataInput,
|
||||
CommandInput,
|
||||
CountCollectionInput,
|
||||
CreateCollectionInput,
|
||||
DeleteCollectionInput,
|
||||
DeleteRecordsInput,
|
||||
ListCollectionsInput,
|
||||
QueryInput,
|
||||
)
|
||||
from chromy.handlers.add_data import handle_add_data
|
||||
from chromy.handlers.count_collection import handle_count_collection
|
||||
from chromy.handlers.create_collection import handle_create_collection
|
||||
from chromy.handlers.delete_collection import (
|
||||
handle_delete_collection,
|
||||
handle_delete_records,
|
||||
)
|
||||
from chromy.handlers.list_collections import handle_list_collections
|
||||
from chromy.handlers.query import handle_query
|
||||
|
||||
CommandT = TypeVar("CommandT", bound=CommandInput)
|
||||
CollectionCommandT = TypeVar(
|
||||
"CollectionCommandT",
|
||||
DeleteCollectionInput,
|
||||
CountCollectionInput,
|
||||
AddDataInput,
|
||||
QueryInput,
|
||||
DeleteRecordsInput,
|
||||
)
|
||||
CommandHandler = Callable[[CommandT], int]
|
||||
ErrorMessageBuilder = Callable[[CommandT, Exception], str]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CliErrorHandler(Generic[CommandT]):
|
||||
exception_type: type[Exception]
|
||||
message: ErrorMessageBuilder[CommandT]
|
||||
|
||||
|
||||
def build_command_input(args: Namespace) -> CommandInput:
|
||||
command = str(args.command)
|
||||
|
||||
match command:
|
||||
case "list-collections":
|
||||
return ListCollectionsInput()
|
||||
case "create-collection":
|
||||
return CreateCollectionInput(collection=str(args.collection))
|
||||
case "delete-collection":
|
||||
return DeleteCollectionInput(collection=str(args.collection))
|
||||
case "count":
|
||||
return CountCollectionInput(collection=str(args.collection))
|
||||
case "add-data":
|
||||
return AddDataInput(collection=str(args.collection), file=str(args.file))
|
||||
case "query":
|
||||
return QueryInput(
|
||||
collection=str(args.collection),
|
||||
query_text=str(args.query_text),
|
||||
)
|
||||
case "delete":
|
||||
return DeleteRecordsInput(
|
||||
collection=str(args.collection),
|
||||
where=str(args.where),
|
||||
)
|
||||
case _:
|
||||
raise ValueError(f"Unknown command: {command}")
|
||||
|
||||
|
||||
def execute_command(args: Namespace) -> int:
|
||||
command_input = build_command_input(args)
|
||||
|
||||
match command_input:
|
||||
case ListCollectionsInput():
|
||||
return _run_command(command_input, handle_list_collections)
|
||||
case CreateCollectionInput():
|
||||
return _run_command(
|
||||
command_input,
|
||||
handle_create_collection,
|
||||
(
|
||||
CliErrorHandler(
|
||||
exception_type=InternalError,
|
||||
message=_collection_already_exists_message,
|
||||
),
|
||||
),
|
||||
)
|
||||
case DeleteCollectionInput():
|
||||
return _run_command(
|
||||
command_input,
|
||||
handle_delete_collection,
|
||||
(_collection_not_found_handler(DeleteCollectionInput),),
|
||||
)
|
||||
case CountCollectionInput():
|
||||
return _run_command(
|
||||
command_input,
|
||||
handle_count_collection,
|
||||
(_collection_not_found_handler(CountCollectionInput),),
|
||||
)
|
||||
case AddDataInput():
|
||||
return _run_command(
|
||||
command_input,
|
||||
handle_add_data,
|
||||
(
|
||||
_collection_not_found_handler(AddDataInput),
|
||||
CliErrorHandler(
|
||||
exception_type=FileNotFoundError,
|
||||
message=_file_not_found_message,
|
||||
),
|
||||
),
|
||||
)
|
||||
case QueryInput():
|
||||
return _run_command(
|
||||
command_input,
|
||||
handle_query,
|
||||
(_collection_not_found_handler(QueryInput),),
|
||||
)
|
||||
case DeleteRecordsInput():
|
||||
return _run_command(
|
||||
command_input,
|
||||
handle_delete_records,
|
||||
(
|
||||
_collection_not_found_handler(DeleteRecordsInput),
|
||||
CliErrorHandler(
|
||||
exception_type=ValueError,
|
||||
message=_exception_message,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
assert_never(command_input)
|
||||
|
||||
|
||||
def _run_command(
|
||||
command_input: CommandT,
|
||||
handler: CommandHandler[CommandT],
|
||||
error_handlers: Sequence[CliErrorHandler[CommandT]] = (),
|
||||
) -> int:
|
||||
try:
|
||||
return handler(command_input)
|
||||
except Exception as exc:
|
||||
for error_handler in error_handlers:
|
||||
if isinstance(exc, error_handler.exception_type):
|
||||
print(error_handler.message(command_input, exc))
|
||||
return 1
|
||||
raise
|
||||
|
||||
|
||||
def _collection_already_exists_message(
|
||||
command: CreateCollectionInput,
|
||||
_: Exception,
|
||||
) -> str:
|
||||
return f"Collection '{command.collection}' already exists."
|
||||
|
||||
|
||||
def _collection_not_found_handler(
|
||||
_: type[CollectionCommandT],
|
||||
) -> CliErrorHandler[CollectionCommandT]:
|
||||
return CliErrorHandler(
|
||||
exception_type=NotFoundError,
|
||||
message=_collection_not_found_message,
|
||||
)
|
||||
|
||||
|
||||
def _collection_not_found_message(command: CollectionCommandT, _: Exception) -> str:
|
||||
return f"Collection '{command.collection}' does not exist."
|
||||
|
||||
|
||||
def _file_not_found_message(command: AddDataInput, _: Exception) -> str:
|
||||
return f"The file {command.file} was not found."
|
||||
|
||||
|
||||
def _exception_message(_: DeleteRecordsInput, exc: Exception) -> str:
|
||||
return str(exc)
|
||||
@@ -1,123 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ArgumentSpec:
|
||||
name: str
|
||||
help: str
|
||||
required: bool = False
|
||||
metavar: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CommandSpec:
|
||||
name: str
|
||||
aliases: tuple[str, ...]
|
||||
help: str
|
||||
arguments: tuple[ArgumentSpec, ...] = ()
|
||||
|
||||
|
||||
COMMAND_SPECS: tuple[CommandSpec, ...] = (
|
||||
CommandSpec(
|
||||
name="list-collections",
|
||||
aliases=("lc",),
|
||||
help="List all collections stored in the local Chroma database.",
|
||||
),
|
||||
CommandSpec(
|
||||
name="create-collection",
|
||||
aliases=("cc",),
|
||||
help="Create a collection in the local Chroma database.",
|
||||
arguments=(ArgumentSpec("collection", "Name of the collection to create."),),
|
||||
),
|
||||
CommandSpec(
|
||||
name="delete-collection",
|
||||
aliases=("dc",),
|
||||
help="Delete a collection from the local Chroma database.",
|
||||
arguments=(ArgumentSpec("collection", "Name of the collection to delete."),),
|
||||
),
|
||||
CommandSpec(
|
||||
name="count",
|
||||
aliases=("co",),
|
||||
help="Count records in a collection from the local Chroma database.",
|
||||
arguments=(ArgumentSpec("collection", "Name of the collection to count."),),
|
||||
),
|
||||
CommandSpec(
|
||||
name="add-data",
|
||||
aliases=("ad",),
|
||||
help=(
|
||||
"Chunk, embed, and add a file to a collection in the local Chroma database."
|
||||
),
|
||||
arguments=(
|
||||
ArgumentSpec("collection", "Name of the target collection."),
|
||||
ArgumentSpec(
|
||||
"file", "Path to the file to chunk and add to the collection."
|
||||
),
|
||||
),
|
||||
),
|
||||
CommandSpec(
|
||||
name="query",
|
||||
aliases=("q",),
|
||||
help="Query a collection with the provided text.",
|
||||
arguments=(
|
||||
ArgumentSpec("collection", "Name of the target collection."),
|
||||
ArgumentSpec("query_text", "The text to query."),
|
||||
),
|
||||
),
|
||||
CommandSpec(
|
||||
name="delete",
|
||||
aliases=("del",),
|
||||
help="Delete records from a collection using a metadata filter.",
|
||||
arguments=(
|
||||
ArgumentSpec("collection", "Name of the target collection."),
|
||||
ArgumentSpec(
|
||||
"--where",
|
||||
"Metadata filter in the format <condition>=<value>.",
|
||||
required=True,
|
||||
metavar="CONDITION=VALUE",
|
||||
),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _add_command(
|
||||
subparsers: argparse._SubParsersAction[argparse.ArgumentParser],
|
||||
command: CommandSpec,
|
||||
) -> None:
|
||||
subparser = subparsers.add_parser(
|
||||
command.name,
|
||||
aliases=list(command.aliases),
|
||||
help=command.help,
|
||||
description=command.help,
|
||||
)
|
||||
|
||||
for argument in command.arguments:
|
||||
if argument.name.startswith("-"):
|
||||
subparser.add_argument(
|
||||
argument.name,
|
||||
help=argument.help,
|
||||
metavar=argument.metavar,
|
||||
required=argument.required,
|
||||
)
|
||||
continue
|
||||
|
||||
subparser.add_argument(
|
||||
argument.name,
|
||||
help=argument.help,
|
||||
metavar=argument.metavar,
|
||||
)
|
||||
|
||||
subparser.set_defaults(command=command.name)
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description="Inspect local Chroma collections.")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
for command in COMMAND_SPECS:
|
||||
_add_command(subparsers, command)
|
||||
|
||||
return parser
|
||||
@@ -1,52 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ListCollectionsInput:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CreateCollectionInput:
|
||||
collection: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DeleteCollectionInput:
|
||||
collection: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CountCollectionInput:
|
||||
collection: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AddDataInput:
|
||||
collection: str
|
||||
file: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class QueryInput:
|
||||
collection: str
|
||||
query_text: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DeleteRecordsInput:
|
||||
collection: str
|
||||
where: str
|
||||
|
||||
|
||||
CommandInput = (
|
||||
ListCollectionsInput
|
||||
| CreateCollectionInput
|
||||
| DeleteCollectionInput
|
||||
| CountCollectionInput
|
||||
| AddDataInput
|
||||
| QueryInput
|
||||
| DeleteRecordsInput
|
||||
)
|
||||
@@ -1,10 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from chromy.command_inputs import AddDataInput
|
||||
from chromy.utilities import ingest_file
|
||||
|
||||
|
||||
def handle_add_data(command: AddDataInput) -> int:
|
||||
records_added = ingest_file(command.collection, command.file)
|
||||
print(f"Added {records_added} records to collection '{command.collection}'.")
|
||||
def handle_add_data(collection: str, file: str) -> int:
|
||||
records_added = ingest_file(collection, file)
|
||||
print(f"Added {records_added} records to collection '{collection}'.")
|
||||
return 0
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from chromy.chroma_functions import count_collection
|
||||
from chromy.command_inputs import CountCollectionInput
|
||||
|
||||
|
||||
def handle_count_collection(command: CountCollectionInput) -> int:
|
||||
print(count_collection(command.collection))
|
||||
def handle_count_collection(collection: str) -> int:
|
||||
print(count_collection(collection))
|
||||
return 0
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from chromy.chroma_functions import create_collection
|
||||
from chromy.command_inputs import CreateCollectionInput
|
||||
|
||||
|
||||
def handle_create_collection(command: CreateCollectionInput) -> int:
|
||||
collection_name = create_collection(command.collection)
|
||||
def handle_create_collection(collection: str) -> int:
|
||||
collection_name = create_collection(collection)
|
||||
print(f"Created collection '{collection_name}'.")
|
||||
return 0
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from chromy.chroma_functions import delete_collection, delete_data
|
||||
from chromy.command_inputs import DeleteCollectionInput, DeleteRecordsInput
|
||||
|
||||
|
||||
def _parse_where_clause(where_clause: str) -> dict[str, str]:
|
||||
@@ -19,18 +18,18 @@ def _parse_where_clause(where_clause: str) -> dict[str, str]:
|
||||
return {condition: value}
|
||||
|
||||
|
||||
def handle_delete_collection(command: DeleteCollectionInput) -> int:
|
||||
delete_collection(command.collection)
|
||||
print(f"Deleted collection '{command.collection}'.")
|
||||
def handle_delete_collection(collection: str) -> int:
|
||||
delete_collection(collection)
|
||||
print(f"Deleted collection '{collection}'.")
|
||||
return 0
|
||||
|
||||
|
||||
def handle_delete_records(command: DeleteRecordsInput) -> int:
|
||||
where = _parse_where_clause(command.where)
|
||||
deleted = delete_data(command.collection, where)
|
||||
def handle_delete_records(collection: str, where_clause: str) -> int:
|
||||
where = _parse_where_clause(where_clause)
|
||||
deleted = delete_data(collection, where)
|
||||
condition, value = next(iter(where.items()))
|
||||
print(
|
||||
f"Deleted {deleted} record(s) from collection '{command.collection}' "
|
||||
f"Deleted {deleted} record(s) from collection '{collection}' "
|
||||
f"where {condition}={value}."
|
||||
)
|
||||
return 0
|
||||
|
||||
@@ -1,11 +1,10 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from chromy.chroma_functions import list_collections
|
||||
from chromy.command_inputs import ListCollectionsInput
|
||||
from chromy.utilities import print_lines
|
||||
|
||||
|
||||
def handle_list_collections(_: ListCollectionsInput) -> int:
|
||||
def handle_list_collections() -> int:
|
||||
collections = list_collections()
|
||||
if not collections:
|
||||
print("No collections found.")
|
||||
|
||||
@@ -1,10 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from chromy.command_inputs import QueryInput
|
||||
from chromy.utilities import format_query_result, print_lines, run_query
|
||||
|
||||
|
||||
def handle_query(command: QueryInput) -> int:
|
||||
result = run_query(command.collection, command.query_text)
|
||||
def handle_query(collection: str, query_text: str) -> int:
|
||||
result = run_query(collection, query_text)
|
||||
print_lines(format_query_result(result))
|
||||
return 0
|
||||
|
||||
+4
-6
@@ -2,15 +2,13 @@ from __future__ import annotations
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
from chromy.cli_app import execute_command
|
||||
from chromy.cli_parser import build_parser
|
||||
from chromy.cli import app
|
||||
|
||||
|
||||
def main() -> int:
|
||||
def main() -> None:
|
||||
load_dotenv()
|
||||
args = build_parser().parse_args()
|
||||
return execute_command(args)
|
||||
app()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
main()
|
||||
|
||||
Reference in New Issue
Block a user