SupplyAgent/supply_infra/db/repositories/agent_document_injection_repo.py

from __future__ import annotations

from sqlalchemy import select

from supply_infra.db.models.agent_document_injection import AgentDocumentInjection
from supply_infra.db.repositories.base import BaseRepository


class AgentDocumentInjectionRepository(BaseRepository[AgentDocumentInjection]):
    model = AgentDocumentInjection

    def get_by_agent_name(self, agent_name: str) -> AgentDocumentInjection | None:
        stmt = select(AgentDocumentInjection).where(
            AgentDocumentInjection.agent_name == agent_name
        )
        return self.session.scalar(stmt)

    def upsert(
        self,
        *,
        agent_name: str,
        content: str,
        enabled: bool,
    ) -> AgentDocumentInjection:
        entity = self.get_by_agent_name(agent_name)
        if entity is None:
            entity = AgentDocumentInjection(
                agent_name=agent_name,
                content=content,
                enabled=int(enabled),
            )
            return self.add(entity)

        entity.content = content
        entity.enabled = int(enabled)
        self.session.flush()
        return entity