Saket Aryan
134 lines
4.0 KiB
Python
134 lines
4.0 KiB
Python
"""Entity management commands."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time as _time
|
|
|
|
import typer
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
|
|
from mem0_cli.backend.base import Backend
|
|
from mem0_cli.branding import (
|
|
ACCENT_COLOR,
|
|
BRAND_COLOR,
|
|
DIM_COLOR,
|
|
print_error,
|
|
print_info,
|
|
print_success,
|
|
timed_status,
|
|
)
|
|
from mem0_cli.output import format_json
|
|
|
|
console = Console()
|
|
err_console = Console(stderr=True)
|
|
|
|
|
|
def cmd_entities_list(backend: Backend, entity_type: str, *, output: str) -> None:
|
|
"""List entities of a given type."""
|
|
valid_types = {"users", "agents", "apps", "runs"}
|
|
if entity_type not in valid_types:
|
|
print_error(err_console, f"Invalid entity type: {entity_type}. Use: {', '.join(valid_types)}")
|
|
raise typer.Exit(1)
|
|
|
|
_start = _time.perf_counter()
|
|
with timed_status(err_console, f"Fetching {entity_type}...") as _ts:
|
|
try:
|
|
results = backend.entities(entity_type)
|
|
except Exception as e:
|
|
print_error(err_console, str(e), hint="This feature may require the mem0 Platform.")
|
|
raise typer.Exit(1) from None
|
|
_elapsed = _time.perf_counter() - _start
|
|
|
|
if output == "json":
|
|
format_json(console, results)
|
|
return
|
|
|
|
if not results:
|
|
print_info(console, f"No {entity_type} found.")
|
|
return
|
|
|
|
table = Table(border_style=BRAND_COLOR, header_style=f"bold {ACCENT_COLOR}", padding=(0, 1))
|
|
table.add_column("Name / ID", style="bold")
|
|
table.add_column("Created", max_width=12)
|
|
|
|
for entity in results:
|
|
name = entity.get("name", entity.get("id", "—"))
|
|
created = str(entity.get("created_at", "—"))[:10]
|
|
table.add_row(str(name), created)
|
|
|
|
console.print()
|
|
console.print(table)
|
|
console.print(f" [{DIM_COLOR}]{len(results)} {entity_type} ({_elapsed:.2f}s)[/]")
|
|
console.print()
|
|
|
|
|
|
def cmd_entities_delete(
|
|
backend: Backend,
|
|
*,
|
|
user_id: str | None,
|
|
agent_id: str | None,
|
|
app_id: str | None,
|
|
run_id: str | None,
|
|
force: bool,
|
|
dry_run: bool = False,
|
|
output: str,
|
|
) -> None:
|
|
"""Delete an entity and all its memories (cascade delete)."""
|
|
if not any([user_id, agent_id, app_id, run_id]):
|
|
print_error(err_console, "Provide at least one of --user-id, --agent-id, --app-id, --run-id.")
|
|
raise typer.Exit(1)
|
|
|
|
if dry_run:
|
|
scope_parts = []
|
|
if user_id:
|
|
scope_parts.append(f"user={user_id}")
|
|
if agent_id:
|
|
scope_parts.append(f"agent={agent_id}")
|
|
if app_id:
|
|
scope_parts.append(f"app={app_id}")
|
|
if run_id:
|
|
scope_parts.append(f"run={run_id}")
|
|
scope = ", ".join(scope_parts)
|
|
print_info(console, f"Would delete entity {scope} and all its memories.")
|
|
print_info(console, "No changes made (dry run).")
|
|
return
|
|
|
|
if not force:
|
|
scope_parts = []
|
|
if user_id:
|
|
scope_parts.append(f"user={user_id}")
|
|
if agent_id:
|
|
scope_parts.append(f"agent={agent_id}")
|
|
if app_id:
|
|
scope_parts.append(f"app={app_id}")
|
|
if run_id:
|
|
scope_parts.append(f"run={run_id}")
|
|
scope = ", ".join(scope_parts)
|
|
|
|
confirm = typer.confirm(
|
|
f"\n \u26a0 Delete entity {scope} AND all its memories? This cannot be undone."
|
|
)
|
|
if not confirm:
|
|
print_info(console, "Cancelled.")
|
|
raise typer.Exit(0)
|
|
|
|
_start = _time.perf_counter()
|
|
with timed_status(err_console, "Deleting entity...") as _ts:
|
|
try:
|
|
result = backend.delete_entities(
|
|
user_id=user_id,
|
|
agent_id=agent_id,
|
|
app_id=app_id,
|
|
run_id=run_id,
|
|
)
|
|
except Exception as e:
|
|
print_error(err_console, str(e))
|
|
raise typer.Exit(1) from None
|
|
_elapsed = _time.perf_counter() - _start
|
|
|
|
if output == "json":
|
|
format_json(console, result)
|
|
elif output != "quiet":
|
|
print_success(console, f"Entity deleted with all memories ({_elapsed:.2f}s)")
|