topoteretes--cognee/cognee/cli/commands/agents_command.py

import argparse
import asyncio
import json
from uuid import UUID

from cognee.cli.reference import SupportsCliCommand
from cognee.cli import DEFAULT_DOCS_URL
import cognee.cli.echo as fmt
from cognee.cli.exceptions import CliCommandException


class AgentsCommand(SupportsCliCommand):
    command_string = "agents"
    help_string = "Manage agents (create, list, get, delete, register, unregister, connections)"
    docs_url = DEFAULT_DOCS_URL
    description = """
Manage Cognee agents and their connections.

Subcommands:
  create        Create a new agent and grant it access to datasets
  list          List all agents you own
  get           Show details for a single agent
  delete        Delete an agent by ID
  register      Register an agent connection (session)
  unregister    Unregister an agent connection (session)
  connections   List active agent connections and their memory sources
"""

    def configure_parser(self, parser: argparse.ArgumentParser) -> None:
        sub = parser.add_subparsers(dest="agents_action", title="actions")

        # create
        p_create = sub.add_parser("create", help="Create a new agent")
        p_create.add_argument("name", help="Agent name")
        p_create.add_argument(
            "--datasets",
            nargs="*",
            default=None,
            help="Dataset names or UUIDs to grant the agent read/write access to",
        )

        # list
        sub.add_parser("list", help="List all agents you own")

        # get
        p_get = sub.add_parser("get", help="Show details for a single agent")
        p_get.add_argument("agent_id", help="Agent UUID")

        # delete
        p_del = sub.add_parser("delete", help="Delete an agent by ID")
        p_del.add_argument("agent_id", help="Agent UUID")
        p_del.add_argument("-f", "--force", action="store_true", help="Skip confirmation")

        # register
        p_reg = sub.add_parser("register", help="Register an agent connection")
        p_reg.add_argument("agent_session_name", help="Agent session name")
        p_reg.add_argument("--type", default="api", help="Connection type (default: api)")
        p_reg.add_argument(
            "--memory-mode",
            dest="memory_mode",
            default="unknown",
            help="Memory mode (default: unknown)",
        )
        p_reg.add_argument(
            "--session-id",
            dest="session_id",
            default=None,
            help="Optional session id",
        )
        p_reg.add_argument(
            "--dataset-ids",
            dest="dataset_ids",
            nargs="*",
            default=None,
            help="Dataset UUIDs to associate with the connection",
        )
        p_reg.add_argument(
            "--dataset-names",
            dest="dataset_names",
            nargs="*",
            default=None,
            help="Dataset names to associate with the connection",
        )

        # unregister
        p_unreg = sub.add_parser("unregister", help="Unregister an agent connection")
        p_unreg.add_argument("agent_session_name", help="Agent session name")

        # connections
        p_conn = sub.add_parser("connections", help="List active agent connections")
        p_conn.add_argument(
            "--agent-id",
            dest="agent_id",
            default=None,
            help="Filter by agent UUID",
        )
        p_conn.add_argument(
            "--range",
            dest="range_key",
            default="30d",
            help="Time range key (default: 30d)",
        )
        p_conn.add_argument(
            "--status",
            dest="status_filter",
            default=None,
            help="Filter by connection status",
        )
        p_conn.add_argument("--limit", type=int, default=50, help="Max results (default: 50)")
        p_conn.add_argument("--offset", type=int, default=0, help="Result offset (default: 0)")

    def execute(self, args: argparse.Namespace) -> None:
        action = getattr(args, "agents_action", None)
        if not action:
            fmt.error("No action specified. Use --help to see available actions.")
            raise CliCommandException("No action specified", error_code=1)

        dispatch = {
            "create": self._create,
            "list": self._list,
            "get": self._get,
            "delete": self._delete,
            "register": self._register,
            "unregister": self._unregister,
            "connections": self._connections,
        }
        dispatch[action](args)

    def _create(self, args: argparse.Namespace) -> None:
        async def run():
            import cognee
            from cognee.cli.user_resolution import resolve_cli_user

            user = await resolve_cli_user(getattr(args, "user_id", None), strict=True)
            result = await cognee.agents.create(args.name, datasets=args.datasets, user=user)

            fmt.success(f"Created agent '{args.name}' ({result['agent_id']})")
            fmt.echo(f"Agent ID:    {result['agent_id']}")
            fmt.echo(f"Agent email: {result['agent_email']}")
            fmt.echo(f"API key:     {result['agent_api_key']}")
            fmt.warning(
                "Store this API key now. It is a one-time secret and cannot be retrieved again."
            )

        asyncio.run(run())

    def _list(self, args: argparse.Namespace) -> None:
        async def run():
            import cognee
            from cognee.cli.user_resolution import resolve_cli_user

            user = await resolve_cli_user(getattr(args, "user_id", None), strict=True)
            agents = await cognee.agents.list(user=user)
            if not agents:
                fmt.echo("No agents found.")
                return
            fmt.echo(f"{'Agent ID':<38} {'Agent Email':<40} {'API Key Label'}")
            fmt.echo("-" * 100)
            for a in agents:
                fmt.echo(f"{str(a['agent_id']):<38} {a['agent_email']:<40} {a['api_key_label']}")

        asyncio.run(run())

    def _get(self, args: argparse.Namespace) -> None:
        async def run():
            import cognee
            from cognee.cli.user_resolution import resolve_cli_user

            user = await resolve_cli_user(getattr(args, "user_id", None), strict=True)
            info = await cognee.agents.get(args.agent_id, user=user)
            fmt.echo(f"Agent ID:      {info['agent_id']}")
            fmt.echo(f"Agent email:   {info['agent_email']}")
            fmt.echo(f"API key label: {info['api_key_label']}")

        asyncio.run(run())

    def _delete(self, args: argparse.Namespace) -> None:
        agent_id = args.agent_id
        if not args.force:
            if not fmt.confirm(f"Delete agent {agent_id}? This cannot be undone"):
                fmt.echo("Cancelled.")
                return

        async def run():
            import cognee
            from cognee.cli.user_resolution import resolve_cli_user

            user = await resolve_cli_user(getattr(args, "user_id", None), strict=True)
            await cognee.agents.delete(agent_id, user=user)
            fmt.success(f"Agent {agent_id} deleted.")

        asyncio.run(run())

    def _register(self, args: argparse.Namespace) -> None:
        async def run():
            import cognee
            from cognee.cli.user_resolution import resolve_cli_user

            user = await resolve_cli_user(getattr(args, "user_id", None), strict=True)
            connection = await cognee.agents.register(
                args.agent_session_name,
                user=user,
                type=args.type,
                memory_mode=args.memory_mode,
                session_id=args.session_id,
                dataset_ids=args.dataset_ids,
                dataset_names=args.dataset_names,
            )
            connection_id = connection.get("id", args.agent_session_name)
            fmt.success(f"Registered agent connection {connection_id}")
            fmt.echo(json.dumps(connection, indent=2, default=str))

        asyncio.run(run())

    def _unregister(self, args: argparse.Namespace) -> None:
        async def run():
            import cognee
            from cognee.cli.user_resolution import resolve_cli_user

            user = await resolve_cli_user(getattr(args, "user_id", None), strict=True)
            count = await cognee.agents.unregister(args.agent_session_name, user=user)
            fmt.success(f"Active connections: {count}")

        asyncio.run(run())

    def _connections(self, args: argparse.Namespace) -> None:
        async def run():
            import cognee
            from cognee.cli.user_resolution import resolve_cli_user

            user = await resolve_cli_user(getattr(args, "user_id", None), strict=True)
            agent_id = UUID(args.agent_id) if args.agent_id else None
            response = await cognee.agents.list_connections(
                user=user,
                agent_id=agent_id,
                range_key=args.range_key,
                status_filter=args.status_filter,
                limit=args.limit,
                offset=args.offset,
            )

            agents = response.get("agents", []) or []
            if not agents:
                fmt.echo("No agent connections found.")
            else:
                fmt.echo(f"{'Agent ID':<38} {'Session':<30} {'Status':<12} {'Last Active'}")
                fmt.echo("-" * 100)
                for a in agents:
                    fmt.echo(
                        f"{str(a.get('agent_id', '')):<38} "
                        f"{str(a.get('agent_session_name', '')):<30} "
                        f"{str(a.get('status', '')):<12} "
                        f"{a.get('last_active_at', '')}"
                    )

            memory_sources = response.get("memory_sources")
            if memory_sources:
                fmt.echo("")
                fmt.echo("Memory sources:")
                fmt.echo(json.dumps(memory_sources, indent=2, default=str))

        asyncio.run(run())