JARVIS/backend/app/agents/learning/store.py

from __future__ import annotations

from sqlalchemy import desc, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.agents.schemas.learning import SessionRetrospective
from app.models.learning import LearningArtifactRecord, SessionRetrospectiveRecord


class SessionRetrospectiveStore:
    def __init__(self, db: AsyncSession):
        self.db = db

    async def save(self, retrospective: SessionRetrospective) -> SessionRetrospectiveRecord:
        payload = retrospective.model_dump(mode="json")
        record = SessionRetrospectiveRecord(
            user_id=retrospective.user_id,
            conversation_id=retrospective.conversation_id,
            request_message_id=retrospective.request_message_id,
            response_message_id=retrospective.response_message_id,
            query_text=retrospective.query_text,
            final_response=retrospective.final_response,
            summary_text=retrospective.summary,
            task_type=retrospective.task_type,
            execution_mode=retrospective.execution_mode,
            primary_agent=retrospective.primary_agent,
            verification_status=retrospective.verification_status,
            verification_summary=retrospective.verification_summary,
            skill_names=retrospective.used_skill_names,
            evidence=retrospective.evidence_refs,
            task_refs=retrospective.task_refs,
            payload=payload,
        )
        self.db.add(record)
        await self.db.commit()
        await self.db.refresh(record)
        return record

    async def list_recent(
        self,
        *,
        user_id: str,
        limit: int = 20,
    ) -> list[SessionRetrospectiveRecord]:
        result = await self.db.execute(
            select(SessionRetrospectiveRecord)
            .where(SessionRetrospectiveRecord.user_id == user_id)
            .order_by(desc(SessionRetrospectiveRecord.recorded_at), desc(SessionRetrospectiveRecord.created_at))
            .limit(limit)
        )
        return list(result.scalars().all())


class LearningArtifactStore:
    def __init__(self, db: AsyncSession):
        self.db = db

    async def save_batch(
        self,
        *,
        user_id: str,
        conversation_id: str,
        retrospective_id: str | None,
        artifacts: list[dict[str, object]],
    ) -> list[LearningArtifactRecord]:
        records: list[LearningArtifactRecord] = []
        for artifact in artifacts:
            record = LearningArtifactRecord(
                user_id=user_id,
                conversation_id=conversation_id,
                retrospective_id=retrospective_id,
                artifact_type=str(artifact.get("artifact_type") or "unknown"),
                artifact_key=str(artifact.get("artifact_key") or "") or None,
                summary_text=str(artifact.get("summary_text") or ""),
                payload=dict(artifact.get("payload") or {}),
            )
            self.db.add(record)
            records.append(record)

        await self.db.commit()
        for record in records:
            await self.db.refresh(record)
        return records

    async def list_recent(
        self,
        *,
        user_id: str,
        artifact_type: str | None = None,
        limit: int = 50,
    ) -> list[LearningArtifactRecord]:
        query = select(LearningArtifactRecord).where(LearningArtifactRecord.user_id == user_id)
        if artifact_type:
            query = query.where(LearningArtifactRecord.artifact_type == artifact_type)
        result = await self.db.execute(
            query.order_by(
                desc(LearningArtifactRecord.recorded_at),
                desc(LearningArtifactRecord.created_at),
            ).limit(limit)
        )
        return list(result.scalars().all())

    async def aggregate_counts_by_key(
        self,
        *,
        user_id: str,
        artifact_type: str,
        limit: int = 100,
    ) -> dict[str, int]:
        records = await self.list_recent(user_id=user_id, artifact_type=artifact_type, limit=limit)
        counts: dict[str, int] = {}
        for record in records:
            key = record.artifact_key or "unknown"
            counts[key] = counts.get(key, 0) + 1
        return counts


def append_retrospective_attachment(
    attachments: list[dict] | None,
    retrospective: SessionRetrospective,
) -> list[dict]:
    next_attachments = list(attachments or [])
    next_attachments.append(
        {
            "kind": "session_retrospective",
            "payload": retrospective.model_dump(mode="json"),
        }
    )
    return next_attachments