From ceb0bfe414a8d2ec89154ea44b40975fad640bf8 Mon Sep 17 00:00:00 2001 From: Charles Packer Date: Sun, 21 Jan 2024 13:26:05 -0800 Subject: [PATCH] feat: add agent rename and agent delete to server + REST (#882) --- memgpt/server/rest_api/agents/config.py | 79 ++++++++++++++++++++++++- memgpt/server/server.py | 78 ++++++++++++++++++++---- 2 files changed, 146 insertions(+), 11 deletions(-) diff --git a/memgpt/server/rest_api/agents/config.py b/memgpt/server/rest_api/agents/config.py index 02c9f98b..87ec1228 100644 --- a/memgpt/server/rest_api/agents/config.py +++ b/memgpt/server/rest_api/agents/config.py @@ -1,5 +1,8 @@ import uuid -from fastapi import APIRouter, Depends, Query +import re + +from fastapi import APIRouter, Body, Depends, Query, HTTPException, status +from fastapi.responses import JSONResponse from pydantic import BaseModel, Field from memgpt.server.rest_api.interface import QueuingInterface @@ -13,10 +16,33 @@ class AgentConfigRequest(BaseModel): agent_id: str = Field(..., description="Identifier of the agent whose config is requested.") +class AgentRenameRequest(BaseModel): + user_id: str = Field(..., description="Unique identifier of the user requesting the config.") + agent_id: str = Field(..., description="Identifier of the agent whose config is requested.") + agent_name: str = Field(..., description="New name for the agent.") + + class AgentConfigResponse(BaseModel): config: dict = Field(..., description="The agent configuration object.") +def validate_agent_name(name: str) -> str: + """Validate the requested new agent name (prevent bad inputs)""" + + # Length check + if not (1 <= len(name) <= 50): + raise HTTPException(status_code=400, detail="Name length must be between 1 and 50 characters.") + + # Regex for allowed characters (alphanumeric, spaces, hyphens, underscores) + if not re.match("^[A-Za-z0-9 _-]+$", name): + raise HTTPException(status_code=400, detail="Name contains invalid characters.") + + # Further checks can be added here... + # TODO + + return name + + def setup_agents_config_router(server: SyncServer, interface: QueuingInterface): @router.get("/agents/config", tags=["agents"], response_model=AgentConfigResponse) def get_agent_config( @@ -40,4 +66,55 @@ def setup_agents_config_router(server: SyncServer, interface: QueuingInterface): config = server.get_agent_config(user_id=user_id, agent_id=agent_id) return AgentConfigResponse(config=config) + @router.patch("/agents/rename", tags=["agents"], response_model=AgentConfigResponse) + def update_agent_name(request: AgentRenameRequest = Body(...)): + """ + Updates the name of a specific agent. + + This changes the name of the agent in the database but does NOT edit the agent's persona. + """ + # TODO remove once chatui adds user selection / pulls user from config + request.user_id = None if request.user_id == "null" else request.user_id + + user_id = uuid.UUID(request.user_id) if request.user_id else None + agent_id = uuid.UUID(request.agent_id) if request.agent_id else None + + valid_name = validate_agent_name(request.agent_name) + + interface.clear() + try: + config = server.rename_agent(user_id=user_id, agent_id=agent_id, new_agent_name=valid_name) + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=500, detail=f"{e}") + return AgentConfigResponse(config=config) + + @router.delete("/agents", tags=["agents"]) + def delete_agent( + user_id: str = Query(..., description="Unique identifier of the user requesting the config."), + agent_id: str = Query(..., description="Identifier of the agent whose config is requested."), + ): + """ + Retrieve the configuration for a specific agent. + + This endpoint fetches the configuration details for a given agent, identified by the user and agent IDs. + """ + request = AgentConfigRequest(user_id=user_id, agent_id=agent_id) + + # TODO remove once chatui adds user selection / pulls user from config + request.user_id = None if request.user_id == "null" else request.user_id + + user_id = uuid.UUID(request.user_id) if request.user_id else None + agent_id = uuid.UUID(request.agent_id) if request.agent_id else None + + interface.clear() + try: + server.delete_agent(user_id=user_id, agent_id=agent_id) + return JSONResponse(status_code=status.HTTP_200_OK, content={"message": f"Agent agent_id={agent_id} successfully deleted"}) + except HTTPException: + raise + except Exception as e: + raise HTTPException(status_code=500, detail=f"{e}") + return router diff --git a/memgpt/server/server.py b/memgpt/server/server.py index 99c724bd..ce70f688 100644 --- a/memgpt/server/server.py +++ b/memgpt/server/server.py @@ -645,6 +645,19 @@ class SyncServer(LockingServer): if agent is not None: self.ms.delete_agent(agent_id=agent_id) + def _agent_state_to_config(self, agent_state: AgentState) -> dict: + """Convert AgentState to a dict for a JSON response""" + assert agent_state is not None + + agent_config = { + "id": agent_state.id, + "name": agent_state.name, + "human": agent_state.human, + "persona": agent_state.persona, + "created_at": agent_state.created_at.isoformat(), + } + return agent_config + def list_agents(self, user_id: uuid.UUID) -> dict: """List all available agents to a user""" if self.ms.get_user(user_id=user_id) is None: @@ -654,16 +667,7 @@ class SyncServer(LockingServer): logger.info(f"Retrieved {len(agents_states)} agents for user {user_id}:\n{[vars(s) for s in agents_states]}") return { "num_agents": len(agents_states), - "agents": [ - { - "id": state.id, - "name": state.name, - "human": state.human, - "persona": state.persona, - "created_at": state.created_at.isoformat(), - } - for state in agents_states - ], + "agents": [self._agent_state_to_config(state) for state in agents_states], } def get_agent(self, user_id: uuid.UUID, agent_id: uuid.UUID): @@ -882,6 +886,60 @@ class SyncServer(LockingServer): "modified": modified, } + def rename_agent(self, user_id: uuid.UUID, agent_id: uuid.UUID, new_agent_name: str) -> dict: + """Update the name of the agent in the database""" + if self.ms.get_user(user_id=user_id) is None: + raise ValueError(f"User user_id={user_id} does not exist") + if self.ms.get_agent(agent_id=agent_id, user_id=user_id) is None: + raise ValueError(f"Agent agent_id={agent_id} does not exist") + + # Get the agent object (loaded in memory) + memgpt_agent = self._get_or_load_agent(user_id=user_id, agent_id=agent_id) + + current_name = memgpt_agent.agent_state.name + if current_name == new_agent_name: + raise ValueError(f"New name ({new_agent_name}) is the same as the current name") + + try: + memgpt_agent.agent_state.name = new_agent_name + self.ms.update_agent(agent=memgpt_agent.agent_state) + except Exception as e: + logger.exception(f"Failed to update agent name with:\n{str(e)}") + raise ValueError(f"Failed to update agent name in database") + + # return the new config (only the name should have been updated) + agent_config = self._agent_state_to_config(agent_state=memgpt_agent.agent_state) + return agent_config + + def delete_agent(self, user_id: uuid.UUID, agent_id: uuid.UUID): + """Delete an agent in the database""" + if self.ms.get_user(user_id=user_id) is None: + raise ValueError(f"User user_id={user_id} does not exist") + if self.ms.get_agent(agent_id=agent_id, user_id=user_id) is None: + raise ValueError(f"Agent agent_id={agent_id} does not exist") + + # Verify that the agent exists and is owned by the user + agent_state = self.ms.get_agent(agent_id=agent_id, user_id=user_id) + if not agent_state: + raise ValueError(f"Could not find agent_id={agent_id} under user_id={user_id}") + if agent_state.user_id != user_id: + raise ValueError(f"Could not authorize agent_id={agent_id} with user_id={user_id}") + + # First, if the agent is in the in-memory cache we should remove it + # List of {'user_id': user_id, 'agent_id': agent_id, 'agent': agent_obj} dicts + try: + self.active_agents = [d for d in self.active_agents if str(d["agent_id"]) != str(agent_id)] + except Exception as e: + logger.exception(f"Failed to delete agent {agent_id} from cache via ID with:\n{str(e)}") + raise ValueError(f"Failed to delete agent {agent_id} from cache") + + # Next, attempt to delete it from the actual database + try: + self.ms.delete_agent(agent_id=agent_id) + except Exception as e: + logger.exception(f"Failed to delete agent {agent_id} via ID with:\n{str(e)}") + raise ValueError(f"Failed to delete agent {agent_id} in database") + def authenticate_user(self) -> uuid.UUID: # TODO: Implement actual authentication to enable multi user setup return uuid.UUID(int=uuid.getnode())