From 200107b0f74766d7813fa0776a16c1fc9128b154 Mon Sep 17 00:00:00 2001 From: Bakar Tavadze Date: Tue, 26 Mar 2024 15:47:45 +0400 Subject: [PATCH] Enable migrating data from Redis to Postgres. --- Dockerfile | 2 + README.md | 37 +++++ backend/Makefile | 3 + backend/poetry.lock | 22 ++- backend/pyproject.toml | 1 + backend/redis_to_postgres.py | 277 +++++++++++++++++++++++++++++++++++ docker-compose-prod.yml | 29 +++- docker-compose.yml | 8 + 8 files changed, 375 insertions(+), 4 deletions(-) create mode 100644 backend/redis_to_postgres.py diff --git a/Dockerfile b/Dockerfile index bc626e68..2ecace33 100644 --- a/Dockerfile +++ b/Dockerfile @@ -45,4 +45,6 @@ COPY ./backend . # Copy the frontend build COPY --from=builder /frontend/dist ./ui +HEALTHCHECK --interval=30s --timeout=5s --start-period=10s --start-interval=1s --retries=3 CMD [ "curl", "-f", "http://localhost:8000/health" ] + ENTRYPOINT [ "uvicorn", "app.server:app", "--host", "0.0.0.0" ] diff --git a/README.md b/README.md index aaa5e162..240ad8d7 100644 --- a/README.md +++ b/README.md @@ -172,6 +172,43 @@ yarn dev Navigate to [http://localhost:5173/](http://localhost:5173/) and enjoy! +## Migrating data from Redis to Postgres + +OpenGPTs previously used Redis for data persistence, but has since switched to Postgres. If you have data in Redis that you would like to migrate to Postgres, you can use the following steps: + +### With Docker + +Add Postgres' environment variables to the `.env` file (by following `.env.example`). Then, start the services with `docker compose` and run the following command to migrate data from Redis to Postgres: + +```shell +docker exec -it opengpts-backend make redis_to_postgres +``` + +### Without Docker + +Make sure the following environment variables are set: + +```shell +export POSTGRES_HOST=... +export POSTGRES_PORT=... +export POSTGRES_DB=... +export POSTGRES_USER=... +export POSTGRES_PASSWORD=... +export REDIS_URL=... +``` + +Install [golang-migrate](https://github.com/golang-migrate/migrate) on your machine if you haven't already and run database schema migrations: + +```shell +make migrate +``` + +Finally, run the following command to migrate data from Redis to Postgres: + +```shell +make redis_to_postgres +``` + ## Features As much as possible, we are striving for feature parity with OpenAI. diff --git a/backend/Makefile b/backend/Makefile index 3e8dc5b2..cae39759 100644 --- a/backend/Makefile +++ b/backend/Makefile @@ -19,6 +19,9 @@ start: migrate: migrate -database postgres://$(POSTGRES_USER):$(POSTGRES_PASSWORD)@$(POSTGRES_HOST):$(POSTGRES_PORT)/$(POSTGRES_DB)?sslmode=disable -path ./migrations up +redis_to_postgres: + poetry run python redis_to_postgres.py + test: # We need to update handling of env variables for tests YDC_API_KEY=placeholder OPENAI_API_KEY=placeholder poetry run pytest $(TEST_FILE) diff --git a/backend/poetry.lock b/backend/poetry.lock index 5061a20a..58cb8e0e 100644 --- a/backend/poetry.lock +++ b/backend/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.1 and should not be changed by hand. [[package]] name = "aiohttp" @@ -2994,6 +2994,24 @@ files = [ [package.extras] full = ["numpy"] +[[package]] +name = "redis" +version = "5.0.1" +description = "Python client for Redis database and key-value store" +optional = false +python-versions = ">=3.7" +files = [ + {file = "redis-5.0.1-py3-none-any.whl", hash = "sha256:ed4802971884ae19d640775ba3b03aa2e7bd5e8fb8dfaed2decce4d0fc48391f"}, + {file = "redis-5.0.1.tar.gz", hash = "sha256:0dab495cd5753069d3bc650a0dde8a8f9edde16fc5691b689a566eda58100d0f"}, +] + +[package.dependencies] +async-timeout = {version = ">=4.0.2", markers = "python_full_version <= \"3.11.2\""} + +[package.extras] +hiredis = ["hiredis (>=1.0.0)"] +ocsp = ["cryptography (>=36.0.1)", "pyopenssl (==20.0.1)", "requests (>=2.26.0)"] + [[package]] name = "regex" version = "2023.10.3" @@ -4166,4 +4184,4 @@ multidict = ">=4.0" [metadata] lock-version = "2.0" python-versions = "^3.9.0,<3.12" -content-hash = "6aba3d05838348cb038b98803ed0d66e70caaa08a56d5a16fdbea3c5c63ec073" +content-hash = "bc7c67284c2b95f2880a5bd9a1a6bff68ddb7484bfa04013874e9166899520e2" diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 148ef410..189e799f 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -43,6 +43,7 @@ pgvector = "^0.2.5" psycopg2-binary = "^2.9.9" asyncpg = "^0.29.0" langchain-core = "^0.1.33" +redis = "5.0.1" [tool.poetry.group.dev.dependencies] uvicorn = "^0.23.2" diff --git a/backend/redis_to_postgres.py b/backend/redis_to_postgres.py new file mode 100644 index 00000000..e010e978 --- /dev/null +++ b/backend/redis_to_postgres.py @@ -0,0 +1,277 @@ +import asyncio +import json +import logging +import os +import pickle +import struct +import uuid +from collections import defaultdict +from datetime import datetime +from typing import Any, Iterator, Optional + +import asyncpg +import orjson +from langchain.utilities.redis import get_client +from langchain_core.runnables import ConfigurableFieldSpec, RunnableConfig +from langgraph.checkpoint import BaseCheckpointSaver +from langgraph.checkpoint.base import ( + Checkpoint, + empty_checkpoint, +) +from redis.client import Redis as RedisType + +from app.checkpoint import PostgresCheckpoint +from app.lifespan import get_pg_pool, lifespan +from app.server import app + +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + +redis_client: RedisType = get_client(os.environ["REDIS_URL"], socket_keepalive=True) + +thread_hash_keys = ["assistant_id", "name", "updated_at"] +assistant_hash_keys = ["name", "config", "updated_at", "public"] +embedding_hash_keys = ["namespace", "source", "content_vector", "title", "content"] +public_user_id = "eef39817-c173-4eb6-8be4-f77cf37054fb" + + +def keys(match: str) -> Iterator[str]: + cursor = 0 + while True: + cursor, keys = redis_client.scan(cursor=cursor, match=match, count=100) + for key in keys: + yield key.decode("utf-8") + if cursor == 0: + break + + +def load(keys: list[str], values: list[bytes]) -> dict: + return {k: orjson.loads(v) if v is not None else None for k, v in zip(keys, values)} + + +class RedisCheckpoint(BaseCheckpointSaver): + class Config: + arbitrary_types_allowed = True + + @property + def config_specs(self) -> list[ConfigurableFieldSpec]: + return [ + ConfigurableFieldSpec( + id="user_id", + annotation=Optional[str], + name="User ID", + description=None, + default=None, + is_shared=True, + ), + ConfigurableFieldSpec( + id="thread_id", + annotation=Optional[str], + name="Thread ID", + description=None, + default=None, + is_shared=True, + ), + ] + + def _dump(self, mapping: dict[str, Any]) -> dict: + return { + k: pickle.dumps(v) if v is not None else None for k, v in mapping.items() + } + + def _load(self, mapping: dict[bytes, bytes]) -> dict: + return { + k.decode(): pickle.loads(v) if v is not None else None + for k, v in mapping.items() + } + + def _hash_key(self, config: RunnableConfig) -> str: + user_id = config["configurable"]["user_id"] + thread_id = config["configurable"]["thread_id"] + return f"opengpts:{user_id}:thread:{thread_id}:checkpoint" + + def get(self, config: RunnableConfig) -> Checkpoint | None: + value = self._load(redis_client.hgetall(self._hash_key(config))) + if value.get("v") == 1: + # langgraph version 1 + return value + elif value.get("__pregel_version") == 1: + # permchain version 1 + value.pop("__pregel_version") + value.pop("__pregel_ts") + checkpoint = empty_checkpoint() + if value.get("messages"): + checkpoint["channel_values"] = {"__root__": value["messages"][1]} + else: + checkpoint["channel_values"] = {} + for key in checkpoint["channel_values"]: + checkpoint["channel_versions"][key] = 1 + return checkpoint + else: + # unknown version + return None + + def put(self, config: RunnableConfig, checkpoint: Checkpoint) -> None: + return redis_client.hmset(self._hash_key(config), self._dump(checkpoint)) + + +async def migrate_assistants(conn: asyncpg.Connection) -> None: + logger.info("Migrating assistants.") + + for key in keys("opengpts:*:assistant:*"): + parts = key.split(":") + user_id, assistant_id = parts[1], parts[3] + if user_id == public_user_id: + continue + + values = redis_client.hmget(key, *assistant_hash_keys) + assistant = load(assistant_hash_keys, values) if any(values) else None + if assistant is not None: + await conn.execute( + ( + "INSERT INTO assistant (assistant_id, user_id, name, config, updated_at, public) " + "VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (assistant_id) DO UPDATE SET " + "user_id = EXCLUDED.user_id, name = EXCLUDED.name, config = EXCLUDED.config, " + "updated_at = EXCLUDED.updated_at, public = EXCLUDED.public;" + ), + assistant_id, + user_id, + assistant["name"], + assistant["config"], + datetime.fromisoformat(assistant["updated_at"]), + assistant["public"], + ) + logger.info(f"Migrated assistant {assistant_id} for user {user_id}.") + + +async def migrate_threads(conn: asyncpg.Connection) -> None: + logger.info("Migrating threads.") + + for key in keys("opengpts:*:thread:*"): + if key.endswith(":checkpoint"): + continue + + parts = key.split(":") + user_id, thread_id = parts[1], parts[3] + + values = redis_client.hmget(key, *thread_hash_keys) + thread = load(thread_hash_keys, values) if any(values) else None + if thread is not None: + await conn.execute( + ( + "INSERT INTO thread (thread_id, assistant_id, user_id, name, updated_at) " + "VALUES ($1, $2, $3, $4, $5) ON CONFLICT (thread_id) DO UPDATE SET " + "assistant_id = EXCLUDED.assistant_id, user_id = EXCLUDED.user_id, " + "name = EXCLUDED.name, updated_at = EXCLUDED.updated_at;" + ), + thread_id, + thread["assistant_id"], + user_id, + thread["name"], + datetime.fromisoformat(thread["updated_at"]), + ) + logger.info(f"Migrated thread {thread_id} for user {user_id}.") + + +async def migrate_checkpoints() -> None: + logger.info("Migrating checkpoints.") + + redis_checkpoint = RedisCheckpoint() + postgres_checkpoint = PostgresCheckpoint() + + for key in keys("opengpts:*:thread:*:checkpoint"): + parts = key.split(":") + user_id, thread_id = parts[1], parts[3] + config = {"configurable": {"user_id": user_id, "thread_id": thread_id}} + checkpoint = redis_checkpoint.get(config) + if checkpoint: + await postgres_checkpoint.aput(config, checkpoint) + logger.info( + f"Migrated checkpoint for thread {thread_id} for user {user_id}." + ) + + +async def migrate_embeddings(conn: asyncpg.Connection) -> None: + logger.info("Migrating embeddings.") + + custom_ids = defaultdict(lambda: str(uuid.uuid4())) + + def _get_custom_id(doc: dict) -> str: + """custom_id is unique for each namespace.""" + return custom_ids[doc["namespace"]] + + def _redis_to_postgres_vector(binary_data: bytes) -> list[float]: + """Deserialize binary data to a list of floats.""" + assert len(binary_data) == 4 * 1536, "Invalid binary data length." + format_str = "<" + "1536f" + return list(struct.unpack(format_str, binary_data)) + + def _load_doc(values: list) -> Optional[str]: + doc = {} + for k, v in zip(embedding_hash_keys, values): + if k == "content_vector": + doc[k] = _redis_to_postgres_vector(v) + else: + doc[k] = v.decode() if v is not None else None + return doc + + def _get_cmetadata(doc: dict) -> str: + return json.dumps( + { + "source": doc["source"] if doc["source"] else None, + "namespace": doc["namespace"], + "title": doc["title"], + } + ) + + def _get_document(doc: dict) -> str: + """Sanitize the content by replacing null bytes.""" + return doc["content"].replace("\x00", "x") + + def _get_embedding(doc: dict) -> str: + return str(doc["content_vector"]) + + default_collection = await conn.fetchrow( + "SELECT uuid FROM langchain_pg_collection WHERE name = $1;", "langchain" + ) + assert ( + default_collection is not None + ), "Default collection not found in the database." + + for key in keys("doc:*"): + values = redis_client.hmget(key, *embedding_hash_keys) + doc = _load_doc(values) + await conn.execute( + ( + "INSERT INTO langchain_pg_embedding (document, collection_id, cmetadata, custom_id, embedding, uuid) " + "VALUES ($1, $2, $3, $4, $5, $6);" + ), + _get_document(doc), + default_collection["uuid"], + _get_cmetadata(doc), + _get_custom_id(doc), + _get_embedding(doc), + str(uuid.uuid4()), + ) + logger.info(f"Migrated embedding for namespace {doc['namespace']}.") + + +async def migrate_data(): + logger.info("Starting to migrate data from Redis to Postgres.") + async with get_pg_pool().acquire() as conn, conn.transaction(): + await migrate_assistants(conn) + await migrate_threads(conn) + await migrate_checkpoints() + await migrate_embeddings(conn) + logger.info("Data was migrated successfully.") + + +async def main(): + async with lifespan(app): + await migrate_data() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/docker-compose-prod.yml b/docker-compose-prod.yml index b8bab2a8..cf419b07 100644 --- a/docker-compose-prod.yml +++ b/docker-compose-prod.yml @@ -1,22 +1,47 @@ version: "3" services: + redis: + container_name: opengpts-redis + image: redis/redis-stack-server:latest + ports: + - "6379:6379" + volumes: + - ./redis-volume:/data postgres: image: pgvector/pgvector:pg16 + healthcheck: + test: pg_isready -U $POSTGRES_USER + start_interval: 1s + start_period: 5s + interval: 5s + retries: 5 ports: - - "5432:5432" + - "5433:5432" env_file: - .env volumes: - ./postgres-volume:/var/lib/postgresql/data + postgres-setup: + image: migrate/migrate + depends_on: + postgres: + condition: service_healthy + volumes: + - ./backend/migrations:/migrations + env_file: + - .env + command: ["-path", "/migrations", "-database", "postgres://$POSTGRES_USER:$POSTGRES_PASSWORD@postgres:$POSTGRES_PORT/$POSTGRES_DB?sslmode=disable", "up"] backend: container_name: opengpts-backend image: docker.io/langchain/open-gpts:latest ports: - "8100:8000" # Backend is accessible on localhost:8100 and serves the frontend depends_on: - - postgres + postgres-setup: + condition: service_completed_successfully env_file: - .env environment: POSTGRES_HOST: "postgres" + REDIS_URL: "redis://opengpts-redis:6379" diff --git a/docker-compose.yml b/docker-compose.yml index 6d59cb42..c4000a98 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,6 +1,13 @@ version: "3" services: + redis: + container_name: opengpts-redis + image: redis/redis-stack-server:latest + ports: + - "6379:6379" + volumes: + - ./redis-volume:/data postgres: image: pgvector/pgvector:pg16 healthcheck: @@ -40,6 +47,7 @@ services: - ./backend:/backend environment: POSTGRES_HOST: "postgres" + REDIS_URL: "redis://opengpts-redis:6379" command: - --reload frontend: