Repository: Cinnamon/kotaemon
File: libs/ktem/ktem/tags/crud.py (Python)

Code before:
        return False


class ChunkTagIndexCRUD:
    def __init__(self, engine):
        self._engine = engine

    @cached_property
    def tag_crud(self) -> TagCRUD:
        return TagCRUD(self._engine)

    def list_all(self) -> list[ChunkTagIndex]:
        with Session(self._engine) as session:
            statement = select(ChunkTagIndex)
            results = session.exec(statement).all()

        return results

    def query_by_id(self, record_id: str) -> ChunkTagIndex | None:
        with Session(self._engine) as session:
            statement = select(ChunkTagIndex).where(ChunkTagIndex.id == record_id)

            result = session.exec(statement).first()
            return result

    def query_by_chunk_ids(self, chunk_ids: list[str]) -> dict[str, dict[str, str]]:
        chunk_id_to_tags = defaultdict(dict)
        tag_ids = set()

        with Session(self._engine) as session:
            statement = select(ChunkTagIndex).where(
                ChunkTagIndex.chunk_id.in_(chunk_ids)
            )

            results = session.exec(statement).all()
            for result in results:
                chunk_id = result.chunk_id
                tag_id = result.tag_id
                tag_ids.add(tag_id)
                chunk_id_to_tags[chunk_id][tag_id] = {
                    "content": result.content,
                    "status": result.status,
                }

        tags: list[Tag] = self.tag_crud.query_by_ids(list(tag_ids))
        tag_id_to_tag: dict[str, str] = {tag.id: tag.name for tag in tags}
        for chunk_id, tags_dict in chunk_id_to_tags.items():
            for tag_id, tag_dict in tags_dict.items():
                tag_dict["name"] = tag_id_to_tag[tag_id]

        return chunk_id_to_tags

Review comment (suggestion):
since the return value is `chunk_id_to_tag`, the expected type should be `dict[str, Tag]` right? no motivation to output `dict[str, str]`

Code after:
        with Session(self._engine) as session:
            session.execute(stmt)
            session.commit()

        return False


class ChunkTagIndexCRUD:
    def __init__(self, engine):
        self._engine = engine

    @cached_property
    def tag_crud(self) -> TagCRUD:
        return TagCRUD(self._engine)

    def list_all(self) -> list[ChunkTagIndex]:
        with Session(self._engine) as session:
            statement = select(ChunkTagIndex)
            results = session.exec(statement).all()

        return results

    def query_by_id(self, record_id: str) -> ChunkTagIndex | None:
        with Session(self._engine) as session:
            statement = select(ChunkTagIndex).where(ChunkTagIndex.id == record_id)

            result = session.exec(statement).first()
            return result

    def query_by_chunk_ids(self, chunk_ids: list[str]) -> dict[str, dict[str, Any]]:
        chunk_id_to_tags = defaultdict(dict)
        tag_ids = set()

        with Session(self._engine) as session:
            statement = select(ChunkTagIndex).where(
                ChunkTagIndex.chunk_id.in_(chunk_ids)
            )

            results = session.exec(statement).all()
            for result in results:
                chunk_id = result.chunk_id
                tag_id = result.tag_id
                tag_ids.add(tag_id)
                chunk_id_to_tags[chunk_id][tag_id] = {
                    "content": result.content,
                    "status": result.status,
                }

        tags: list[Tag] = self.tag_crud.query_by_ids(list(tag_ids))
        tag_id_to_tag: dict[str, str] = {tag.id: tag.name for tag in tags}
        for chunk_id, tags_dict in chunk_id_to_tags.items():