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

from __future__ import annotations

import asyncio
import logging
from typing import Any

from app.config import settings
from app.database import async_session
from app.agents.learning.bridge import update_learning_decision_with_bridge
from app.agents.learning.pattern_miner import LearningPatternMiner
from app.agents.learning.audit import build_learning_audit_entry
from app.agents.learning.retrospector import build_session_retrospective
from app.agents.learning.signal_extractor import RetrospectiveSignalExtractor
from app.agents.learning.skill_candidate_builder import SkillCandidateBuilder
from app.agents.learning.store import LearningArtifactStore, SessionRetrospectiveStore
from app.agents.schemas.learning import LearningDecision, SessionRetrospective
from app.agents.skills.evaluator import SkillPromotionEvaluator

logger = logging.getLogger(__name__)


def _enrich_retrospective(retrospective: SessionRetrospective) -> SessionRetrospective:
    signals = RetrospectiveSignalExtractor().extract(retrospective)
    patterns = LearningPatternMiner().mine(signals)
    skill_candidates = SkillCandidateBuilder().build(patterns)

    decision = LearningDecision(
        decision="create_candidate" if skill_candidates else ("reinforce_memory" if signals else "defer"),
        explanation=(
            "Retrospective produced reusable candidate skills."
            if skill_candidates
            else "Retrospective only reinforces memory-like evidence."
            if signals
            else "No stable signal was extracted from this retrospective."
        ),
        evidence_refs=(skill_candidates[0].evidence_refs if skill_candidates else retrospective.evidence_refs[:3]),
        metadata={
            "signal_count": len(signals),
            "pattern_count": len(patterns),
            "skill_candidate_count": len(skill_candidates),
        },
    )

    retrospective.learning_signals = signals
    retrospective.pattern_candidates = patterns
    retrospective.skill_candidates = skill_candidates
    retrospective.learning_decision = update_learning_decision_with_bridge(decision, signals)
    return retrospective


def _build_learning_artifacts(retrospective: SessionRetrospective) -> list[dict[str, object]]:
    artifacts: list[dict[str, object]] = []
    for signal in retrospective.learning_signals:
        artifacts.append(
            {
                "artifact_type": "signal",
                "artifact_key": signal.signal_type,
                "summary_text": signal.explanation or signal.signal_type,
                "payload": signal.model_dump(mode="json"),
            }
        )
    for pattern in retrospective.pattern_candidates:
        artifacts.append(
            {
                "artifact_type": "pattern_candidate",
                "artifact_key": pattern.pattern_type,
                "summary_text": pattern.description,
                "payload": pattern.model_dump(mode="json"),
            }
        )
    for candidate in retrospective.skill_candidates:
        artifacts.append(
            {
                "artifact_type": "skill_candidate",
                "artifact_key": candidate.name,
                "summary_text": candidate.summary,
                "payload": candidate.model_dump(mode="json"),
            }
        )
    if retrospective.learning_decision is not None:
        artifacts.append(
            {
                "artifact_type": "learning_decision",
                "artifact_key": retrospective.learning_decision.decision,
                "summary_text": retrospective.learning_decision.explanation,
                "payload": retrospective.learning_decision.model_dump(mode="json"),
            }
        )
        artifacts.append(
            {
                "artifact_type": "learning_audit",
                "artifact_key": retrospective.retrospective_id or "retrospective",
                "summary_text": retrospective.learning_decision.explanation,
                "payload": build_learning_audit_entry(retrospective),
            }
        )
    return artifacts


def _build_lifecycle_artifacts(decisions: list) -> list[dict[str, object]]:
    artifacts: list[dict[str, object]] = []
    for decision in decisions:
        artifacts.append(
            {
                "artifact_type": "skill_lifecycle_decision",
                "artifact_key": getattr(decision, "skill_name", None) or "skill",
                "summary_text": getattr(decision, "reason", ""),
                "payload": decision.model_dump(mode="json"),
            }
        )
    return artifacts


async def persist_retrospective(
    *,
    user_id: str,
    conversation_id: str,
    request_message_id: str | None,
    response_message_id: str | None,
    query_text: str,
    final_response: str | None,
    state: dict[str, Any] | None,
) -> None:
    retrospective = build_session_retrospective(
        request_id=response_message_id or request_message_id or conversation_id,
        session_id=conversation_id,
        user_query=query_text,
        state=state,
        runtime_context={"user_id": user_id},
    )
    retrospective = _enrich_retrospective(retrospective)

    async with async_session() as session:
        saved = await SessionRetrospectiveStore(session).save(retrospective)
        lifecycle_decisions = []
        if settings.ENABLE_SKILL_PROMOTION:
            lifecycle_decisions = await SkillPromotionEvaluator(session).sync_retrospective(
                user_id=user_id,
                retrospective=retrospective,
            )
        if settings.ENABLE_LEARNING_SIGNALS:
            await LearningArtifactStore(session).save_batch(
                user_id=user_id,
                conversation_id=conversation_id,
                retrospective_id=saved.id,
                artifacts=[
                    *_build_learning_artifacts(retrospective),
                    *_build_lifecycle_artifacts(lifecycle_decisions),
                ],
            )


def schedule_retrospective_job(**kwargs) -> asyncio.Task[None] | None:
    if not settings.ENABLE_RETROSPECTIVE:
        return None
    try:
        task = asyncio.create_task(persist_retrospective(**kwargs))
    except RuntimeError:
        return None

    def _handle_completion(done_task: asyncio.Task[None]) -> None:
        try:
            done_task.result()
        except Exception:
            logger.exception("retrospective_job_failed")

    task.add_done_callback(_handle_completion)
    return task


def schedule_retrospective_learning_event(
    *,
    user_id: str,
    conversation_id: str,
    retrospective: SessionRetrospective,
    session_factory=async_session,
) -> asyncio.Task[None] | None:
    if not settings.ENABLE_RETROSPECTIVE:
        return None

    async def _persist_existing() -> None:
        async with session_factory() as session:
            enriched = _enrich_retrospective(retrospective)
            saved = await SessionRetrospectiveStore(session).save(enriched)
            lifecycle_decisions = []
            if settings.ENABLE_SKILL_PROMOTION:
                lifecycle_decisions = await SkillPromotionEvaluator(session).sync_retrospective(
                    user_id=user_id,
                    retrospective=enriched,
                )
            if settings.ENABLE_LEARNING_SIGNALS:
                await LearningArtifactStore(session).save_batch(
                    user_id=user_id,
                    conversation_id=conversation_id,
                    retrospective_id=saved.id,
                    artifacts=[
                        *_build_learning_artifacts(enriched),
                        *_build_lifecycle_artifacts(lifecycle_decisions),
                    ],
                )

    try:
        task = asyncio.create_task(_persist_existing())
    except RuntimeError:
        return None

    def _handle_completion(done_task: asyncio.Task[None]) -> None:
        try:
            done_task.result()
        except Exception:
            logger.exception(
                "retrospective_learning_event_failed",
                extra={
                    "details": {
                        "user_id": user_id,
                        "conversation_id": conversation_id,
                    }
                },
            )

    task.add_done_callback(_handle_completion)
    return task