clarification_agent.py

from __future__ import annotations

import json
import logging
from typing import Awaitable, Callable, Optional

from backend.agents.base import BaseAgent
from backend.models.enums import AgentPhase
from backend.models.schemas import (
    AmbiguityReport,
    ClarificationQuestion,
    ClarificationRequest,
    ClarifiedIntent,
)
from backend.prompts.clarification_prompt import (
    AUTO_RESOLVE_PROMPT,
    CLARIFICATION_SYSTEM_PROMPT,
    CLARIFICATION_WITH_CONTEXT_PROMPT,
)
from backend.services.exa_service import ExaService, SearchResult
from backend.services.llm_service import LLMService

logger = logging.getLogger(__name__)

AUTO_RESOLVE_CONFIDENCE_THRESHOLD = 0.7

# System prompt used to derive Exa search queries from the intent
_QUERY_DERIVATION_PROMPT = """\
You are a search query specialist. Given a user's intent and a list of ambiguity \
dimensions, produce 2-3 precise web search queries that would surface current \
best-practice guidance, real-world examples, or domain standards relevant to the task.

Output ONLY a JSON object: {"queries": ["query1", "query2", "query3"]}
Queries should be specific, domain-aware, and suitable for a neural web search engine.
Do NOT include generic queries.
"""


class ClarificationAgent(BaseAgent):
    phase = AgentPhase.CLARIFICATION
    name = "Clarification Agent"

    def __init__(self, llm_service: LLMService, exa_service: ExaService) -> None:
        super().__init__(llm_service)
        self.exa = exa_service

    def get_system_prompt(self) -> str:
        return CLARIFICATION_SYSTEM_PROMPT

    # ------------------------------------------------------------------
    # Public methods
    # ------------------------------------------------------------------

    async def generate_questions(
        self,
        input_data: AmbiguityReport,
        on_stream: Optional[Callable[[str], Awaitable[None]]] = None,
    ) -> tuple[ClarificationRequest | None, list[SearchResult]]:
        """Generate clarification questions enriched with Exa search context.

        Returns a (ClarificationRequest | None, search_results) tuple so the
        caller can pass the same search_results into auto_resolve_questions()
        without performing a second round of searches.
        """
        if not input_data.needs_clarification:
            return None, []

        search_results: list[SearchResult] = []
        if self.exa.available:
            try:
                queries = await self._derive_search_queries(input_data)
                if queries:
                    search_results = await self.exa.search_many(queries)
                    logger.info(
                        "Exa returned %d results for queries: %s",
                        len(search_results),
                        queries,
                    )
            except Exception:
                logger.exception("Exa enrichment failed — falling back to LLM-only")

        if search_results:
            req = await self._generate_with_context(input_data, search_results)
        else:
            req = await self._generate_without_context(input_data)

        return req, search_results

    async def auto_resolve_questions(
        self,
        request: ClarificationRequest,
        search_results: list[SearchResult],
        ambiguity_report: AmbiguityReport,
    ) -> ClarificationRequest:
        """Try to auto-answer each question from search results via LLM reasoning.

        Questions that can be confidently answered (confidence >= threshold) are
        marked auto_resolved=True with their answer and source stored.
        Questions that need personal/org-specific context remain unresolved.
        """
        if not search_results:
            return request

        context_block = self._format_context_block(search_results)

        for q in request.questions:
            try:
                resolved = await self._attempt_resolve(q, context_block, ambiguity_report)
                if resolved:
                    answer, source = resolved
                    q.auto_resolved = True
                    q.auto_answer = answer
                    q.auto_answer_source = source
                    logger.info(
                        "Auto-resolved question [%s]: %s => %r",
                        q.dimension, q.question, answer
                    )
                else:
                    logger.info(
                        "Could not auto-resolve question [%s]: %s",
                        q.dimension, q.question
                    )
            except Exception:
                logger.exception("Auto-resolve failed for question: %s", q.question)

        request.auto_resolved_count = sum(1 for q in request.questions if q.auto_resolved)
        return request

    @staticmethod
    def get_unresolved(request: ClarificationRequest) -> list[ClarificationQuestion]:
        """Return only the questions that could not be auto-resolved."""
        return [q for q in request.questions if not q.auto_resolved]

    @staticmethod
    def build_auto_answers(request: ClarificationRequest) -> dict[str, str]:
        """Build a dimension→answer dict from all auto-resolved questions."""
        return {
            q.dimension: q.auto_answer
            for q in request.questions
            if q.auto_resolved and q.dimension
        }

    async def process_answers(
        self,
        ambiguity_report: AmbiguityReport,
        answers: dict[str, str],
        on_stream: Optional[Callable[[str], Awaitable[None]]] = None,
    ) -> ClarifiedIntent:
        """Take user answers (possibly merged with auto-answers) and refine intent."""
        user_msg = (
            f"Original intent:\n{ambiguity_report.intent.model_dump_json(indent=2)}\n\n"
            f"Ambiguity flags:\n{ambiguity_report.model_dump_json(indent=2)}\n\n"
            f"Clarifications:\n{json.dumps(answers, indent=2)}\n\n"
            "Produce a refined, unambiguous intent.\n"
            "Respond ONLY with JSON containing: refined_goal, refined_constraints, "
            "refined_output_format."
        )
        result = await self.llm.generate_json(self.get_system_prompt(), user_msg)
        result["original_intent"] = ambiguity_report.intent.model_dump()
        result["clarifications"] = answers
        return ClarifiedIntent(**result)

    async def process(self, input_data, on_stream=None):
        raise NotImplementedError(
            "Use generate_questions() and process_answers() directly"
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    async def _attempt_resolve(
        self,
        question: ClarificationQuestion,
        context_block: str,
        report: AmbiguityReport,
    ) -> tuple[str, str] | None:
        """Ask LLM whether this question can be answered from search results.

        Returns (answer, source) if resolvable with sufficient confidence,
        or None if the question needs the user.
        """
        user_msg = (
            f"User's goal: {report.intent.goal}\n"
            f"User's context: {report.intent.context}\n\n"
            f"Clarification question: {question.question}\n"
            f"Ambiguity dimension: {question.dimension}\n"
            f"Why it matters: {question.why_needed}\n\n"
            f"{context_block}\n\n"
            "Can this question be confidently answered from the web search results alone?\n"
            "Remember: only answer from the search results — do NOT use your own knowledge.\n"
            'Respond ONLY with JSON: {"resolvable": bool, "answer": "...", '
            '"confidence": 0.0-1.0, "source": "..."}'
        )
        result = await self.llm.generate_json(AUTO_RESOLVE_PROMPT, user_msg)

        if not result.get("resolvable", False):
            return None
        confidence = float(result.get("confidence", 0.0))
        if confidence < AUTO_RESOLVE_CONFIDENCE_THRESHOLD:
            return None
        answer = str(result.get("answer", "")).strip()
        source = str(result.get("source", "web search")).strip()
        if not answer:
            return None
        return answer, source

    async def _derive_search_queries(self, report: AmbiguityReport) -> list[str]:
        """Ask the LLM to produce targeted Exa search queries from the intent."""
        dimensions = ", ".join(f.dimension for f in report.flags) or "general"
        user_msg = (
            f"User's goal: {report.intent.goal}\n"
            f"Context: {report.intent.context}\n"
            f"Ambiguous dimensions: {dimensions}\n\n"
            "Produce 2-3 specific web search queries to find current best practices, "
            "real-world examples, or domain standards that would help clarify these dimensions.\n\n"
            'Respond ONLY with JSON: {"queries": ["...", "...", "..."]}'
        )
        try:
            result = await self.llm.generate_json(_QUERY_DERIVATION_PROMPT, user_msg)
            queries = result.get("queries", [])
            return [q for q in queries if isinstance(q, str) and q.strip()][:3]
        except Exception:
            logger.exception("Failed to derive Exa search queries")
            return []

    async def _generate_with_context(
        self,
        report: AmbiguityReport,
        search_results: list[SearchResult],
    ) -> ClarificationRequest:
        """Generate questions using Exa search context injected into the prompt."""
        context_block = self._format_context_block(search_results)
        user_msg = (
            "Based on the ambiguity report below AND the real-world web search context "
            "provided, generate the minimum set of essential, high-impact clarification "
            "questions. Use the search results to ground options in current best practices "
            "and real-world examples.\n\n"
            f"Ambiguity report:\n{report.model_dump_json(indent=2)}\n\n"
            f"{context_block}\n\n"
            "Principles: minimal, specific, decision-oriented. Max 3 questions.\n"
            "Use actual examples from the search results in the 'options' field where relevant.\n\n"
            "Respond ONLY with JSON matching:\n"
            f"{ClarificationRequest.model_json_schema()}"
        )
        result = await self.llm.generate_json(CLARIFICATION_WITH_CONTEXT_PROMPT, user_msg)
        req = ClarificationRequest(**result)
        req.search_context = [dict(r) for r in search_results]
        return req

    async def _generate_without_context(
        self, report: AmbiguityReport
    ) -> ClarificationRequest:
        """Fallback: generate questions using LLM only."""
        user_msg = (
            "Based on the ambiguity report below, generate the minimum set of "
            "essential, high-impact clarification questions.\n\n"
            f"Ambiguity report:\n{report.model_dump_json(indent=2)}\n\n"
            "Principles: minimal, specific, decision-oriented. Max 3 questions.\n\n"
            "Respond ONLY with JSON matching:\n"
            f"{ClarificationRequest.model_json_schema()}"
        )
        result = await self.llm.generate_json(self.get_system_prompt(), user_msg)
        return ClarificationRequest(**result)

    @staticmethod
    def _format_context_block(results: list[SearchResult]) -> str:
        """Format search results into a readable context block for the LLM."""
        if not results:
            return ""
        lines = ["## Web Search Context"]
        for i, r in enumerate(results, 1):
            lines.append(f"{i}. **{r['title']}** ({r['url']})")
            if r["snippet"]:
                lines.append(f"   > {r['snippet']}")
        return "\n".join(lines)