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

Code before:
                return data

    async def _save_data(
        self,
        data: dict[str, list[dict[str, Any]]],
    ) -> None:
        async with self._lock:
            async with aiofiles.open(self.file_path, mode="w") as file:
                json_string = json.dumps(
                    data,
                    ensure_ascii=False,
                    indent=4,
                )
                await file.write(json_string)

    async def insert_one(
        self,
        collection: CollectionDescriptor,
        document: dict[str, Any],
    ) -> dict[str, Any]:
        result = await self.transient_db.insert_one(
            collection,
            collection.schema(**document).model_dump(mode="json"),
        )
        await self._process_operation_counter()
        return collection.schema.model_validate(result).model_dump()

    async def find(
        self,
        collection: CollectionDescriptor,
        filters: dict[str, FieldFilter],
    ) -> Sequence[dict[str, Any]]:
        return [
            collection.schema.model_validate(doc).model_dump()
            for doc in await self.transient_db.find(collection, filters)
        ]

    async def find_one(
        self,
        collection: CollectionDescriptor,
        filters: dict[str, FieldFilter],
    ) -> dict[str, Any]:
        result = await self.transient_db.find_one(collection, filters)
        return collection.schema.model_validate(result).model_dump()

    async def update_one(
        self,
        collection: CollectionDescriptor,
        filters: dict[str, FieldFilter],
        updated_document: dict[str, Any],
        upsert: bool = False,

Review comment (bug):
I think the contract here is a little bit weird.

The internal Pydantic model does not necessarily have to match the input document structure perfectly.
This code seems to assume that it does, while in reality I think the caller for `insert_one()` will expect to have the same structure that it sent returned to it.

I think that in general the contract here, of returning an ID-augmented document (of an unspecified structure) is wrong for `insert_one()`. Instead, PyMongo's API is something that works better IMO, where you basically just return the ID of the newly created document—as a string in this case.

So change the signature to `insert_one(self, collection: CollectionDescriptor, document: dict[str, Any]) -> ObjectID`

And at the top add a `NewType(...)` for `ObjectID` as we do for things like `SessionID`, `EventID`, etc.

Code after:
                    indent=4,
                )
                await file.write(json_string)

    async def find(
        self,
        collection: CollectionDescriptor,
        filters: dict[str, FieldFilter],
    ) -> Sequence[dict[str, Any]]:
        # TODO MC-101
        self._collection_descriptors[collection.name] = collection
        return [
            collection.schema.model_validate(doc).model_dump()
            for doc in await self.transient_db.find(collection, filters)
        ]

    async def find_one(
        self,
        collection: CollectionDescriptor,
        filters: dict[str, FieldFilter],
    ) -> dict[str, Any]:
        # TODO MC-101
        self._collection_descriptors[collection.name] = collection
        result = await self.transient_db.find_one(collection, filters)
        return collection.schema.model_validate(result).model_dump()

    async def insert_one(
        self,
        collection: CollectionDescriptor,
        document: dict[str, Any],
    ) -> ObjectId:
        # TODO MC-101
        self._collection_descriptors[collection.name] = collection
        document_id = await self.transient_db.insert_one(
            collection,
            collection.schema(**document).model_dump(mode="json"),
        )
        await self._process_operation_counter()
        return document_id

    async def update_one(
        self,
        collection: CollectionDescriptor,
        filters: dict[str, FieldFilter],
        updated_document: dict[str, Any],
        upsert: bool = False,
    ) -> ObjectId:
        # TODO MC-101
        self._collection_descriptors[collection.name] = collection
        result = await self.transient_db.update_one(
            collection,