Skip to content

Commit

Permalink
fix: add DB prefill for default user, preset, humans, and persona for…
Browse files Browse the repository at this point in the history
… server (#1273)
  • Loading branch information
sarahwooders authored Apr 19, 2024
1 parent b17c653 commit 6767153
Show file tree
Hide file tree
Showing 7 changed files with 78 additions and 8 deletions.
5 changes: 3 additions & 2 deletions compose.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -26,12 +26,13 @@ services:
env_file:
- .env
environment:
- MEMGPT_PGURI=postgresql://${MEMGPT_PG_USER}:${MEMGPT_PG_PASSWORD}@pgvector_db:5432/${MEMGPT_PG_DB}
- POSTGRES_URI=postgresql://${MEMGPT_PG_USER}:${MEMGPT_PG_PASSWORD}@pgvector_db:5432/${MEMGPT_PG_DB} # TODO: deprecate
- MEMGPT_SERVER_PASS=${MEMGPT_SERVER_PASS} # memgpt server password
- MEMGPT_PG_DB=${MEMGPT_PG_DB}
- MEMGPT_PG_USER=${MEMGPT_PG_USER}
- MEMGPT_PG_PASSWORD=${MEMGPT_PG_PASSWORD}
- MEMGPT_PG_URL=pgvector_db
- MEMGPT_PG_HOST=pgvector_db
- MEMGPT_PG_PORT=5432
volumes:
- ./configs/server_config.yaml:/root/.memgpt/config # config file
- ~/.memgpt/credentials:/root/.memgpt/credentials # credentials file
Expand Down
9 changes: 4 additions & 5 deletions configs/server_config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -18,21 +18,20 @@ embedding_chunk_size = 300
[archival_storage]
type = postgres
path = /root/.memgpt/chroma
uri = postgresql+pg8000://swis_memgpt:swis_memgpt@memgpt-db/swis_memgpt
uri = postgresql+pg8000://memgpt:memgpt@pgvector_db:5432/memgpt

[recall_storage]
type = postgres
path = /root/.memgpt
uri = postgresql+pg8000://swis_memgpt:swis_memgpt@memgpt-db/swis_memgpt
uri = postgresql+pg8000://memgpt:memgpt@pgvector_db:5432/memgpt

[metadata_storage]
type = postgres
path = /root/.memgpt
uri = postgresql+pg8000://swis_memgpt:swis_memgpt@memgpt-db/swis_memgpt
uri = postgresql+pg8000://memgpt:memgpt@pgvector_db:5432/memgpt

[version]
memgpt_version = 0.3.10
memgpt_version = 0.3.11

[client]
anon_clientid = 00000000-0000-0000-0000-000000000000

48 changes: 48 additions & 0 deletions dev-compose.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
services:
memgpt_db:
image: ankane/pgvector:v0.5.1
networks:
default:
aliases:
- pgvector_db
- memgpt-db
environment:
- POSTGRES_USER=${MEMGPT_PG_USER}
- POSTGRES_PASSWORD=${MEMGPT_PG_PASSWORD}
- POSTGRES_DB=${MEMGPT_PG_DB}
volumes:
- ./.persist/pgdata:/var/lib/postgresql/data
- ./init.sql:/docker-entrypoint-initdb.d/init.sql
ports:
- "5432:5432"
memgpt_server:
#image: memgpt/memgpt-server:latest
image: memgpt-server
hostname: memgpt-server
build:
context: .
dockerfile: Dockerfile
depends_on:
- memgpt_db
ports:
- "8083:8083"
- "8283:8283"
env_file:
- .env
environment:
- MEMGPT_SERVER_PASS=${MEMGPT_SERVER_PASS} # memgpt server password
- MEMGPT_PG_DB=${MEMGPT_PG_DB}
- MEMGPT_PG_USER=${MEMGPT_PG_USER}
- MEMGPT_PG_PASSWORD=${MEMGPT_PG_PASSWORD}
- MEMGPT_PG_HOST=pgvector_db
- MEMGPT_PG_PORT=5432
volumes:
- ./configs/server_config.yaml:/root/.memgpt/config # config file
- ~/.memgpt/credentials:/root/.memgpt/credentials # credentials file
memgpt_nginx:
hostname: memgpt-nginx
image: nginx:stable-alpine3.17-slim
volumes:
- ./nginx.conf:/etc/nginx/nginx.conf
ports:
- "80:80"
2 changes: 2 additions & 0 deletions memgpt/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -387,6 +387,8 @@ def get_user_from_api_key(self, api_key: str) -> Optional[User]:
def create_agent(self, agent: AgentState):
# insert into agent table
# make sure agent.name does not already exist for user user_id
assert agent.state is not None, "Agent state must be provided"
assert len(list(agent.state.keys())) > 0, "Agent state must not be empty"
with self.session_maker() as session:
if session.query(AgentModel).filter(AgentModel.name == agent.name).filter(AgentModel.user_id == agent.user_id).count() > 0:
raise ValueError(f"Agent with name {agent.name} already exists")
Expand Down
13 changes: 13 additions & 0 deletions memgpt/server/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -247,6 +247,19 @@ def __init__(
# Initialize the metadata store
self.ms = MetadataStore(self.config)

# pre-fill database (users, presets, humans, personas)
# TODO: figure out how to handle default users (server is technically multi-user)
user_id = uuid.UUID(self.config.anon_clientid)
user = User(
id=uuid.UUID(self.config.anon_clientid),
)
if self.ms.get_user(user_id):
# update user
self.ms.update_user(user)
else:
self.ms.create_user(user)
presets.add_default_presets(user_id, self.ms)

# NOTE: removed, since server should be multi-user
## Create the default user
# base_user_id = uuid.UUID(self.config.anon_clientid)
Expand Down
2 changes: 1 addition & 1 deletion tests/test_load_archival.py
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,7 @@ def test_load_directory(
)
ms.delete_user(user.id)
ms.create_user(user)
ms.create_agent(agent)
# ms.create_agent(agent)
user = ms.get_user(user.id)
print("Got user:", user, embedding_config)

Expand Down
7 changes: 7 additions & 0 deletions tests/test_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -190,6 +190,13 @@ def test_storage(
human=get_human_text(TEST_MEMGPT_CONFIG.human),
llm_config=TEST_MEMGPT_CONFIG.default_llm_config,
embedding_config=TEST_MEMGPT_CONFIG.default_embedding_config,
state={
"persona": "",
"human": "",
"system": "",
"functions": [],
"messages": [],
},
)
ms.create_user(user)
ms.create_agent(agent)
Expand Down

0 comments on commit 6767153

Please sign in to comment.