Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

feat: client tests for docker compose server (#1189)

Browse Source
This commit is contained in:
Sarah Wooders
2024-03-26 18:58:00 -07:00
committed by GitHub
parent c8297f53e4
commit c0bd66c957
15 changed files with 77 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
.github/workflows/tests.yml vendored
View File

@@ -1,7 +1,7 @@
name: Run All pytest Tests name: Run All pytest Tests
env: env:
PGVECTOR_TEST_DB_URL: ${{ secrets.PGVECTOR_TEST_DB_URL }} MEMGPT_PGURI: ${{ secrets.MEMGPT_PGURI }}
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
on: on:
@@ -21,6 +21,7 @@ jobs:
- name: Build and run container - name: Build and run container
run: bash db/run_postgres.sh run: bash db/run_postgres.sh
- name: "Setup Python, Poetry and Dependencies" - name: "Setup Python, Poetry and Dependencies"
uses: packetcoders/action-setup-cache-python-poetry@main uses: packetcoders/action-setup-cache-python-poetry@main
with: with:
@@ -28,9 +29,25 @@ jobs:
poetry-version: "1.7.1" poetry-version: "1.7.1"
install-args: "--all-extras" install-args: "--all-extras"
- name: Initialize credentials
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
run: |
poetry run memgpt quickstart --backend openai
- name: Run docker compose server
env:
MEMGPT_PG_DB: memgpt
MEMGPT_PG_USER: memgpt
MEMGPT_PG_PASSWORD: memgpt
MEMGPT_SERVER_PASS: test_server_token
MEMGPT_CONFIG_PATH: configs/server_config.yaml
run: docker compose up &
- name: Run server tests - name: Run server tests
env: env:
PGVECTOR_TEST_DB_URL: postgresql+pg8000://memgpt:memgpt@localhost:5432/memgpt MEMGPT_PGURI: postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
MEMGPT_SERVER_PASS: test_server_token MEMGPT_SERVER_PASS: test_server_token
run: | run: |
@@ -38,7 +55,7 @@ jobs:
- name: Run tests with pytest - name: Run tests with pytest
env: env:
PGVECTOR_TEST_DB_URL: postgresql+pg8000://memgpt:memgpt@localhost:5432/memgpt MEMGPT_PGURI: postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
MEMGPT_SERVER_PASS: test_server_token MEMGPT_SERVER_PASS: test_server_token
PYTHONPATH: ${{ github.workspace }}:${{ env.PYTHONPATH }} PYTHONPATH: ${{ github.workspace }}:${{ env.PYTHONPATH }}
@@ -47,7 +64,7 @@ jobs:
- name: Run storage tests - name: Run storage tests
env: env:
PGVECTOR_TEST_DB_URL: postgresql+pg8000://memgpt:memgpt@localhost:5432/memgpt MEMGPT_PGURI: postgresql+pg8000://memgpt:memgpt@localhost:5432/memgpt
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
MEMGPT_SERVER_PASS: test_server_token MEMGPT_SERVER_PASS: test_server_token
run: | run: |

6
compose.yaml
View File

@@ -13,12 +13,16 @@ services:
- "5432:5432" - "5432:5432"
memgpt_server: memgpt_server:
image: memgpt/memgpt-server:0.3.6 image: memgpt/memgpt-server:0.3.7
depends_on: depends_on:
- pgvector_db - pgvector_db
environment: environment:
- POSTGRES_URI=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}
- MEMGPT_SERVER_PASS=${MEMGPT_SERVER_PASS} # memgpt server password - 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
volumes: volumes:
- ./configs/server_config.yaml:/root/.memgpt/config # config file - ./configs/server_config.yaml:/root/.memgpt/config # config file
- ~/.memgpt/credentials:/root/.memgpt/credentials # credentials file - ~/.memgpt/credentials:/root/.memgpt/credentials # credentials file

2
db/run_postgres.sh
View File

@@ -4,7 +4,7 @@ docker build -f db/Dockerfile.simple -t pg-test .
# run container # run container
docker run -d --rm \ docker run -d --rm \
--name memgpt-db-test \ --name memgpt-db-test \
-p 5432:5432 \ -p 8888:5432 \
-e POSTGRES_PASSWORD=password \ -e POSTGRES_PASSWORD=password \
-v memgpt_db_test:/var/lib/postgresql/data \ -v memgpt_db_test:/var/lib/postgresql/data \
pg-test:latest pg-test:latest

4
docs/storage.md
View File

@@ -38,13 +38,13 @@ To run the Postgres backend, you will need a URI to a Postgres database that sup
- Add the following line to your shell profile (e.g., `~/.bashrc`, `~/.zshrc`): - Add the following line to your shell profile (e.g., `~/.bashrc`, `~/.zshrc`):
```sh ```sh
export PGVECTOR_TEST_DB_URL=postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt export MEMGPT_PGURI=postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
``` ```
- Or create a `.env` file in the root project directory with: - Or create a `.env` file in the root project directory with:
```sh ```sh
PGVECTOR_TEST_DB_URL=postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt MEMGPT_PGURI=postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt
``` ```
4. Run the script from the root project directory: 4. Run the script from the root project directory:

1
init.sql
View File

@@ -6,7 +6,6 @@
\set db_password `([ -r /var/run/secrets/memgpt-password ] && cat /var/run/secrets/memgpt-password) || echo "${POSTGRES_PASSWORD:-memgpt}"` \set db_password `([ -r /var/run/secrets/memgpt-password ] && cat /var/run/secrets/memgpt-password) || echo "${POSTGRES_PASSWORD:-memgpt}"`
\set db_name `([ -r /var/run/secrets/memgpt-db ] && cat /var/run/secrets/memgpt-db) || echo "${POSTGRES_DB:-memgpt}"` \set db_name `([ -r /var/run/secrets/memgpt-db ] && cat /var/run/secrets/memgpt-db) || echo "${POSTGRES_DB:-memgpt}"`
-- CREATE USER :"db_user" -- CREATE USER :"db_user"
-- WITH PASSWORD :'db_password' -- WITH PASSWORD :'db_password'
-- NOCREATEDB -- NOCREATEDB

17
memgpt/agent_store/db.py
View File

@@ -439,15 +439,18 @@ class PostgresStorageConnector(SQLStorageConnector):
self.db_model = get_db_model(config, self.table_name, table_type, user_id, agent_id) self.db_model = get_db_model(config, self.table_name, table_type, user_id, agent_id)
# construct URI from enviornment variables # construct URI from enviornment variables
db = os.getenv("MEMGPT_PG_DB", "memgpt") if os.getenv("MEMGPT_PGURI"):
user = os.getenv("MEMGPT_PG_USER", "memgpt") self.uri = os.getenv("MEMGPT_PGURI")
password = os.getenv("MEMGPT_PG_PASSWORD", "memgpt") else:
port = os.getenv("MEMGPT_PG_PORT", "5432") db = os.getenv("MEMGPT_PG_DB", "memgpt")
url = os.getenv("MEMGPT_PG_URL", "localhost") user = os.getenv("MEMGPT_PG_USER", "memgpt")
uri = f"postgresql+pg8000://{user}:{password}@{url}:{port}/{db}" password = os.getenv("MEMGPT_PG_PASSWORD", "memgpt")
port = os.getenv("MEMGPT_PG_PORT", "5432")
url = os.getenv("MEMGPT_PG_URL", "localhost")
self.uri = f"postgresql+pg8000://{user}:{password}@{url}:{port}/{db}"
# create engine # create engine
self.engine = create_engine(uri) self.engine = create_engine(self.uri)
for c in self.db_model.__table__.columns: for c in self.db_model.__table__.columns:
if c.name == "embedding": if c.name == "embedding":

15
memgpt/metadata.py
View File

@@ -305,12 +305,15 @@ class MetadataStore:
# TODO: get DB URI or path # TODO: get DB URI or path
if config.metadata_storage_type == "postgres": if config.metadata_storage_type == "postgres":
# construct URI from enviornment variables # construct URI from enviornment variables
db = os.getenv("MEMGPT_PG_DB", "memgpt") if os.getenv("MEMGPT_PGURI"):
user = os.getenv("MEMGPT_PG_USER", "memgpt") self.uri = os.getenv("MEMGPT_PGURI")
password = os.getenv("MEMGPT_PG_PASSWORD", "memgpt") else:
port = os.getenv("MEMGPT_PG_PORT", "5432") db = os.getenv("MEMGPT_PG_DB", "memgpt")
url = os.getenv("MEMGPT_PG_URL", "localhost") user = os.getenv("MEMGPT_PG_USER", "memgpt")
self.uri = f"postgresql+pg8000://{user}:{password}@{url}:{port}/{db}" password = os.getenv("MEMGPT_PG_PASSWORD", "memgpt")
port = os.getenv("MEMGPT_PG_PORT", "5432")
url = os.getenv("MEMGPT_PG_URL", "localhost")
self.uri = f"postgresql+pg8000://{user}:{password}@{url}:{port}/{db}"
elif config.metadata_storage_type == "sqlite": elif config.metadata_storage_type == "sqlite":
path = os.path.join(config.metadata_storage_path, "sqlite.db") path = os.path.join(config.metadata_storage_path, "sqlite.db")
self.uri = f"sqlite:///{path}" self.uri = f"sqlite:///{path}"

4
scripts/pack_docker.sh Normal file
View File

@@ -0,0 +1,4 @@
export MEMGPT_VERSION=$(memgpt version)
#docker build -t memgpt/memgpt-server:${MEMGPT_VERSION} --platform linux/x86_64 .
docker buildx build --platform=linux/amd64,linux/arm64,linux/x86_64 -t memgpt/memgpt-server:${MEMGPT_VERSION} .
docker push memgpt/memgpt-server:${MEMGPT_VERSION}

2
tests/clear_postgres_db.py
View File

@@ -5,7 +5,7 @@ from sqlalchemy import create_engine, MetaData
def main(): def main():
uri = os.environ.get( uri = os.environ.get(
"PGVECTOR_TEST_DB_URL", "MEMGPT_PGURI",
"postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt", "postgresql+pg8000://memgpt:memgpt@localhost:8888/memgpt",
) )

4
tests/test_admin_client.py
View File

@@ -32,8 +32,8 @@ def run_server():
load_dotenv() load_dotenv()
# Use os.getenv with a fallback to os.environ.get # Use os.getenv with a fallback to os.environ.get
db_url = os.getenv("PGVECTOR_TEST_DB_URL") or os.environ.get("PGVECTOR_TEST_DB_URL") db_url = os.getenv("MEMGPT_PGURI") or os.environ.get("MEMGPT_PGURI")
assert db_url, "Missing PGVECTOR_TEST_DB_URL" assert db_url, "Missing MEMGPT_PGURI"
if os.getenv("OPENAI_API_KEY"): if os.getenv("OPENAI_API_KEY"):
config = TestMGPTConfig( config = TestMGPTConfig(

12
tests/test_client.py
View File

@@ -43,8 +43,8 @@ def run_server():
load_dotenv() load_dotenv()
# Use os.getenv with a fallback to os.environ.get # Use os.getenv with a fallback to os.environ.get
db_url = os.getenv("PGVECTOR_TEST_DB_URL") or os.environ.get("PGVECTOR_TEST_DB_URL") db_url = os.getenv("MEMGPT_PGURI") or os.environ.get("MEMGPT_PGURI")
assert db_url, "Missing PGVECTOR_TEST_DB_URL" assert db_url, "Missing MEMGPT_PGURI"
if os.getenv("OPENAI_API_KEY"): if os.getenv("OPENAI_API_KEY"):
config = TestMGPTConfig( config = TestMGPTConfig(
@@ -97,6 +97,7 @@ def run_server():
config.save() config.save()
credentials.save() credentials.save()
print("Starting server...")
start_server(debug=True) start_server(debug=True)
@@ -104,22 +105,23 @@ def run_server():
@pytest.fixture( @pytest.fixture(
params=[ params=[
{"base_url": local_service_url}, {"base_url": local_service_url},
# {"base_url": docker_compose_url}, # TODO: add when docker compose added to tests {"base_url": docker_compose_url}, # TODO: add when docker compose added to tests
# {"base_url": None} # TODO: add when implemented # {"base_url": None} # TODO: add when implemented
], ],
scope="module", scope="module",
) )
# @pytest.fixture(params=[{"base_url": test_base_url}], scope="module") # @pytest.fixture(params=[{"base_url": test_base_url}], scope="module")
def client(request): def client(request):
print("CLIENT", request.param["base_url"])
if request.param["base_url"]: if request.param["base_url"]:
if request.param["base_url"] == local_service_url: if request.param["base_url"] == local_service_url:
# start server # start server
print("Starting server...") print("Starting server thread")
thread = threading.Thread(target=run_server, daemon=True) thread = threading.Thread(target=run_server, daemon=True)
thread.start() thread.start()
time.sleep(5) time.sleep(5)
admin = Admin(local_service_url, test_server_token) admin = Admin(request.param["base_url"], test_server_token)
response = admin.create_user(test_user_id) # Adjust as per your client's method response = admin.create_user(test_user_id) # Adjust as per your client's method
user_id = response.user_id user_id = response.user_id
token = response.api_key token = response.api_key

8
tests/test_load_archival.py
View File

@@ -46,10 +46,10 @@ def test_load_directory(
wipe_config() wipe_config()
# setup config # setup config
if metadata_storage_connector == "postgres": if metadata_storage_connector == "postgres":
if not os.getenv("PGVECTOR_TEST_DB_URL"): if not os.getenv("MEMGPT_PGURI"):
print("Skipping test, missing PG URI") print("Skipping test, missing PG URI")
return return
TEST_MEMGPT_CONFIG.metadata_storage_uri = os.getenv("PGVECTOR_TEST_DB_URL") TEST_MEMGPT_CONFIG.metadata_storage_uri = os.getenv("MEMGPT_PGURI")
TEST_MEMGPT_CONFIG.metadata_storage_type = "postgres" TEST_MEMGPT_CONFIG.metadata_storage_type = "postgres"
elif metadata_storage_connector == "sqlite": elif metadata_storage_connector == "sqlite":
print("testing sqlite metadata") print("testing sqlite metadata")
@@ -57,10 +57,10 @@ def test_load_directory(
else: else:
raise NotImplementedError(f"Storage type {metadata_storage_connector} not implemented") raise NotImplementedError(f"Storage type {metadata_storage_connector} not implemented")
if passage_storage_connector == "postgres": if passage_storage_connector == "postgres":
if not os.getenv("PGVECTOR_TEST_DB_URL"): if not os.getenv("MEMGPT_PGURI"):
print("Skipping test, missing PG URI") print("Skipping test, missing PG URI")
return return
TEST_MEMGPT_CONFIG.archival_storage_uri = os.getenv("PGVECTOR_TEST_DB_URL") TEST_MEMGPT_CONFIG.archival_storage_uri = os.getenv("MEMGPT_PGURI")
TEST_MEMGPT_CONFIG.archival_storage_type = "postgres" TEST_MEMGPT_CONFIG.archival_storage_type = "postgres"
elif passage_storage_connector == "chroma": elif passage_storage_connector == "chroma":
print("testing chroma passage storage") print("testing chroma passage storage")

12
tests/test_metadata_store.py
View File

@@ -16,11 +16,11 @@ from memgpt.models.pydantic_models import HumanModel, PersonaModel
@pytest.mark.parametrize("storage_connector", ["sqlite"]) @pytest.mark.parametrize("storage_connector", ["sqlite"])
def test_storage(storage_connector): def test_storage(storage_connector):
if storage_connector == "postgres": if storage_connector == "postgres":
if not os.getenv("PGVECTOR_TEST_DB_URL"): if not os.getenv("MEMGPT_PGURI"):
print("Skipping test, missing PG URI") print("Skipping test, missing PG URI")
return return
TEST_MEMGPT_CONFIG.archival_storage_uri = os.environ["PGVECTOR_TEST_DB_URL"] TEST_MEMGPT_CONFIG.archival_storage_uri = os.environ["MEMGPT_PGURI"]
TEST_MEMGPT_CONFIG.recall_storage_uri = os.environ["PGVECTOR_TEST_DB_URL"] TEST_MEMGPT_CONFIG.recall_storage_uri = os.environ["MEMGPT_PGURI"]
TEST_MEMGPT_CONFIG.archival_storage_type = "postgres" TEST_MEMGPT_CONFIG.archival_storage_type = "postgres"
TEST_MEMGPT_CONFIG.recall_storage_type = "postgres" TEST_MEMGPT_CONFIG.recall_storage_type = "postgres"
if storage_connector == "sqlite": if storage_connector == "sqlite":
@@ -49,8 +49,8 @@ def test_storage(storage_connector):
user_id=user_1.id, user_id=user_1.id,
name="agent_1", name="agent_1",
preset=DEFAULT_PRESET, preset=DEFAULT_PRESET,
persona=get_persona_text(DEFAULT_PERSONA), persona=DEFAULT_PERSONA,
human=get_human_text(DEFAULT_HUMAN), human=DEFAULT_HUMAN,
llm_config=TEST_MEMGPT_CONFIG.default_llm_config, llm_config=TEST_MEMGPT_CONFIG.default_llm_config,
embedding_config=TEST_MEMGPT_CONFIG.default_embedding_config, embedding_config=TEST_MEMGPT_CONFIG.default_embedding_config,
) )
@@ -73,7 +73,7 @@ def test_storage(storage_connector):
from memgpt.presets.presets import add_default_presets from memgpt.presets.presets import add_default_presets
add_default_presets(user_1.id, ms) add_default_presets(user_1.id, ms)
preset_obj = ms.get_preset(preset_name=DEFAULT_PRESET, user_id=user_1.id) preset_obj = ms.get_preset(name=DEFAULT_PRESET, user_id=user_1.id)
from memgpt.interface import CLIInterface as interface # for printing to terminal from memgpt.interface import CLIInterface as interface # for printing to terminal
# Overwrite fields in the preset if they were specified # Overwrite fields in the preset if they were specified

4
tests/test_server.py
View File

@@ -20,8 +20,8 @@ def server():
wipe_memgpt_home() wipe_memgpt_home()
# Use os.getenv with a fallback to os.environ.get # Use os.getenv with a fallback to os.environ.get
db_url = os.getenv("PGVECTOR_TEST_DB_URL") or os.environ.get("PGVECTOR_TEST_DB_URL") db_url = os.getenv("MEMGPT_PGURI") or os.environ.get("MEMGPT_PGURI")
assert db_url, "Missing PGVECTOR_TEST_DB_URL" assert db_url, "Missing MEMGPT_PGURI"
if os.getenv("OPENAI_API_KEY"): if os.getenv("OPENAI_API_KEY"):
config = TestMGPTConfig( config = TestMGPTConfig(

6
tests/test_storage.py
View File

@@ -122,11 +122,11 @@ def test_storage(
# del globals()['Message'] # del globals()['Message']
if storage_connector == "postgres": if storage_connector == "postgres":
if not os.getenv("PGVECTOR_TEST_DB_URL"): if not os.getenv("MEMGPT_PGURI"):
print("Skipping test, missing PG URI") print("Skipping test, missing PG URI")
return return
TEST_MEMGPT_CONFIG.archival_storage_uri = os.environ["PGVECTOR_TEST_DB_URL"] TEST_MEMGPT_CONFIG.archival_storage_uri = os.environ["MEMGPT_PGURI"]
TEST_MEMGPT_CONFIG.recall_storage_uri = os.environ["PGVECTOR_TEST_DB_URL"] TEST_MEMGPT_CONFIG.recall_storage_uri = os.environ["MEMGPT_PGURI"]
TEST_MEMGPT_CONFIG.archival_storage_type = "postgres" TEST_MEMGPT_CONFIG.archival_storage_type = "postgres"
TEST_MEMGPT_CONFIG.recall_storage_type = "postgres" TEST_MEMGPT_CONFIG.recall_storage_type = "postgres"
if storage_connector == "lancedb": if storage_connector == "lancedb":