r2r/cli/commands/users.py

import json
from builtins import list as _list
from uuid import UUID

import asyncclick as click
from asyncclick import pass_context

from cli.utils.timer import timer
from r2r import R2RAsyncClient, R2RException


@click.group()
def users():
    """Users commands."""
    pass


@users.command()
@click.argument("email", required=True, type=str)
@click.argument("password", required=True, type=str)
@pass_context
async def create(ctx: click.Context, email: str, password: str):
    """Create a new user."""
    client: R2RAsyncClient = ctx.obj

    try:
        with timer():
            response = await client.users.create(
                email=email,
                password=password,
            )
        click.echo(json.dumps(response, indent=2))
    except R2RException as e:
        click.echo(str(e), err=True)
    except Exception as e:
        click.echo(str(f"An unexpected error occurred: {e}"), err=True)


@users.command()
@click.option("--ids", multiple=True, help="Document IDs to fetch", type=str)
@click.option(
    "--offset",
    default=0,
    help="The offset to start from. Defaults to 0.",
    type=int,
)
@click.option(
    "--limit",
    default=100,
    help="The maximum number of nodes to return. Defaults to 100.",
    type=int,
)
@pass_context
async def list(
    ctx: click.Context,
    ids: tuple[str, ...],
    offset: int,
    limit: int,
):
    """Get an overview of users."""
    uuids: _list[str | UUID] | None = (
        [UUID(id_) for id_ in ids] if ids else None
    )

    client: R2RAsyncClient = ctx.obj
    try:
        with timer():
            response = await client.users.list(
                ids=uuids,
                offset=offset,
                limit=limit,
            )

        for user in response["results"]:  # type: ignore
            click.echo(json.dumps(user, indent=2))
    except R2RException as e:
        click.echo(str(e), err=True)
    except Exception as e:
        click.echo(str(f"An unexpected error occurred: {e}"), err=True)


@users.command()
@click.argument("id", required=True, type=str)
@pass_context
async def retrieve(ctx: click.Context, id):
    """Retrieve a user by ID."""
    client: R2RAsyncClient = ctx.obj

    try:
        with timer():
            response = await client.users.retrieve(id=id)
        click.echo(json.dumps(response, indent=2))
    except R2RException as e:
        click.echo(str(e), err=True)
    except Exception as e:
        click.echo(str(f"An unexpected error occurred: {e}"), err=True)


@users.command()
@pass_context
async def me(ctx: click.Context):
    """Retrieve the current user."""
    client: R2RAsyncClient = ctx.obj

    try:
        with timer():
            response = await client.users.me()

        click.echo(json.dumps(response, indent=2))
    except R2RException as e:
        click.echo(str(e), err=True)
    except Exception as e:
        click.echo(str(f"An unexpected error occurred: {e}"), err=True)


@users.command()
@click.argument("id", required=True, type=str)
@click.option(
    "--offset",
    default=0,
    help="The offset to start from. Defaults to 0.",
)
@click.option(
    "--limit",
    default=100,
    help="The maximum number of nodes to return. Defaults to 100.",
)
@pass_context
async def list_collections(ctx: click.Context, id, offset, limit):
    """List collections for a specific user."""
    client: R2RAsyncClient = ctx.obj

    try:
        with timer():
            response = await client.users.list_collections(
                id=id,
                offset=offset,
                limit=limit,
            )

        for collection in response["results"]:  # type: ignore
            click.echo(json.dumps(collection, indent=2))
    except R2RException as e:
        click.echo(str(e), err=True)
    except Exception as e:
        click.echo(str(f"An unexpected error occurred: {e}"), err=True)


@users.command()
@click.argument("id", required=True, type=str)
@click.argument("collection_id", required=True, type=str)
@pass_context
async def add_to_collection(ctx: click.Context, id, collection_id):
    """Retrieve a user by ID."""
    client: R2RAsyncClient = ctx.obj

    try:
        with timer():
            response = await client.users.add_to_collection(
                id=id,
                collection_id=collection_id,
            )
        click.echo(json.dumps(response, indent=2))
    except R2RException as e:
        click.echo(str(e), err=True)
    except Exception as e:
        click.echo(str(f"An unexpected error occurred: {e}"), err=True)


@users.command()
@click.argument("id", required=True, type=str)
@click.argument("collection_id", required=True, type=str)
@pass_context
async def remove_from_collection(ctx: click.Context, id, collection_id):
    """Retrieve a user by ID."""
    client: R2RAsyncClient = ctx.obj

    try:
        with timer():
            response = await client.users.remove_from_collection(
                id=id,
                collection_id=collection_id,
            )
        click.echo(json.dumps(response, indent=2))
    except R2RException as e:
        click.echo(str(e), err=True)
    except Exception as e:
        click.echo(str(f"An unexpected error occurred: {e}"), err=True)