MemGPT/memgpt/client/client.py
Sarah Wooders 8c06cc4bf7
refactor!: Migrate users + agent information into storage connectors (#785)
Co-authored-by: cpacker <packercharles@gmail.com>
2024-01-08 15:59:49 -08:00

115 lines
4.6 KiB
Python

import os
from typing import Dict, List, Union
from memgpt.data_types import AgentState
from memgpt.cli.cli import QuickstartChoice
from memgpt.cli.cli import set_config_with_dict, quickstart as quickstart_func, str_to_quickstart_choice
from memgpt.config import MemGPTConfig, AgentConfig
from memgpt.persistence_manager import PersistenceManager
from memgpt.server.rest_api.interface import QueuingInterface
from memgpt.server.server import SyncServer
class Client(object):
def __init__(
self,
user_id: str = None,
auto_save: bool = False,
quickstart: Union[QuickstartChoice, str, None] = None,
config: Union[Dict, MemGPTConfig] = None, # not the same thing as AgentConfig
debug: bool = False,
):
"""
Initializes a new instance of Client class.
:param auto_save: indicates whether to automatically save after every message.
:param quickstart: allows running quickstart on client init.
:param config: optional config settings to apply after quickstart
:param debug: indicates whether to display debug messages.
"""
self.auto_save = auto_save
# make sure everything is set up properly
MemGPTConfig.create_config_dir()
# If this is the first ever start, do basic initialization
if not MemGPTConfig.exists() and config is None and quickstart is None:
# Default to openai
print("Detecting uninitialized MemGPT, defaulting to quickstart == openai")
quickstart = "openai"
if quickstart:
# api key passed in config has priority over env var
if isinstance(config, dict) and "openai_api_key" in config:
openai_key = config["openai_api_key"]
else:
openai_key = os.environ.get("OPENAI_API_KEY", None)
# throw an error if we can't resolve the key
if openai_key:
os.environ["OPENAI_API_KEY"] = openai_key
elif quickstart == QuickstartChoice.openai or quickstart == "openai":
raise ValueError("Please set OPENAI_API_KEY or pass 'openai_api_key' in config dict")
if isinstance(quickstart, str):
quickstart = str_to_quickstart_choice(quickstart)
quickstart_func(backend=quickstart, debug=debug)
if config is not None:
set_config_with_dict(config)
self.user_id = MemGPTConfig.load().anon_clientid if user_id is None else user_id
self.interface = QueuingInterface(debug=debug)
self.server = SyncServer(default_interface=self.interface)
def list_agents(self):
self.interface.clear()
return self.server.list_agents(user_id=self.user_id)
def agent_exists(self, agent_id: str) -> bool:
existing = self.list_agents()
return agent_id in existing["agent_names"]
def create_agent(
self,
agent_config: dict,
# persistence_manager: Union[PersistenceManager, None] = None,
throw_if_exists: bool = True,
) -> AgentState:
if isinstance(agent_config, dict):
agent_name = agent_config.get("name")
else:
raise TypeError(f"agent_config must be of type dict")
if not self.agent_exists(agent_id=agent_name):
self.interface.clear()
agent_state = self.server.create_agent(user_id=self.user_id, agent_config=agent_config)
return agent_state
else:
raise ValueError(f"Agent {agent_name} already exists")
def get_agent_config(self, agent_id: str) -> Dict:
self.interface.clear()
return self.server.get_agent_config(user_id=self.user_id, agent_id=agent_id)
def get_agent_memory(self, agent_id: str) -> Dict:
self.interface.clear()
return self.server.get_agent_memory(user_id=self.user_id, agent_id=agent_id)
def update_agent_core_memory(self, agent_id: str, new_memory_contents: Dict) -> Dict:
self.interface.clear()
return self.server.update_agent_core_memory(user_id=self.user_id, agent_id=agent_id, new_memory_contents=new_memory_contents)
def user_message(self, agent_id: str, message: str) -> List[Dict]:
self.interface.clear()
self.server.user_message(user_id=self.user_id, agent_id=agent_id, message=message)
if self.auto_save:
self.save()
return self.interface.to_list()
def run_command(self, agent_id: str, command: str) -> Union[str, None]:
self.interface.clear()
return self.server.run_command(user_id=self.user_id, agent_id=agent_id, command=command)
def save(self):
self.server.save_agents()