feat: tests and bug fixes AgentState.state (#1058)

Author: cpacker

memgpt/agent.py

@@ -9,14 +9,15 @@
 
 from memgpt.metadata import MetadataStore
 from memgpt.agent_store.storage import StorageConnector, TableType
-from memgpt.data_types import AgentState, Message, EmbeddingConfig, Passage
+from memgpt.data_types import AgentState, Message, LLMConfig, EmbeddingConfig, Passage, Preset
 from memgpt.models import chat_completion_response
 from memgpt.interface import AgentInterface
 from memgpt.persistence_manager import LocalStateManager
 from memgpt.system import get_login_event, package_function_response, package_summarize_message, get_initial_boot_messages
 from memgpt.memory import CoreMemory as InContextMemory, summarize_messages
 from memgpt.llm_api_tools import create, is_context_overflow_error
 from memgpt.utils import (
+    create_random_username,
     get_tool_call_id,
     get_local_time,
     parse_json,
@@ -167,37 +168,81 @@ def initialize_message_sequence(
 class Agent(object):
     def __init__(
         self,
-        agent_state: AgentState,
         interface: AgentInterface,
+        # agents can be created from providing agent_state
+        agent_state: Optional[AgentState] = None,
+        # or from providing a preset (requires preset + extra fields)
+        preset: Optional[Preset] = None,
+        created_by: Optional[uuid.UUID] = None,
+        name: Optional[str] = None,
+        llm_config: Optional[LLMConfig] = None,
+        embedding_config: Optional[EmbeddingConfig] = None,
         # extras
         messages_total: Optional[int] = None,  # TODO remove?
         first_message_verify_mono: bool = True,  # TODO move to config?
     ):
+
+        # An agent can be created from a Preset object
+        if preset is not None:
+            assert agent_state is None, "Can create an agent from a Preset or AgentState (but both were provided)"
+            assert created_by is not None, "Must provide created_by field when creating an Agent from a Preset"
+            assert llm_config is not None, "Must provide llm_config field when creating an Agent from a Preset"
+            assert embedding_config is not None, "Must provide embedding_config field when creating an Agent from a Preset"
+
+            # if agent_state is also provided, override any preset values
+            init_agent_state = AgentState(
+                name=name if name else create_random_username(),
+                user_id=created_by,
+                persona=preset.persona,
+                human=preset.human,
+                llm_config=llm_config,
+                embedding_config=embedding_config,
+                preset=preset.name,  # TODO link via preset.id instead of name?
+                state={
+                    "persona": preset.persona,
+                    "human": preset.human,
+                    "system": preset.system,
+                    "functions": preset.functions_schema,
+                    "messages": None,
+                },
+            )
+
+        # An agent can also be created directly from AgentState
+        elif agent_state is not None:
+            assert preset is None, "Can create an agent from a Preset or AgentState (but both were provided)"
+            assert agent_state.state is not None and agent_state.state != {}, "AgentState.state cannot be empty"
+
+            # Assume the agent_state passed in is formatted correctly
+            init_agent_state = agent_state
+
+        else:
+            raise ValueError("Both Preset and AgentState were null (must provide one or the other)")
+
         # Hold a copy of the state that was used to init the agent
-        self.agent_state = agent_state
+        self.agent_state = init_agent_state
 
         # gpt-4, gpt-3.5-turbo, ...
-        self.model = agent_state.llm_config.model
+        self.model = self.agent_state.llm_config.model
 
         # Store the system instructions (used to rebuild memory)
-        if "system" not in agent_state.state:
+        if "system" not in self.agent_state.state:
             raise ValueError(f"'system' not found in provided AgentState")
-        self.system = agent_state.state["system"]
+        self.system = self.agent_state.state["system"]
 
-        if "functions" not in agent_state.state:
+        if "functions" not in self.agent_state.state:
             raise ValueError(f"'functions' not found in provided AgentState")
         # Store the functions schemas (this is passed as an argument to ChatCompletion)
-        self.functions = agent_state.state["functions"]  # these are the schema
+        self.functions = self.agent_state.state["functions"]  # these are the schema
         # Link the actual python functions corresponding to the schemas
         self.functions_python = {k: v["python_function"] for k, v in link_functions(function_schemas=self.functions).items()}
         assert all([callable(f) for k, f in self.functions_python.items()]), self.functions_python
 
         # Initialize the memory object
-        if "persona" not in agent_state.state:
+        if "persona" not in self.agent_state.state:
             raise ValueError(f"'persona' not found in provided AgentState")
-        if "human" not in agent_state.state:
+        if "human" not in self.agent_state.state:
             raise ValueError(f"'human' not found in provided AgentState")
-        self.memory = initialize_memory(ai_notes=agent_state.state["persona"], human_notes=agent_state.state["human"])
+        self.memory = initialize_memory(ai_notes=self.agent_state.state["persona"], human_notes=self.agent_state.state["human"])
 
         # Interface must implement:
         # - internal_monologue
@@ -210,7 +255,7 @@ def __init__(
 
         # Create the persistence manager object based on the AgentState info
         # TODO
-        self.persistence_manager = LocalStateManager(agent_state=agent_state)
+        self.persistence_manager = LocalStateManager(agent_state=self.agent_state)
 
         # State needed for heartbeat pausing
         self.pause_heartbeats_start = None
@@ -226,17 +271,17 @@ def __init__(
         self._messages: List[Message] = []
 
         # Once the memory object is initialized, use it to "bake" the system message
-        if "messages" in agent_state.state and agent_state.state["messages"] is not None:
+        if "messages" in self.agent_state.state and self.agent_state.state["messages"] is not None:
             # print(f"Agent.__init__ :: loading, state={agent_state.state['messages']}")
-            if not isinstance(agent_state.state["messages"], list):
-                raise ValueError(f"'messages' in AgentState was bad type: {type(agent_state.state['messages'])}")
-            assert all([isinstance(msg, str) for msg in agent_state.state["messages"]])
+            if not isinstance(self.agent_state.state["messages"], list):
+                raise ValueError(f"'messages' in AgentState was bad type: {type(self.agent_state.state['messages'])}")
+            assert all([isinstance(msg, str) for msg in self.agent_state.state["messages"]])
 
             # Convert to IDs, and pull from the database
             raw_messages = [
-                self.persistence_manager.recall_memory.storage.get(id=uuid.UUID(msg_id)) for msg_id in agent_state.state["messages"]
+                self.persistence_manager.recall_memory.storage.get(id=uuid.UUID(msg_id)) for msg_id in self.agent_state.state["messages"]
             ]
-            assert all([isinstance(msg, Message) for msg in raw_messages]), (raw_messages, agent_state.state["messages"])
+            assert all([isinstance(msg, Message) for msg in raw_messages]), (raw_messages, self.agent_state.state["messages"])
             self._messages.extend([cast(Message, msg) for msg in raw_messages if msg is not None])
 
         else:

memgpt/cli/cli.py

@@ -412,7 +412,7 @@ def server(
 
 def run(
     persona: Annotated[Optional[str], typer.Option(help="Specify persona")] = None,
-    agent: Annotated[Optional[str], typer.Option(help="Specify agent save file")] = None,
+    agent: Annotated[Optional[str], typer.Option(help="Specify agent name")] = None,
     human: Annotated[Optional[str], typer.Option(help="Specify human")] = None,
     preset: Annotated[Optional[str], typer.Option(help="Specify preset")] = None,
     # model flags
@@ -605,18 +605,13 @@ def run(
         ms.update_agent(agent_state)
 
         # create agent
-        memgpt_agent = Agent(agent_state, interface=interface())
+        memgpt_agent = Agent(agent_state=agent_state, interface=interface())
 
     else:  # create new agent
         # create new agent config: override defaults with args if provided
         typer.secho("\n🧬 Creating new agent...", fg=typer.colors.WHITE)
 
-        if agent is None:
-            # determine agent name
-            # agent_count = len(ms.list_agents(user_id=user.id))
-            # agent = f"agent_{agent_count}"
-            agent = utils.create_random_username()
-
+        agent_name = agent if agent else utils.create_random_username()
         llm_config = config.default_llm_config
         embedding_config = config.default_embedding_config  # TODO allow overriding embedding params via CLI run
 
@@ -649,68 +644,43 @@ def run(
             )
             llm_config.model_endpoint_type = model_endpoint_type
 
-        agent_state = AgentState(
-            name=agent,
-            user_id=user.id,
-            persona=persona if persona else config.persona,
-            human=human if human else config.human,
-            preset=preset if preset else config.preset,
-            llm_config=llm_config,
-            embedding_config=embedding_config,
-        )
-        # ms.create_agent(agent_state)
-
-        typer.secho(f"->  🤖 Using persona profile '{agent_state.persona}'", fg=typer.colors.WHITE)
-        typer.secho(f"->  🧑 Using human profile '{agent_state.human}'", fg=typer.colors.WHITE)
-
-        # Supress llama-index noise
-        # TODO(swooders) add persistence manager code? or comment out?
-        # with suppress_stdout():
-        # TODO: allow configrable state manager (only local is supported right now)
-        # persistence_manager = LocalStateManager(agent_config)  # TODO: insert dataset/pre-fill
-
         # create agent
         try:
-            preset = ms.get_preset(preset_name=agent_state.preset, user_id=user.id)
-            if preset is None:
+            preset_obj = ms.get_preset(preset_name=preset if preset else config.preset, user_id=user.id)
+            if preset_obj is None:
                 # create preset records in metadata store
                 from memgpt.presets.presets import add_default_presets
 
                 add_default_presets(user.id, ms)
                 # try again
-                preset = ms.get_preset(preset_name=agent_state.preset, user_id=user.id)
-                assert preset is not None, "Couldn't find presets in database, please run `memgpt configure`"
+                preset_obj = ms.get_preset(preset_name=preset if preset else config.preset, user_id=user.id)
+                if preset_obj is None:
+                    typer.secho("Couldn't find presets in database, please run `memgpt configure`", fg=typer.colors.RED)
+                    sys.exit(1)
 
-            memgpt_agent = presets.create_agent_from_preset(
-                agent_state=agent_state,
-                preset=preset,
+            # Overwrite fields in the preset if they were specified
+            preset_obj.human = human if human else config.human
+            preset_obj.persona = persona if persona else config.persona
+
+            typer.secho(f"->  🤖 Using persona profile '{preset_obj.persona}'", fg=typer.colors.WHITE)
+            typer.secho(f"->  🧑 Using human profile '{preset_obj.human}'", fg=typer.colors.WHITE)
+
+            memgpt_agent = Agent(
                 interface=interface(),
+                name=agent_name,
+                created_by=user.id,
+                preset=preset_obj,
+                llm_config=llm_config,
+                embedding_config=embedding_config,
+                # gpt-3.5-turbo tends to omit inner monologue, relax this requirement for now
+                first_message_verify_mono=True if (model is not None and "gpt-4" in model) else False,
             )
             save_agent(agent=memgpt_agent, ms=ms)
+
         except ValueError as e:
-            # TODO(swooders) what's the equivalent cleanup code for the new DB refactor?
             typer.secho(f"Failed to create agent from provided information:\n{e}", fg=typer.colors.RED)
-            # # Delete the directory of the failed agent
-            # try:
-            #     # Path to the specific file
-            #     agent_config_file = agent_config.agent_config_path
-
-            #     # Check if the file exists
-            #     if os.path.isfile(agent_config_file):
-            #         # Delete the file
-            #         os.remove(agent_config_file)
-
-            #     # Now, delete the directory along with any remaining files in it
-            #     agent_save_dir = os.path.join(MEMGPT_DIR, "agents", agent_config.name)
-            #     shutil.rmtree(agent_save_dir)
-            # except:
-            #     typer.secho(f"Failed to delete agent directory during cleanup:\n{e}", fg=typer.colors.RED)
             sys.exit(1)
-        typer.secho(f"🎉 Created new agent '{agent_state.name}' (id={agent_state.id})", fg=typer.colors.GREEN)
-
-    # pretty print agent config
-    # printd(json.dumps(vars(agent_config), indent=4, sort_keys=True, ensure_ascii=JSON_ENSURE_ASCII))
-    # printd(json.dumps(agent_init_state), indent=4, sort_keys=True, ensure_ascii=JSON_ENSURE_ASCII))
+        typer.secho(f"🎉 Created new agent '{memgpt_agent.agent_state.name}' (id={memgpt_agent.agent_state.id})", fg=typer.colors.GREEN)
 
     # start event loop
     from memgpt.main import run_agent_loop

memgpt/presets/presets.py

@@ -68,6 +68,7 @@ def create_agent_from_preset(
     agent_state: AgentState, preset: Preset, interface: AgentInterface, persona_is_file: bool = True, human_is_file: bool = True
 ):
     """Initialize a new agent from a preset (combination of system + function)"""
+    raise DeprecationWarning("Function no longer supported - pass a Preset object to Agent.__init__ instead")
 
     # Input validation
     if agent_state.persona is None:

memgpt/server/server.py

@@ -598,34 +598,43 @@ def create_agent(
         if not user:
             raise ValueError(f"cannot find user with associated client id: {user_id}")
 
-        agent_state = AgentState(
-            user_id=user.id,
-            name=name if name else utils.create_random_username(),
-            preset=preset if preset else self.config.preset,
-            # TODO we need to allow passing raw persona/human text via the server request
-            persona=persona if persona else self.config.persona,
-            human=human if human else self.config.human,
-            llm_config=llm_config if llm_config else self.server_llm_config,
-            embedding_config=embedding_config if embedding_config else self.server_embedding_config,
-        )
         # NOTE: you MUST add to the metadata store before creating the agent, otherwise the storage connectors will error on creation
         # TODO: fix this db dependency and remove
         # self.ms.create_agent(agent_state)
 
-        logger.debug(f"Attempting to create agent from agent_state:\n{agent_state}")
         try:
-            preset = self.ms.get_preset(preset_name=agent_state.preset, user_id=user_id)
-            assert preset is not None, f"preset {agent_state.preset} does not exist"
-
-            agent = presets.create_agent_from_preset(agent_state=agent_state, preset=preset, interface=interface)
+            preset_obj = self.ms.get_preset(preset_name=preset if preset else self.config.preset, user_id=user_id)
+            assert preset_obj is not None, f"preset {preset if preset else self.config.preset} does not exist"
+            logger.debug(f"Attempting to create agent from preset:\n{preset_obj}")
+
+            # Overwrite fields in the preset if they were specified
+            preset_obj.human = human if human else self.config.human
+            preset_obj.persona = persona if persona else self.config.persona
+
+            llm_config = llm_config if llm_config else self.server_llm_config
+            embedding_config = embedding_config if embedding_config else self.server_embedding_config
+
+            agent = Agent(
+                interface=interface,
+                preset=preset_obj,
+                name=name,
+                created_by=user.id,
+                llm_config=llm_config,
+                embedding_config=embedding_config,
+                # gpt-3.5-turbo tends to omit inner monologue, relax this requirement for now
+                first_message_verify_mono=True if (llm_config.model is not None and "gpt-4" in llm_config.model) else False,
+            )
             save_agent(agent=agent, ms=self.ms)
 
             # FIXME: this is a hacky way to get the system prompts injected into agent into the DB
             # self.ms.update_agent(agent.agent_state)
         except Exception as e:
             logger.exception(e)
-            self.ms.delete_agent(agent_id=agent_state.id)
-            raise
+            try:
+                self.ms.delete_agent(agent_id=agent.agent_state.id)
+            except Exception as delete_e:
+                logger.exception(f"Failed to delete_agent:\n{delete_e}")
+            raise e
 
         save_agent(agent, self.ms)
 

tests/test_metadata_store.py

@@ -2,9 +2,11 @@
 from memgpt.constants import DEFAULT_HUMAN, DEFAULT_PERSONA, DEFAULT_PRESET
 import pytest
 
+from memgpt.agent import Agent, save_agent
 from memgpt.metadata import MetadataStore
 from memgpt.config import MemGPTConfig
 from memgpt.data_types import User, AgentState, Source, LLMConfig, EmbeddingConfig
+from memgpt.utils import get_human_text, get_persona_text
 
 
 # @pytest.mark.parametrize("storage_connector", ["postgres", "sqlite"])
@@ -50,6 +52,39 @@ def test_storage(storage_connector):
     len(ms.list_sources(user_id=user_1.id)) == 1
     len(ms.list_sources(user_id=user_2.id)) == 0
 
+    # test agent_state saving
+    agent_state = ms.get_agent(agent_1.id).state
+    assert agent_state is None, agent_state  # when created via create_agent, it should be empty
+
+    from memgpt.presets.presets import add_default_presets
+
+    add_default_presets(user_1.id, ms)
+    preset_obj = ms.get_preset(preset_name=DEFAULT_PRESET, user_id=user_1.id)
+    from memgpt.interface import CLIInterface as interface  # for printing to terminal
+
+    # Overwrite fields in the preset if they were specified
+    preset_obj.human = get_human_text(DEFAULT_HUMAN)
+    preset_obj.persona = get_persona_text(DEFAULT_PERSONA)
+
+    # Create the agent
+    agent = Agent(
+        interface=interface(),
+        created_by=user_1.id,
+        name="agent_test_agent_state",
+        preset=preset_obj,
+        llm_config=config.default_llm_config,
+        embedding_config=config.default_embedding_config,
+        # gpt-3.5-turbo tends to omit inner monologue, relax this requirement for now
+        first_message_verify_mono=(
+            True if (config.default_llm_config.model is not None and "gpt-4" in config.default_llm_config.model) else False
+        ),
+    )
+    agent_with_agent_state = agent.agent_state
+    save_agent(agent=agent, ms=ms)
+
+    agent_state = ms.get_agent(agent_with_agent_state.id).state
+    assert agent_state is not None, agent_state  # when created via create_agent_from_preset, it should be non-empty
+
     # test: updating
 
     # test: update JSON-stored LLMConfig class