PMFlow/backend/app/services/adaptive_rag_service/mixin_answer_precheck.py

"""Adaptive RAG mixin module."""
from __future__ import annotations

from .shared import *


class AdaptiveRAGAnswerPrecheckMixin:
    async def _maybe_return_clarification_response(
        self,
        *,
        clarification_answer: Optional[str],
        started: float,
        clean_q: str,
        effective_q: str,
        query_type: str,
        asked_by: Optional[int],
        asked_via: str,
        project_id: Optional[int],
        session_id: Optional[int],
        turn_index: int,
        state_row: Optional[QASessionState],
        rewrite_meta: dict[str, Any],
        context_session_id: Optional[int],
        active_thinking_mode: str,
    ) -> Optional[dict[str, Any]]:
        if not clarification_answer:
            return None

        elapsed = int((time.perf_counter() - started) * 1000)
        clarify_sources = [
            {
                "source": "system",
                "label": "Session State",
                "score": 0.95,
                "metadata": {"mode": "clarification", "reason": "missing_subject"},
            }
        ]
        qa = QARecord(
            tenant_id=self.tenant_id,
            question=clean_q,
            answer=clarification_answer,
            sources=clarify_sources,
            intent="question",
            domains_used=["system_context"],
            response_time_ms=elapsed,
            asked_by=asked_by,
            asked_via=(asked_via or "pmflow")[:20],
            project_id=project_id,
            session_id=session_id,
            turn_index=turn_index,
        )
        self.db.add(qa)
        await self._update_session_state_after_turn(
            session_state=state_row,
            question=clean_q,
            effective_question=effective_q,
            answer=clarification_answer,
            query_type=query_type,
            ambiguity="needs_clarification",
            rewrite_meta=rewrite_meta,
            sources=clarify_sources,
        )
        await self.db.commit()
        await self.db.refresh(qa)
        await self.context_cache.append_session_messages(
            tenant_id=self.tenant_id,
            session_id=context_session_id,
            question=clean_q,
            answer=clarification_answer,
        )
        return self._build_response_payload(
            answer=clarification_answer,
            sources=clarify_sources,
            domains_used=["system_context"],
            response_time_ms=elapsed,
            db_query_ms=0,
            vector_search_ms=0,
            db_timeout=False,
            vector_timeout=False,
            qa_id=int(qa.id),
            session_id=int(session_id) if session_id is not None else None,
            turn_index=int(turn_index),
            relevance_max=0.95,
            query_type=query_type,
            confidence=0.92,
            missing_fields=["subject"],
            thinking_mode=active_thinking_mode,
        )

    async def _maybe_return_snapshot_direct_response(
        self,
        *,
        direct_answer: Optional[str],
        started: float,
        clean_q: str,
        effective_q: str,
        asked_by: Optional[int],
        asked_via: str,
        project_id: Optional[int],
        session_id: Optional[int],
        turn_index: int,
        state_row: Optional[QASessionState],
        rewrite_meta: dict[str, Any],
        context_session_id: Optional[int],
        active_thinking_mode: str,
        user_snapshot: Optional[dict[str, Any]],
    ) -> Optional[dict[str, Any]]:
        if not direct_answer:
            return None

        elapsed = int((time.perf_counter() - started) * 1000)
        direct_sources = [
            {
                "source": "system",
                "label": "System Snapshot",
                "score": 0.99,
                "metadata": {"mode": "snapshot_direct", "topic": "identity_or_architecture"},
            }
        ]
        if user_snapshot:
            direct_sources.append(
                {
                    "source": "system",
                    "label": "User Snapshot",
                    "score": 0.98,
                    "metadata": {"mode": "snapshot_direct", "topic": "current_user"},
                }
            )

        qa = QARecord(
            tenant_id=self.tenant_id,
            question=clean_q,
            answer=direct_answer,
            sources=direct_sources,
            intent="question",
            domains_used=["system_context"],
            response_time_ms=elapsed,
            asked_by=asked_by,
            asked_via=(asked_via or "pmflow")[:20],
            project_id=project_id,
            session_id=session_id,
            turn_index=turn_index,
        )
        self.db.add(qa)
        await self._update_session_state_after_turn(
            session_state=state_row,
            question=clean_q,
            effective_question=effective_q,
            answer=direct_answer,
            query_type=QUERY_TYPE_EXPLAIN,
            ambiguity="clear",
            rewrite_meta=rewrite_meta,
            sources=direct_sources,
        )
        await self.db.commit()
        await self.db.refresh(qa)
        await self.context_cache.append_session_messages(
            tenant_id=self.tenant_id,
            session_id=context_session_id,
            question=clean_q,
            answer=direct_answer,
        )
        return self._build_response_payload(
            answer=direct_answer,
            sources=direct_sources,
            domains_used=["system_context"],
            response_time_ms=elapsed,
            db_query_ms=0,
            vector_search_ms=0,
            db_timeout=False,
            vector_timeout=False,
            qa_id=int(qa.id),
            session_id=int(session_id) if session_id is not None else None,
            turn_index=int(turn_index),
            relevance_max=1.0,
            query_type=QUERY_TYPE_EXPLAIN,
            confidence=0.99,
            missing_fields=[],
            thinking_mode=active_thinking_mode,
        )