src/agents/handlers/problem_statement.py

"""Problem Statement skill — guide the user to a usable problem statement.

Routed by the orchestrator (intent `problem_statement`) and callable as a skill.
An LLM drafts/refines the statement from the analysis title + the user's message and
declares what's still `missing`; a check validates only when nothing is missing. The
model is instructed to fill `objective`/`metric` ONLY from what the user explicitly
stated — a bare data question ("which X has the most Y?") leaves them in `missing`, so
it does not auto-validate (the gate stays meaningful). On a valid draft it persists
`problem_statement` + `problem_validated=True`; otherwise it streams guidance and
leaves the analysis un-validated.

NOTE: completeness is still a (hardened) LLM judgment — the truly deterministic gate
is an explicit user confirmation, planned with the frontend (see T3b / #11).

See `ORCHESTRATOR_REWORK_PLAN.md` §4 and the 2026-06-18 checkpoint.
"""

from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

from langchain_core.messages import BaseMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import Runnable
from langchain_openai import AzureChatOpenAI
from pydantic import BaseModel, Field

from src.middlewares.logging import get_logger

if TYPE_CHECKING:
    from src.agents.state_store import AnalysisStateStore

logger = get_logger("problem_statement")

_PROMPT_PATH = (
    Path(__file__).resolve().parent.parent.parent
    / "config"
    / "prompts"
    / "problem_statement.md"
)


class ProblemStatementDraft(BaseModel):
    """LLM output for the Problem Statement skill."""

    problem_statement: str = Field(
        ..., description="The refined, standalone problem statement (never empty)."
    )
    objective: str = Field(
        "", description="What success looks like — fill ONLY when the user explicitly "
        "stated it; never inferred from a data question. Empty otherwise."
    )
    metric: str = Field(
        "", description="The KPI to move/investigate — fill ONLY when the user "
        "explicitly stated it; never inferred from a data question. Empty otherwise."
    )
    missing: list[str] = Field(
        default_factory=list,
        description="Which of 'objective' / 'metric' the user has NOT explicitly stated "
        "yet. A bare data question leaves both here. Empty list = complete.",
    )
    feedback: str = Field(
        ...,
        description="Message to the user — guidance if incomplete, confirmation if complete.",
    )


def is_valid(draft: ProblemStatementDraft) -> bool:
    """Complete iff there's a statement and the model flagged nothing missing.

    Keying on the model's explicit `missing` list (rather than 'are objective/metric
    non-empty') is what stops a bare data question from auto-validating: the hardened
    prompt puts the un-stated parts in `missing`, so this returns False for it.
    """
    return bool(draft.problem_statement.strip()) and not draft.missing


def _load_prompt_text() -> str:
    return _PROMPT_PATH.read_text(encoding="utf-8")


def _build_default_chain() -> Runnable:
    from src.config.settings import settings

    llm = AzureChatOpenAI(
        azure_deployment=settings.azureai_deployment_name_4o,
        openai_api_version=settings.azureai_api_version_4o,
        azure_endpoint=settings.azureai_endpoint_url_4o,
        api_key=settings.azureai_api_key_4o,
        temperature=0,
    )
    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", _load_prompt_text()),
            MessagesPlaceholder(variable_name="history", optional=True),
            (
                "human",
                "Analysis title: {analysis_title}\n"
                "Current problem statement: {current}\n\n"
                "User message: {message}",
            ),
        ]
    )
    return prompt | llm.with_structured_output(ProblemStatementDraft)


class ProblemStatementAgent:
    """Single LLM call that drafts/refines a problem statement.

    Inject `chain` for tests; the default builds the Azure OpenAI chain on first use.
    """

    def __init__(self, chain: Runnable | None = None) -> None:
        self._chain = chain

    def _ensure_chain(self) -> Runnable:
        if self._chain is None:
            self._chain = _build_default_chain()
        return self._chain

    async def draft(
        self,
        message: str,
        analysis_title: str,
        current: str,
        history: list[BaseMessage] | None = None,
    ) -> ProblemStatementDraft:
        chain = self._ensure_chain()
        return await chain.ainvoke(
            {
                "message": message,
                "analysis_title": analysis_title,
                "current": current,
                "history": history or [],
            }
        )


async def run_problem_statement(
    message: str,
    analysis_id: str | None,
    *,
    agent: ProblemStatementAgent,
    store: AnalysisStateStore,
    history: list[BaseMessage] | None = None,
) -> str:
    """Draft + validate the problem statement; persist on a valid draft.

    Loads the current title/statement (if the analysis exists), drafts a refinement,
    runs the deterministic completeness check, and writes `problem_statement` +
    `problem_validated` back. Returns the user-facing feedback. When `analysis_id` is
    missing (e.g. a legacy room), it still drafts + returns guidance but cannot persist.
    """
    analysis_title, current = "New analysis", ""
    if analysis_id:
        state = await store.get(analysis_id)
        if state is not None:
            analysis_title, current = state.analysis_title, state.problem_statement

    draft = await agent.draft(message, analysis_title, current, history)
    validated = is_valid(draft)

    if analysis_id:
        await store.update(
            analysis_id,
            problem_statement=draft.problem_statement,
            problem_validated=validated,
        )
    logger.info("problem_statement drafted", analysis_id=analysis_id, validated=validated)
    return draft.feedback