backend/pipeline/reference_mapper.py

"""In-place fact update on REFERENCE-tier baseline text — no scratch generation."""

from __future__ import annotations

import logging

from backend.config import settings
from backend.llm import openai_client
from backend.models.schema import TemplateSchema
from backend.pipeline.composition_output import (
    accept_narrative_section_output,
    sanitize_section_prose,
)
from backend.pipeline.paragraph_merge import merge_observations_into_paragraph
from backend.prompts.mapping_prompt import (
    FACT_GROUNDING_RULES,
    MAPPING_SYSTEM_BASE_MAXIMUM,
    MAPPING_SYSTEM_BASE_MEDIUM,
    MAPPING_SYSTEM_BASE_MINIMUM,
    MAPPING_USER_TEMPLATE,
    RATING_HINT_TEMPLATE,
    RICS_DOMAIN_RULES,
    _observations_bulleted,
)
from backend.prompts.prompt_few_shot_examples import (
    MAPPING_COT_PROTOCOL,
    MAPPING_FEW_SHOT,
)
from backend.prompts.prompt_message_assembly import (
    append_cot_to_system,
    apply_dynamic_literature,
    inject_few_shot_turns,
)
from backend.rag.retriever import InterferenceLevel

logger = logging.getLogger(__name__)

# Relocated verbatim to backend/prompts/mapping_prompt.py (test-exempt). Aliased
# here to preserve the existing internal/private reference and import surface.
_RICS_DOMAIN_RULES = RICS_DOMAIN_RULES


def select_mapping_prompt(interference_level: str) -> str:
    """Return the tier-specific mapping system prompt base for the given level."""
    level = (interference_level or "maximum").strip().lower()
    if level == "minimum":
        return MAPPING_SYSTEM_BASE_MINIMUM
    if level == "medium":
        return MAPPING_SYSTEM_BASE_MEDIUM
    return MAPPING_SYSTEM_BASE_MAXIMUM


def _normalize_interference_level(
    interference_level: str | InterferenceLevel | None,
) -> InterferenceLevel:
    """Map caller input to a supported composition tier; default maximum."""
    raw = str(interference_level or "").strip().lower()
    if raw in ("minimum", "medium", "maximum"):
        return raw  # type: ignore[return-value]
    return "maximum"


def compose_mapping_system_prompt(
    interference_level: str | InterferenceLevel | None,
) -> str:
    """Full mapping system prompt: tier base + shared RICS domain rules."""
    level = _normalize_interference_level(interference_level)
    return (
        select_mapping_prompt(level).strip()
        + "\n\n"
        + FACT_GROUNDING_RULES.strip()
        + "\n\n"
        + _RICS_DOMAIN_RULES.strip()
    )


def build_interference_messages(
    interference_level: str | InterferenceLevel | None,
    *,
    observations: list[str],
    baseline: str,
    schema: TemplateSchema,
    section_id: str = "",
    section_title: str = "",
    rating_value: str | None = None,
    extra_references: list[str] | None = None,
) -> list[dict[str, str]]:
    """Select the prompt builder for minimum / medium / maximum AI involvement."""
    level = _normalize_interference_level(interference_level)
    baseline_text = baseline.strip()
    if extra_references:
        extras = "\n\n".join(item.strip() for item in extra_references if item.strip())
        if extras:
            baseline_text = f"{baseline_text}\n\n{extras}".strip()

    rating_line = ""
    if schema.rating_system.detected and rating_value:
        rating_line = RATING_HINT_TEMPLATE.format(rating_value=rating_value)

    system = append_cot_to_system(
        compose_mapping_system_prompt(level), MAPPING_COT_PROTOCOL
    )
    user = MAPPING_USER_TEMPLATE.format(
        section_id=section_id or "—",
        section_label=section_title or section_id or "—",
        rating_line=rating_line,
        first_reference_baseline_paragraph=baseline_text or "(none)",
        observations_bulleted=_observations_bulleted(observations),
    )

    messages = inject_few_shot_turns(
        [
            {"role": "system", "content": system.strip()},
            {"role": "user", "content": user.strip()},
        ],
        MAPPING_FEW_SHOT,
    )
    query = "\n".join(
        [section_title or section_id or "", *(o.strip() for o in observations)]
    ).strip()
    return apply_dynamic_literature(messages, query=query)


def map_inplace_baseline(
    baseline_text: str,
    observations: list[str],
    schema: TemplateSchema,
    *,
    section_id: str = "",
    section_title: str = "",
    rating_value: str | None = None,
    messages: list[dict[str, str]] | None = None,
) -> str:
    """Apply in-place fact updates to the retrieved past-report baseline only."""
    baseline = (baseline_text or "").strip()
    if not baseline:
        return ""
    if not observations:
        return baseline

    merged = merge_observations_into_paragraph(baseline, observations, schema)

    if settings.use_llm_paragraph_mapping and openai_client.is_available():
        llm_messages = messages or build_interference_messages(
            "maximum",
            observations=observations,
            baseline=baseline,
            schema=schema,
            section_id=section_id,
            section_title=section_title,
            rating_value=rating_value,
        )
        try:
            out = openai_client.chat_text(
                llm_messages,
                model=settings.mapping_model,
                max_tokens=settings.max_tokens_mapping,
                temperature=0.0,
                call_label="mapping",
            )
            text = sanitize_section_prose((out or "").strip())
            if accept_narrative_section_output(text, observations):
                return text
        except Exception as exc:  # noqa: BLE001
            logger.warning(
                "In-place LLM mapping failed (%s); using deterministic merge.", exc
            )

    return merged


def map_reference_paragraph(
    reference_paragraph: str,
    observations: list[str],
    schema: TemplateSchema,
    interference_level: InterferenceLevel,
    *,
    section_id: str = "",
    section_title: str = "",
    rating_value: str | None = None,
    extra_references: list[str] | None = None,
) -> str:
    """Map notes onto the assembled REFERENCE baseline using tier-specific prompts."""
    level = _normalize_interference_level(interference_level)
    messages = build_interference_messages(
        level,
        observations=observations,
        baseline=reference_paragraph,
        schema=schema,
        section_id=section_id,
        section_title=section_title,
        rating_value=rating_value,
        extra_references=extra_references,
    )
    return map_inplace_baseline(
        reference_paragraph,
        observations,
        schema,
        section_id=section_id,
        section_title=section_title,
        rating_value=rating_value,
        messages=messages,
    )