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

Code before:
from datetime import datetime
from functools import cached_property

from ktem.db.base_models import TagProcessStatus, TagScope, TagType
from ktem.db.models import ChunkTagIndex, Tag
from sqlmodel import Session, select


class TagCRUD:
    def __init__(self, engine):
        self._engine = engine

    def list_all(self) -> list[Tag]:
        with Session(self._engine) as session:
            statement = select(Tag)
            results = session.exec(statement).all()

            return results

    def create(
        self,
        name: str,
        prompt: str,
        config: str,
        type: str = TagType.text.value,
        scope: str = TagScope.chunk.value,
        valid_classes: str | None = None,
    ) -> str | None:
        meta = {}
        if type == TagType.classification:
            assert (
                valid_classes is not None and valid_classes != ""
            ), "Invalid valid classes!"
            meta["valid_classes"] = valid_classes

        with Session(self._engine) as session:
            # tag name must be unique
            existed_result = self.query_by_name(name)
            if existed_result:
                raise Exception(f"Tag name: {name} has already been existed!")

            tag: Tag = Tag(
                name=name,
                prompt=prompt,
                config=config,
                type=type,
                scope=scope,
                meta=meta,
            )
            session.add(tag)
            session.commit()

Review comment (suggestion):
ideally it should be `type: TagType = TagType.text`

Code after:
from functools import cached_property
from typing import Any

from ktem.db.base_models import TagProcessStatus, TagScope, TagType
from ktem.db.models import ChunkTagIndex, Tag
from sqlmodel import Session, select, update


class TagCRUD:
    def __init__(self, engine):
        self._engine = engine

    def list_all(self) -> list[Tag]:
        with Session(self._engine) as session:
            statement = select(Tag)
            results = session.exec(statement).all()

            return results

    def create(
        self,
        name: str,
        prompt: str,
        config: str,
        type: str | TagType = TagType.text,
        scope: str | TagType = TagScope.chunk,
        valid_classes: str | None = None,
    ) -> str | None:
        meta = {}
        if type == TagType.classification:
            assert (
                valid_classes is not None and valid_classes != ""
            ), "Invalid valid classes!"
            meta["valid_classes"] = valid_classes

        with Session(self._engine) as session:
            # tag name must be unique
            existed_result = self.query_by_name(name)
            if existed_result:
                raise Exception(f"Tag name: {name} has already been existed!")

            tag: Tag = Tag(
                name=name,
                prompt=prompt,
                config=config,
                type=type.value if isinstance(type, TagType) else type,
                scope=scope.value if isinstance(type, TagScope) else scope,
                meta=meta,
            )
            session.add(tag)
            session.commit()