backend/utils/session_db.py

from utils.pipeline import RAGPipeline


class SessionDB:
    data: dict[str, RAGPipeline] = {}

    def add(self, key: str, value: RAGPipeline) -> None:
        if key in self.data:
            raise Exception("Already exists")
        self.data[key] = value

    def get(self, key: str) -> RAGPipeline:
        try:
            return self.data[key]
        except KeyError:
            raise Exception(f"Not found for {key}")

    def remove(self, key: str) -> None:
        self.data.pop(key)