Repository: emcie-co/parlant
File: src/parlant/core/guidelines.py (Python)

Code before:
    condition: str
    action: str


class _GuidelineDocumentV0_2_0(TypedDict, total=False):
    id: ObjectId
    version: Version.String
    creation_utc: str
    guideline_set: str
    condition: str
    action: str
    enabled: bool


class GuidelineDocumentStore(GuidelineStore):
    VERSION = Version.from_string("0.2.0")

    def __init__(self, database: DocumentDatabase, allow_migration: bool = False) -> None:
        self._database = database
        self._collection: DocumentCollection[_GuidelineDocumentV0_2_0]
        self._allow_migration = allow_migration
        self._lock = ReaderWriterLock()

    async def _document_loader(self, doc: BaseDocument) -> Optional[_GuidelineDocumentV0_2_0]:
        if doc["version"] == "0.1.0":
            doc_v0_1_0 = cast(_GuidelineDocument, doc)

            return _GuidelineDocumentV0_2_0(
                id=doc_v0_1_0["id"],
                version=Version.String("0.2.0"),
                creation_utc=doc_v0_1_0["creation_utc"],
                guideline_set=doc_v0_1_0["guideline_set"],
                condition=doc_v0_1_0["condition"],
                action=doc_v0_1_0["action"],
                enabled=True,
            )

        elif doc["version"] == "0.2.0":
            return cast(_GuidelineDocumentV0_2_0, doc)

        return None

    async def __aenter__(self) -> Self:
        async with MigrationHelper(
            store=self,
            database=self._database,
            allow_migration=self._allow_migration,
        ):
            self._collection = await self._database.get_or_create_collection(
                name="guidelines",
                schema=_GuidelineDocumentV0_2_0,

Review comment (suggestion):
The naming should be different.

`_GuidelineDocument` should always point to the latest version.
It's the old `_GuidelineDocument` that should be renamed to `_GuidelineDocument_V0_1_0`.

Code after:
    version: Version.String
    creation_utc: str
    guideline_set: str
    condition: str
    action: str


class _GuidelineDocument(TypedDict, total=False):
    id: ObjectId
    version: Version.String
    creation_utc: str
    guideline_set: str
    condition: str
    action: str
    enabled: bool


class GuidelineDocumentStore(GuidelineStore):
    VERSION = Version.from_string("0.2.0")

    def __init__(self, database: DocumentDatabase, allow_migration: bool = False) -> None:
        self._database = database
        self._collection: DocumentCollection[_GuidelineDocument]
        self._allow_migration = allow_migration
        self._lock = ReaderWriterLock()

    async def _document_loader(self, doc: BaseDocument) -> Optional[_GuidelineDocument]:
        async def v0_1_0_to_v_0_2_0(doc: BaseDocument) -> Optional[BaseDocument]:
            d = cast(_GuidelineDocument_V0_1_0, doc)
            return _GuidelineDocument(
                id=d["id"],
                version=Version.String("0.2.0"),
                creation_utc=d["creation_utc"],
                guideline_set=d["guideline_set"],
                condition=d["condition"],
                action=d["action"],
                enabled=True,
            )

        return await DocumentMigrationHelper[_GuidelineDocument](
            self, {"0.1.0": v0_1_0_to_v_0_2_0}
        ).migrate(doc)

    async def __aenter__(self) -> Self:
        async with DocumentStoreMigrationHelper(
            store=self,
            database=self._database,
            allow_migration=self._allow_migration,
        ):
            self._collection = await self._database.get_or_create_collection(
                name="guidelines",