src/feedback/refinement.py

from __future__ import annotations
"""Refinement strategies for improving responses."""

from dataclasses import dataclass
from typing import Any

from src.feedback.gaps import InformationGap


@dataclass
class RefinementAction:
    """An action to refine the response."""

    action_type: str  # "search", "verify", "expand", "clarify", "restructure"
    description: str
    parameters: dict[str, Any]
    priority: int  # 1-5, lower is higher priority


class RefinementStrategy:
    """Determines strategies for refining responses."""

    def __init__(self, max_iterations: int = 3):
        """Initialize the refinement strategy.
        
        Args:
            max_iterations: Maximum refinement iterations allowed
        """
        self.max_iterations = max_iterations
        self._current_iteration = 0

    def analyze(
        self,
        query: str,
        answer: str,
        gaps: list[InformationGap],
        quality_score: float,
    ) -> list[RefinementAction]:
        """Analyze response and determine refinement actions.
        
        Args:
            query: Original query
            answer: Current answer
            gaps: Identified information gaps
            quality_score: Overall quality score
            
        Returns:
            List of refinement actions
        """
        actions = []

        # If quality is already high, minimal refinement needed
        if quality_score >= 0.8:
            return []

        # Handle each gap type
        for gap in gaps:
            if gap.gap_type == "missing_fact":
                actions.append(RefinementAction(
                    action_type="search",
                    description=f"Search for missing information: {gap.description}",
                    parameters={
                        "search_query": gap.suggested_search or query,
                        "gap_description": gap.description,
                    },
                    priority=1 if gap.severity == "high" else 2,
                ))

            elif gap.gap_type == "unverified":
                actions.append(RefinementAction(
                    action_type="verify",
                    description=f"Verify claim: {gap.description}",
                    parameters={
                        "claim": gap.description,
                    },
                    priority=2,
                ))

            elif gap.gap_type == "unclear":
                actions.append(RefinementAction(
                    action_type="clarify",
                    description=f"Clarify: {gap.description}",
                    parameters={
                        "unclear_part": gap.description,
                    },
                    priority=3,
                ))

            elif gap.gap_type == "outdated":
                actions.append(RefinementAction(
                    action_type="search",
                    description=f"Find recent information: {gap.description}",
                    parameters={
                        "search_query": gap.suggested_search or f"{query} latest recent",
                        "temporal": True,
                    },
                    priority=1,
                ))

        # General improvements based on quality score
        if quality_score < 0.5:
            actions.append(RefinementAction(
                action_type="expand",
                description="Expand answer with more detail",
                parameters={"reason": "Low quality score"},
                priority=2,
            ))

        # Sort by priority
        actions.sort(key=lambda a: a.priority)

        # Limit number of actions
        return actions[:5]

    def should_continue(self, quality_score: float) -> bool:
        """Determine if refinement should continue.
        
        Args:
            quality_score: Current quality score
            
        Returns:
            True if refinement should continue
        """
        if self._current_iteration >= self.max_iterations:
            return False

        # Continue if quality is below threshold
        return quality_score < 0.7

    def increment_iteration(self) -> int:
        """Increment the iteration counter.
        
        Returns:
            Current iteration number
        """
        self._current_iteration += 1
        return self._current_iteration

    def reset(self) -> None:
        """Reset the iteration counter."""
        self._current_iteration = 0

    def get_iteration(self) -> int:
        """Get current iteration number.
        
        Returns:
            Current iteration
        """
        return self._current_iteration

    def create_refined_query(
        self,
        original_query: str,
        action: RefinementAction,
    ) -> str:
        """Create a refined query for additional search.
        
        Args:
            original_query: Original user query
            action: Refinement action to execute
            
        Returns:
            Refined search query
        """
        if action.action_type == "search":
            return action.parameters.get("search_query", original_query)

        elif action.action_type == "verify":
            claim = action.parameters.get("claim", "")
            return f"verify {original_query} {claim}"

        elif action.action_type == "expand":
            return f"{original_query} detailed explanation"

        elif action.action_type == "clarify":
            unclear = action.parameters.get("unclear_part", "")
            return f"{original_query} {unclear} meaning definition"

        return original_query

    def merge_answers(
        self,
        original_answer: str,
        new_information: str,
        action: RefinementAction,
    ) -> str:
        """Merge new information into existing answer.
        
        Args:
            original_answer: Original answer text
            new_information: New information to incorporate
            action: The refinement action that produced this info
            
        Returns:
            Merged answer
        """
        # Simple merging strategy - append with context
        if action.action_type == "expand":
            return f"{original_answer}\n\n**Additional Details:**\n{new_information}"

        elif action.action_type == "verify":
            return f"{original_answer}\n\n**Verification:**\n{new_information}"

        elif action.action_type == "search":
            return f"{original_answer}\n\n**Additional Information:**\n{new_information}"

        elif action.action_type == "clarify":
            return f"{original_answer}\n\n**Clarification:**\n{new_information}"

        return f"{original_answer}\n\n{new_information}"

    def prioritize_actions(
        self,
        actions: list[RefinementAction],
        time_budget: float | None = None,
    ) -> list[RefinementAction]:
        """Prioritize refinement actions within constraints.
        
        Args:
            actions: List of actions to prioritize
            time_budget: Optional time budget in seconds
            
        Returns:
            Prioritized list of actions
        """
        # Estimate time per action type
        time_estimates = {
            "search": 2.0,
            "verify": 3.0,
            "expand": 1.5,
            "clarify": 1.0,
            "restructure": 0.5,
        }

        if time_budget is None:
            return sorted(actions, key=lambda a: a.priority)

        # Select actions that fit within budget
        prioritized = []
        remaining_time = time_budget

        sorted_actions = sorted(actions, key=lambda a: a.priority)

        for action in sorted_actions:
            estimated_time = time_estimates.get(action.action_type, 1.0)
            if estimated_time <= remaining_time:
                prioritized.append(action)
                remaining_time -= estimated_time

        return prioritized