AI_engine/core/interactive_question.py

"""
Interactive Question Asking (Human-in-the-Loop) subsystem for LogiFlow AI / Doormile Agent.

Allows agents and skills to ask structured questions to human operators
with options, validation, and multi-select support.
"""
from dataclasses import dataclass, field, asdict
from typing import List, Optional, Dict, Any, Callable, Awaitable
import asyncio
import uuid
import time
from core.tool_registry import Tool, ToolParameter, register_tool, tool_registry
from core.logger import logger


@dataclass
class QuestionOption:
    id: str
    label: str
    description: Optional[str] = None
    badge: Optional[str] = None  # e.g. "12 AM–9 AM", "Step 1", "Recommended"
    metadata: Dict[str, Any] = field(default_factory=dict)


@dataclass
class Question:
    question_id: str
    prompt: str
    options: List[QuestionOption] = field(default_factory=list)
    is_multi_select: bool = False
    allow_custom_input: bool = True
    context: Dict[str, Any] = field(default_factory=dict)
    created_at: float = field(default_factory=time.time)
    answered: bool = False
    answer: Optional[Any] = None

    def to_dict(self) -> Dict[str, Any]:
        return {
            "question_id": self.question_id,
            "prompt": self.prompt,
            "options": [asdict(opt) for opt in self.options],
            "is_multi_select": self.is_multi_select,
            "allow_custom_input": self.allow_custom_input,
            "context": self.context,
            "created_at": self.created_at,
            "answered": self.answered,
            "answer": self.answer,
        }


class QuestionManager:
    """Tracks pending questions and handles human operator callbacks."""

    def __init__(self):
        self._pending: Dict[str, Question] = {}
        self._futures: Dict[str, asyncio.Future] = {}

    def create_question(
        self,
        prompt: str,
        options: Optional[List[Dict[str, Any]]] = None,
        is_multi_select: bool = False,
        allow_custom_input: bool = True,
        context: Optional[Dict[str, Any]] = None,
    ) -> Question:
        qid = f"q-{uuid.uuid4().hex[:8]}"
        formatted_options = []
        if options:
            for opt in options:
                if isinstance(opt, QuestionOption):
                    formatted_options.append(opt)
                elif isinstance(opt, dict):
                    formatted_options.append(
                        QuestionOption(
                            id=str(opt.get("id", opt.get("value", ""))),
                            label=str(opt.get("label", opt.get("text", ""))),
                            description=opt.get("description"),
                            badge=opt.get("badge"),
                            metadata=opt.get("metadata", {}),
                        )
                    )
                else:
                    formatted_options.append(QuestionOption(id=str(opt), label=str(opt)))

        question = Question(
            question_id=qid,
            prompt=prompt,
            options=formatted_options,
            is_multi_select=is_multi_select,
            allow_custom_input=allow_custom_input,
            context=context or {},
        )
        self._pending[qid] = question
        return question

    async def wait_for_answer(self, question: Question, timeout_s: float = 300.0) -> Any:
        """Asynchronously wait for human response to this question."""
        loop = asyncio.get_running_loop()
        fut = loop.create_future()
        self._futures[question.question_id] = fut

        try:
            answer = await asyncio.wait_for(fut, timeout=timeout_s)
            question.answered = True
            question.answer = answer
            return answer
        except asyncio.TimeoutError:
            logger.warning(f"Question {question.question_id} timed out after {timeout_s}s")
            question.answered = False
            raise TimeoutError(f"Question {question.question_id} timed out waiting for human input")
        finally:
            self._futures.pop(question.question_id, None)
            self._pending.pop(question.question_id, None)

    def answer_question(self, question_id: str, answer: Any) -> bool:
        """Called when human operator submits an answer from the UI.

        Returns False if the question is unknown or was already answered."""
        question = self._pending.get(question_id)
        if not question:
            logger.warning(f"No pending question found for id {question_id}")
            return False
        if question.answered:
            logger.warning(f"Question {question_id} already answered; ignoring")
            return False

        question.answered = True
        question.answer = answer

        fut = self._futures.get(question_id)
        if fut and not fut.done():
            fut.set_result(answer)  # wait_for_answer() removes it from _pending
        else:
            self._pending.pop(question_id, None)  # nobody waiting; don't leak
        return True

    def list_pending(self) -> List[Question]:
        return [q for q in self._pending.values() if not q.answered]

    def get_pending(self, question_id: str) -> Optional[Question]:
        return self._pending.get(question_id)


# Global singleton
question_manager = QuestionManager()


# Expose ask_question as a standard tool in ToolRegistry
async def ask_question_handler(
    prompt: str,
    options: Optional[List[Dict[str, Any]]] = None,
    is_multi_select: bool = False,
    allow_custom_input: bool = True,
    context: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """Tool handler for agents asking questions to the human operator."""
    q = question_manager.create_question(
        prompt=prompt,
        options=options,
        is_multi_select=is_multi_select,
        allow_custom_input=allow_custom_input,
        context=context,
    )
    logger.info(f"Agent asked question [{q.question_id}]: {prompt}")
    return {
        "status": "question_asked",
        "question": q.to_dict(),
    }


tool_registry.register(
    Tool(
        name="ask_question",
        description="Ask a clarifying or choice-based question to the human operator when input is missing or ambiguous.",
        handler=ask_question_handler,
        parameters=[
            ToolParameter(name="prompt", type="string", description="The question text to ask the operator"),
            ToolParameter(name="options", type="array", description="List of selectable options with id, label, badge", required=False),
            ToolParameter(name="is_multi_select", type="boolean", description="Whether multiple options can be chosen", required=False, default=False),
            ToolParameter(name="allow_custom_input", type="boolean", description="Whether user can write their own custom text", required=False, default=True),
        ],
        category="human_interaction",
    )
)