Add memory export / import feature (#3345)
Co-authored-by: Vikram Iyer <vikramiyer@Vikrams-MacBook-Pro.local>
This commit is contained in:
VikramIyer125
@@ -1,6 +1,7 @@
|
||||
from .apps import router as apps_router
|
||||
from .backup import router as backup_router
|
||||
from .config import router as config_router
|
||||
from .memories import router as memories_router
|
||||
from .stats import router as stats_router
|
||||
|
||||
__all__ = ["memories_router", "apps_router", "stats_router", "config_router"]
|
||||
__all__ = ["memories_router", "apps_router", "stats_router", "config_router", "backup_router"]
|
||||
|
||||
@@ -0,0 +1,499 @@
|
||||
from datetime import UTC, datetime
|
||||
import io
|
||||
import json
|
||||
import gzip
|
||||
import zipfile
|
||||
from typing import Optional, List, Dict, Any
|
||||
from uuid import UUID
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Query, Form
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
from sqlalchemy.orm import Session, joinedload
|
||||
from sqlalchemy import and_
|
||||
|
||||
from app.database import get_db
|
||||
from app.models import (
|
||||
User, App, Memory, MemoryState, Category, memory_categories,
|
||||
MemoryStatusHistory, AccessControl
|
||||
)
|
||||
from app.utils.memory import get_memory_client
|
||||
|
||||
from uuid import uuid4
|
||||
|
||||
router = APIRouter(prefix="/api/v1/backup", tags=["backup"])
|
||||
|
||||
class ExportRequest(BaseModel):
|
||||
user_id: str
|
||||
app_id: Optional[UUID] = None
|
||||
from_date: Optional[int] = None
|
||||
to_date: Optional[int] = None
|
||||
include_vectors: bool = True
|
||||
|
||||
def _iso(dt: Optional[datetime]) -> Optional[str]:
|
||||
if isinstance(dt, datetime):
|
||||
try:
|
||||
return dt.astimezone(UTC).isoformat()
|
||||
except:
|
||||
return dt.replace(tzinfo=UTC).isoformat()
|
||||
return None
|
||||
|
||||
def _parse_iso(dt: Optional[str]) -> Optional[datetime]:
|
||||
if not dt:
|
||||
return None
|
||||
try:
|
||||
return datetime.fromisoformat(dt)
|
||||
except Exception:
|
||||
try:
|
||||
return datetime.fromisoformat(dt.replace("Z", "+00:00"))
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
def _export_sqlite(db: Session, req: ExportRequest) -> Dict[str, Any]:
|
||||
user = db.query(User).filter(User.user_id == req.user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(status_code=404, detail="User not found")
|
||||
|
||||
time_filters = []
|
||||
if req.from_date:
|
||||
time_filters.append(Memory.created_at >= datetime.fromtimestamp(req.from_date, tz=UTC))
|
||||
if req.to_date:
|
||||
time_filters.append(Memory.created_at <= datetime.fromtimestamp(req.to_date, tz=UTC))
|
||||
|
||||
mem_q = (
|
||||
db.query(Memory)
|
||||
.options(joinedload(Memory.categories), joinedload(Memory.app))
|
||||
.filter(
|
||||
Memory.user_id == user.id,
|
||||
*(time_filters or []),
|
||||
* ( [Memory.app_id == req.app_id] if req.app_id else [] ),
|
||||
)
|
||||
)
|
||||
|
||||
memories = mem_q.all()
|
||||
memory_ids = [m.id for m in memories]
|
||||
|
||||
app_ids = sorted({m.app_id for m in memories if m.app_id})
|
||||
apps = db.query(App).filter(App.id.in_(app_ids)).all() if app_ids else []
|
||||
|
||||
cats = sorted({c for m in memories for c in m.categories}, key = lambda c: str(c.id))
|
||||
|
||||
mc_rows = db.execute(
|
||||
memory_categories.select().where(memory_categories.c.memory_id.in_(memory_ids))
|
||||
).fetchall() if memory_ids else []
|
||||
|
||||
history = db.query(MemoryStatusHistory).filter(MemoryStatusHistory.memory_id.in_(memory_ids)).all() if memory_ids else []
|
||||
|
||||
acls = db.query(AccessControl).filter(
|
||||
AccessControl.subject_type == "app",
|
||||
AccessControl.subject_id.in_(app_ids) if app_ids else False
|
||||
).all() if app_ids else []
|
||||
|
||||
return {
|
||||
"user": {
|
||||
"id": str(user.id),
|
||||
"user_id": user.user_id,
|
||||
"name": user.name,
|
||||
"email": user.email,
|
||||
"metadata": user.metadata_,
|
||||
"created_at": _iso(user.created_at),
|
||||
"updated_at": _iso(user.updated_at)
|
||||
},
|
||||
"apps": [
|
||||
{
|
||||
"id": str(a.id),
|
||||
"owner_id": str(a.owner_id),
|
||||
"name": a.name,
|
||||
"description": a.description,
|
||||
"metadata": a.metadata_,
|
||||
"is_active": a.is_active,
|
||||
"created_at": _iso(a.created_at),
|
||||
"updated_at": _iso(a.updated_at),
|
||||
}
|
||||
for a in apps
|
||||
],
|
||||
"categories": [
|
||||
{
|
||||
"id": str(c.id),
|
||||
"name": c.name,
|
||||
"description": c.description,
|
||||
"created_at": _iso(c.created_at),
|
||||
"updated_at": _iso(c.updated_at),
|
||||
}
|
||||
for c in cats
|
||||
],
|
||||
"memories": [
|
||||
{
|
||||
"id": str(m.id),
|
||||
"user_id": str(m.user_id),
|
||||
"app_id": str(m.app_id) if m.app_id else None,
|
||||
"content": m.content,
|
||||
"metadata": m.metadata_,
|
||||
"state": m.state.value,
|
||||
"created_at": _iso(m.created_at),
|
||||
"updated_at": _iso(m.updated_at),
|
||||
"archived_at": _iso(m.archived_at),
|
||||
"deleted_at": _iso(m.deleted_at),
|
||||
"category_ids": [str(c.id) for c in m.categories], #TODO: figure out a way to add category names simply to this
|
||||
}
|
||||
for m in memories
|
||||
],
|
||||
"memory_categories": [
|
||||
{"memory_id": str(r.memory_id), "category_id": str(r.category_id)}
|
||||
for r in mc_rows
|
||||
],
|
||||
"status_history": [
|
||||
{
|
||||
"id": str(h.id),
|
||||
"memory_id": str(h.memory_id),
|
||||
"changed_by": str(h.changed_by),
|
||||
"old_state": h.old_state.value,
|
||||
"new_state": h.new_state.value,
|
||||
"changed_at": _iso(h.changed_at),
|
||||
}
|
||||
for h in history
|
||||
],
|
||||
"access_controls": [
|
||||
{
|
||||
"id": str(ac.id),
|
||||
"subject_type": ac.subject_type,
|
||||
"subject_id": str(ac.subject_id) if ac.subject_id else None,
|
||||
"object_type": ac.object_type,
|
||||
"object_id": str(ac.object_id) if ac.object_id else None,
|
||||
"effect": ac.effect,
|
||||
"created_at": _iso(ac.created_at),
|
||||
}
|
||||
for ac in acls
|
||||
],
|
||||
"export_meta": {
|
||||
"app_id_filter": str(req.app_id) if req.app_id else None,
|
||||
"from_date": req.from_date,
|
||||
"to_date": req.to_date,
|
||||
"version": "1",
|
||||
"generated_at": datetime.now(UTC).isoformat(),
|
||||
},
|
||||
}
|
||||
|
||||
def _export_logical_memories_gz(
|
||||
db: Session,
|
||||
*,
|
||||
user_id: str,
|
||||
app_id: Optional[UUID] = None,
|
||||
from_date: Optional[int] = None,
|
||||
to_date: Optional[int] = None
|
||||
) -> bytes:
|
||||
"""
|
||||
Export a provider-agnostic backup of memories so they can be restored to any vector DB
|
||||
by re-embedding content. One JSON object per line, gzip-compressed.
|
||||
|
||||
Schema (per line):
|
||||
{
|
||||
"id": "<uuid>",
|
||||
"content": "<text>",
|
||||
"metadata": {...},
|
||||
"created_at": "<iso8601 or null>",
|
||||
"updated_at": "<iso8601 or null>",
|
||||
"state": "active|paused|archived|deleted",
|
||||
"app": "<app name or null>",
|
||||
"categories": ["catA", "catB", ...]
|
||||
}
|
||||
"""
|
||||
|
||||
user = db.query(User).filter(User.user_id == user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(status_code=404, detail="User not found")
|
||||
|
||||
time_filters = []
|
||||
if from_date:
|
||||
time_filters.append(Memory.created_at >= datetime.fromtimestamp(from_date, tz=UTC))
|
||||
if to_date:
|
||||
time_filters.append(Memory.created_at <= datetime.fromtimestamp(to_date, tz=UTC))
|
||||
|
||||
q = (
|
||||
db.query(Memory)
|
||||
.options(joinedload(Memory.categories), joinedload(Memory.app))
|
||||
.filter(
|
||||
Memory.user_id == user.id,
|
||||
*(time_filters or []),
|
||||
)
|
||||
)
|
||||
if app_id:
|
||||
q = q.filter(Memory.app_id == app_id)
|
||||
|
||||
buf = io.BytesIO()
|
||||
with gzip.GzipFile(fileobj=buf, mode="wb") as gz:
|
||||
for m in q.all():
|
||||
record = {
|
||||
"id": str(m.id),
|
||||
"content": m.content,
|
||||
"metadata": m.metadata_ or {},
|
||||
"created_at": _iso(m.created_at),
|
||||
"updated_at": _iso(m.updated_at),
|
||||
"state": m.state.value,
|
||||
"app": m.app.name if m.app else None,
|
||||
"categories": [c.name for c in m.categories],
|
||||
}
|
||||
gz.write((json.dumps(record) + "\n").encode("utf-8"))
|
||||
return buf.getvalue()
|
||||
|
||||
@router.post("/export")
|
||||
async def export_backup(req: ExportRequest, db: Session = Depends(get_db)):
|
||||
sqlite_payload = _export_sqlite(db=db, req=req)
|
||||
memories_blob = _export_logical_memories_gz(
|
||||
db=db,
|
||||
user_id=req.user_id,
|
||||
app_id=req.app_id,
|
||||
from_date=req.from_date,
|
||||
to_date=req.to_date,
|
||||
|
||||
)
|
||||
|
||||
#TODO: add vector store specific exports in future for speed
|
||||
|
||||
zip_buf = io.BytesIO()
|
||||
with zipfile.ZipFile(zip_buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
|
||||
zf.writestr("memories.json", json.dumps(sqlite_payload, indent=2))
|
||||
zf.writestr("memories.jsonl.gz", memories_blob)
|
||||
|
||||
zip_buf.seek(0)
|
||||
return StreamingResponse(
|
||||
zip_buf,
|
||||
media_type="application/zip",
|
||||
headers={"Content-Disposition": f'attachment; filename="memories_export_{req.user_id}.zip"'},
|
||||
)
|
||||
|
||||
@router.post("/import")
|
||||
async def import_backup(
|
||||
file: UploadFile = File(..., description="Zip with memories.json and memories.jsonl.gz"),
|
||||
user_id: str = Form(..., description="Import memories into this user_id"),
|
||||
mode: str = Query("overwrite"),
|
||||
db: Session = Depends(get_db)
|
||||
):
|
||||
if not file.filename.endswith(".zip"):
|
||||
raise HTTPException(status_code=400, detail="Expected a zip file.")
|
||||
|
||||
if mode not in {"skip", "overwrite"}:
|
||||
raise HTTPException(status_code=400, detail="Invalid mode. Must be 'skip' or 'overwrite'.")
|
||||
|
||||
user = db.query(User).filter(User.user_id == user_id).first()
|
||||
if not user:
|
||||
raise HTTPException(status_code=404, detail="User not found")
|
||||
|
||||
content = await file.read()
|
||||
try:
|
||||
with zipfile.ZipFile(io.BytesIO(content), "r") as zf:
|
||||
names = zf.namelist()
|
||||
|
||||
def find_member(filename: str) -> Optional[str]:
|
||||
for name in names:
|
||||
# Skip directory entries
|
||||
if name.endswith('/'):
|
||||
continue
|
||||
if name.rsplit('/', 1)[-1] == filename:
|
||||
return name
|
||||
return None
|
||||
|
||||
sqlite_member = find_member("memories.json")
|
||||
if not sqlite_member:
|
||||
raise HTTPException(status_code=400, detail="memories.json missing in zip")
|
||||
|
||||
memories_member = find_member("memories.jsonl.gz")
|
||||
|
||||
sqlite_data = json.loads(zf.read(sqlite_member))
|
||||
memories_blob = zf.read(memories_member) if memories_member else None
|
||||
except Exception:
|
||||
raise HTTPException(status_code=400, detail="Invalid zip file")
|
||||
|
||||
default_app = db.query(App).filter(App.owner_id == user.id, App.name == "openmemory").first()
|
||||
if not default_app:
|
||||
default_app = App(owner_id=user.id, name="openmemory", is_active=True, metadata_={})
|
||||
db.add(default_app)
|
||||
db.commit()
|
||||
db.refresh(default_app)
|
||||
|
||||
cat_id_map: Dict[str, UUID] = {}
|
||||
for c in sqlite_data.get("categories", []):
|
||||
cat = db.query(Category).filter(Category.name == c["name"]).first()
|
||||
if not cat:
|
||||
cat = Category(name=c["name"], description=c.get("description"))
|
||||
db.add(cat)
|
||||
db.commit()
|
||||
db.refresh(cat)
|
||||
cat_id_map[c["id"]] = cat.id
|
||||
|
||||
old_to_new_id: Dict[str, UUID] = {}
|
||||
for m in sqlite_data.get("memories", []):
|
||||
incoming_id = UUID(m["id"])
|
||||
existing = db.query(Memory).filter(Memory.id == incoming_id).first()
|
||||
|
||||
# Cross-user collision: always mint a new UUID and import as a new memory
|
||||
if existing and existing.user_id != user.id:
|
||||
target_id = uuid4()
|
||||
else:
|
||||
target_id = incoming_id
|
||||
|
||||
old_to_new_id[m["id"]] = target_id
|
||||
|
||||
# Same-user collision + skip mode: leave existing row untouched
|
||||
if existing and (existing.user_id == user.id) and mode == "skip":
|
||||
continue
|
||||
|
||||
# Same-user collision + overwrite mode: treat import as ground truth
|
||||
if existing and (existing.user_id == user.id) and mode == "overwrite":
|
||||
incoming_state = m.get("state", "active")
|
||||
existing.user_id = user.id
|
||||
existing.app_id = default_app.id
|
||||
existing.content = m.get("content") or ""
|
||||
existing.metadata_ = m.get("metadata") or {}
|
||||
try:
|
||||
existing.state = MemoryState(incoming_state)
|
||||
except Exception:
|
||||
existing.state = MemoryState.active
|
||||
# Update state-related timestamps from import (ground truth)
|
||||
existing.archived_at = _parse_iso(m.get("archived_at"))
|
||||
existing.deleted_at = _parse_iso(m.get("deleted_at"))
|
||||
existing.created_at = _parse_iso(m.get("created_at")) or existing.created_at
|
||||
existing.updated_at = _parse_iso(m.get("updated_at")) or existing.updated_at
|
||||
db.add(existing)
|
||||
db.commit()
|
||||
continue
|
||||
|
||||
new_mem = Memory(
|
||||
id=target_id,
|
||||
user_id=user.id,
|
||||
app_id=default_app.id,
|
||||
content=m.get("content") or "",
|
||||
metadata_=m.get("metadata") or {},
|
||||
state=MemoryState(m.get("state", "active")) if m.get("state") else MemoryState.active,
|
||||
created_at=_parse_iso(m.get("created_at")) or datetime.now(UTC),
|
||||
updated_at=_parse_iso(m.get("updated_at")) or datetime.now(UTC),
|
||||
archived_at=_parse_iso(m.get("archived_at")),
|
||||
deleted_at=_parse_iso(m.get("deleted_at")),
|
||||
)
|
||||
db.add(new_mem)
|
||||
db.commit()
|
||||
|
||||
for link in sqlite_data.get("memory_categories", []):
|
||||
mid = old_to_new_id.get(link["memory_id"])
|
||||
cid = cat_id_map.get(link["category_id"])
|
||||
if not (mid and cid):
|
||||
continue
|
||||
exists = db.execute(
|
||||
memory_categories.select().where(
|
||||
(memory_categories.c.memory_id == mid) & (memory_categories.c.category_id == cid)
|
||||
)
|
||||
).first()
|
||||
|
||||
if not exists:
|
||||
db.execute(memory_categories.insert().values(memory_id=mid, category_id=cid))
|
||||
db.commit()
|
||||
|
||||
for h in sqlite_data.get("status_history", []):
|
||||
hid = UUID(h["id"])
|
||||
mem_id = old_to_new_id.get(h["memory_id"], UUID(h["memory_id"]))
|
||||
exists = db.query(MemoryStatusHistory).filter(MemoryStatusHistory.id == hid).first()
|
||||
if exists and mode == "skip":
|
||||
continue
|
||||
rec = exists if exists else MemoryStatusHistory(id=hid)
|
||||
rec.memory_id = mem_id
|
||||
rec.changed_by = user.id
|
||||
try:
|
||||
rec.old_state = MemoryState(h.get("old_state", "active"))
|
||||
rec.new_state = MemoryState(h.get("new_state", "active"))
|
||||
except Exception:
|
||||
rec.old_state = MemoryState.active
|
||||
rec.new_state = MemoryState.active
|
||||
rec.changed_at = _parse_iso(h.get("changed_at")) or datetime.now(UTC)
|
||||
db.add(rec)
|
||||
db.commit()
|
||||
|
||||
memory_client = get_memory_client()
|
||||
vector_store = getattr(memory_client, "vector_store", None) if memory_client else None
|
||||
|
||||
if vector_store and memory_client and hasattr(memory_client, "embedding_model"):
|
||||
def iter_logical_records():
|
||||
if memories_blob:
|
||||
gz_buf = io.BytesIO(memories_blob)
|
||||
with gzip.GzipFile(fileobj=gz_buf, mode="rb") as gz:
|
||||
for raw in gz:
|
||||
yield json.loads(raw.decode("utf-8"))
|
||||
else:
|
||||
for m in sqlite_data.get("memories", []):
|
||||
yield {
|
||||
"id": m["id"],
|
||||
"content": m.get("content"),
|
||||
"metadata": m.get("metadata") or {},
|
||||
"created_at": m.get("created_at"),
|
||||
"updated_at": m.get("updated_at"),
|
||||
}
|
||||
|
||||
for rec in iter_logical_records():
|
||||
old_id = rec["id"]
|
||||
new_id = old_to_new_id.get(old_id, UUID(old_id))
|
||||
content = rec.get("content") or ""
|
||||
metadata = rec.get("metadata") or {}
|
||||
created_at = rec.get("created_at")
|
||||
updated_at = rec.get("updated_at")
|
||||
|
||||
if mode == "skip":
|
||||
try:
|
||||
get_fn = getattr(vector_store, "get", None)
|
||||
if callable(get_fn) and vector_store.get(str(new_id)):
|
||||
continue
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
payload = dict(metadata)
|
||||
payload["data"] = content
|
||||
if created_at:
|
||||
payload["created_at"] = created_at
|
||||
if updated_at:
|
||||
payload["updated_at"] = updated_at
|
||||
payload["user_id"] = user_id
|
||||
payload.setdefault("source_app", "openmemory")
|
||||
|
||||
try:
|
||||
vec = memory_client.embedding_model.embed(content, "add")
|
||||
vector_store.insert(vectors=[vec], payloads=[payload], ids=[str(new_id)])
|
||||
except Exception as e:
|
||||
print(f"Vector upsert failed for memory {new_id}: {e}")
|
||||
continue
|
||||
|
||||
return {"message": f'Import completed into user "{user_id}"'}
|
||||
|
||||
return {"message": f'Import completed into user "{user_id}"'}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@ from app.config import DEFAULT_APP_ID, USER_ID
|
||||
from app.database import Base, SessionLocal, engine
|
||||
from app.mcp_server import setup_mcp_server
|
||||
from app.models import App, User
|
||||
from app.routers import apps_router, config_router, memories_router, stats_router
|
||||
from app.routers import apps_router, backup_router, config_router, memories_router, stats_router
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi_pagination import add_pagination
|
||||
@@ -83,6 +83,7 @@ app.include_router(memories_router)
|
||||
app.include_router(apps_router)
|
||||
app.include_router(stats_router)
|
||||
app.include_router(config_router)
|
||||
app.include_router(backup_router)
|
||||
|
||||
# Add pagination support
|
||||
add_pagination(app)
|
||||
|
||||
Reference in New Issue
Block a user