Llama Stack 0.1.7 to 0.2.9 MongoDB persistent data migration

migrate.py

"""
Adds `created_at` field to all Agents in persisten store.
Adds empty list `turns` to every Session under each Agent.

Run with:
  $ export MONGODB_URI=mongodb://.....
  $ export MONGODB_DB_NAME=...
  $ export MONGODB_USERNAME=...
  $ export MONGODB_PASSWORD=...
  $ python migrate.py 
"""
import asyncio
import json
import os

from pymongo import AsyncMongoClient

CREATED_AT_KEY = "created_at"
TURNS_KEY = "turns"
CREATED_AT = "2025-07-01T00:00:00.000000+00:00"


def get_mongo_collection(
        connection_uri: str,
        database_name: str,
        username: str,
        password: str,
        collection_name: str
):
    client = AsyncMongoClient(
        connection_uri,
        username=username,
        password=password,
        uuidRepresentation="standard",
        tz_aware=True,
    )
    db = client[database_name]
    return db[collection_name]


async def main():
    print("Migrating Llama Stack MongoDB from 0.1.7 to 0.2.9")
    print("Sleeping 5s...")
    await asyncio.sleep(5)
    llamastack_kvstore = get_mongo_collection(
        os.environ.get('MONGODB_URI'),
        os.environ.get('MONGODB_DB_NAME'),
        os.environ.get('MONGODB_USERNAME'),
        os.environ.get('MONGODB_PASSWORD'),
        "llamastack_kvstore"
    )

    # Get a list of all Agent IDs
    ls_agents = llamastack_kvstore.find(
        {"key": {"$regex": "agent:.*"}}
    )

    async for agent in ls_agents:
        # Check if Agent has a created_at property
        agent_properties = json.loads(agent["value"])
        agent_key = agent["key"]
        print(f"Checking Agent for created_at: {agent_key}")
        if CREATED_AT_KEY not in agent_properties:
            # Add created_at property
            print(f"Adding created_at field to {agent_key}")
            agent_properties[CREATED_AT_KEY] = CREATED_AT
            result = await llamastack_kvstore.update_one(
                {"key": agent['key']},
                {"$set": {"value": json.dumps(agent_properties, separators=(",", ":")), }}
            )
            if result.modified_count > 0:
                print(f"Modified {agent_key}")
            else:
                print(f"Did not modify {agent_key}")

        # Get a list of all Sessions under this Agent
        print(f"Checking Session turns field for Agent {agent_key}")
        agent_id = agent_key.split(":")[1]
        ls_agent_sessions = llamastack_kvstore.find(
            {"key": {"$regex": f"session:{agent_id}:[^:]+$"}}
        )
        async for session in ls_agent_sessions:
            session_key = session["key"]
            session_properties = json.loads(session["value"])
            print(f"Checking Session {session_key} for {agent_key}")
            if TURNS_KEY not in session_properties:
                # Add empty turns list to session
                session_properties[TURNS_KEY] = []
                result = await llamastack_kvstore.update_one(
                    {"key": session_key},
                    {"$set": {"value": json.dumps(session_properties, separators=(",", ":"))}}
                )
                if result.modified_count > 0:
                    print(f"Modified {session_key}")
                else:
                    print(f"Did not modify {session_key}")


if __name__ == "__main__":
    asyncio.run(main())