utils/model_generation.py

import json
import requests
from typing import List, Dict, Any, Union
import time
import numpy as np

import os


PROMPT_TEMPLATES = {
    "verbatim_sentiment": {
        "system": (
            "You are a compliance-grade policy analyst assistant. Prime directive: be faithful to the provided sources. "
            "Do NOT speculate. If the answer is not supported by the sources, say 'Not found in sources' and stop. "
            "Every non-trivial claim MUST be grounded with an inline citation in the form (filename p.X). "
            "Prefer 'unknown/not stated' over guessing. "
            "Follow this Grounding Protocol before answering: (1) read Context Sources; (2) extract exact quotes; "
            "(3) map each assertion to a citation; (4) list gaps and unknowns. "
            
            "Avoid hallucinations. Base everything strictly on the content provided. "
            "Output must be  complete sentences and adequate context. "
            "If sentiment or coherence inputs are disabled or empty, omit those sections entirely (do not mention they were omitted)."
            "Do not even write anything in sentiment and coherence if it is not available"
            "Try to meet the user's specification as much as possible where if they only want items from a certain page only give out data from that page or if it is from a certain document please only retrieve just from that document"
            "Order by page"
            "The context is already searched, retrieved and reranked when handed to you."


    ),
    # dynamic assembly; placeholders kept for backward compatibility but sections may be removed
    "user_template": "DYNAMIC"
    },

    "abstractive_summary": {
        "system": (
            "You are a policy analyst summarizing government documents for a general audience. "
            "Faithfulness is mandatory: paraphrase only what is supported by the sources and cite key claims inline (filename p.X). "
            "Avoid quotes unless legally binding language is essential. "
            "Bias toward completeness over brevity; use full sentences and helpful structure. "
            "If critical info is absent, say 'Not found in sources'—do not infer."
        ),
        "user_template": """Query: {query}

Write a comprehensive, plain-language summary with these sections:
- What It Covers (scope, entities, timelines) [cite]
- Key Requirements & Controls (what must be done) [cite]
- Enforcement & Penalties (who enforces, how, consequences) [cite]
- Deadlines & Effective Dates (explicit dates or 'not stated') [cite]
- Exemptions/Thresholds (if any; otherwise 'not stated') [cite]
- Risks & Open Questions (gaps/ambiguities; no speculation)
- Action Checklist (practical steps derived strictly from the sources) [cite]

Rules:
- Use citations for non-obvious claims (filename p.X).
- Avoid quotes unless a phrase is legally binding.
- If the sources do not answer the query, state 'Not found in sources'.

Topic hint: {topic_hint}

Context DOCS:
{context_block}
"""
    },

    "followup_reasoning": {
        "system": (
            "You are an assistant that explains policy documents interactively, reasoning step-by-step. "
            "Be strictly faithful to the documents; if a detail is absent, say so. "
            "Cite document filename and page for each factual claim. "
            "Favor clarity and completeness over brevity; full sentences only."
        ),
        "user_template": """User query: {query}

Answer step-by-step:
1) Direct Answer (what the sources actually support) with inline citations (filename p.X).
2) Why (short reasoning mapped to specific passages) with citations.
3) Edge Cases & Exceptions (only if present; otherwise 'not stated') with citations.
4) What’s Missing (explicitly note absent info; no speculation).

Then list 3–6 Follow-up Questions a reader might ask, and answer each using the docs.
- If a follow-up cannot be answered with the docs, respond: 'Not found in sources.'

Topic: {topic_hint}

DOCS:
{context_block}
"""
    },
}



# --- LLM client ---
def get_do_completion(api_key, model_name, messages, temperature=0.2, max_tokens=800):
    url = "https://inference.do-ai.run/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    data = {
        "model": model_name,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens
    }
    try:
        resp = requests.post(url, headers=headers, json=data, timeout=90)
        resp.raise_for_status()
        return resp.json()
    except requests.exceptions.HTTPError as e:
        print(f"HTTP error occurred: {e}")
        print(f"Response body: {e.response.text if e.response is not None else ''}")
        return None
    except requests.exceptions.RequestException as e:
        print(f"Request error: {e}")
        return None
    except json.JSONDecodeError as e:
        print(f"Failed to decode JSON: {e}")
        print(f"Response text: {resp.text if 'resp' in locals() else ''}")
        return None



# --- Prompt context builder ---
def _clip(text: str, max_chars: int = 1400) -> str:
    """Trim content to limit prompt size."""
    if not text:
        return ""
    text = str(text).strip()
    return text[:max_chars] + ("..." if len(text) > max_chars else "")


def build_context_block(top_docs: List[Dict[str, Any]]) -> str:
    """
    Formats each document with real citation:
    - Extracts file name from 'source' path
    - Uses 'page_label' or falls back to 'page'
    - Returns: <<<SOURCE: {filename}, p. {page_label}>>>
    """
    blocks = []
    for i, item in enumerate(top_docs):
        if hasattr(item, "page_content"):
            text = item.page_content
            meta = getattr(item, "metadata", {})
        else:
            text = item.get("text") or item.get("page_content", "")
            meta = item.get("metadata", {})

        # Get file name from path
        full_path = meta.get("source", "")
        filename = os.path.basename(full_path) if full_path else f"Document_{i+1}"

        # Prefer page_label if available, else fallback to raw page
        page_label = meta.get("page_label") or meta.get("page") or "unknown"

        citation = f"{filename}, p. {page_label}"

    blocks.append(f"<<<SOURCE: {citation}>>>\n{_clip(text)}\n</SOURCE>")

    return "\n".join(blocks)


# --- Message builder ---
def build_messages(
    query: str,
    top_docs: List[Dict[str, Any]],
    task_mode: str,
    sentiment_rollup: Dict[str, List[str]],
    coherence_report: str = "",
    topic_hint: str = "energy policy",
    allowlist_meta: Dict[str, Any] = None
) -> List[Dict[str, str]]:
    template = PROMPT_TEMPLATES.get(task_mode)
    if not template:
        raise ValueError(f"Unknown task mode: {task_mode}")

    context_block = build_context_block(top_docs)
    sentiment_present = bool(sentiment_rollup)
    coherence_present = bool(coherence_report)
    sentiment_json = json.dumps(sentiment_rollup or {}, ensure_ascii=False)

    # Build user prompt dynamically to truly omit absent sections
    parts = [
        f"Query: {query}\n",
        "Deliverables (omit any section whose input is empty/disabled):",
        "1) Quoted Policy Excerpts\n   - Quote the necessary text and append citations like (filename p.X). Group by subtopic.\n   - Honor any page or document restriction from the query strictly.\n   - Order by page",
    ]
    if sentiment_present:
        parts.append("2) Sentiment Summary\n   - Using the Sentiment JSON, explain tone, gaps, penalties, and enforcement clarity in plain English. Do not invent fields that aren't present.")
    if coherence_present:
        idx = 3 if sentiment_present else 2
        parts.append(f"{idx}) Coherence Assessment\n   - From the coherence report: on-topic vs off-topic; note coherent/off-topic/repeated sections only if present.")
    parts.append(
        "\nConstraints:\n- No external knowledge. No speculation. If a user ask is outside the sources, state 'Not found in sources.'\n- Use full sentences.\n- Each substantive statement has a citation."
    )
    parts.append(f"\nTopic hint: {topic_hint}\n")
    if sentiment_present:
        parts.append(f"Sentiment JSON (rolled-up across top docs):\n{sentiment_json}\n")
    if coherence_present:
        parts.append(f"Coherence report:\n{coherence_report}\n")
    guard = ""
    if allowlist_meta:
        doc_id = allowlist_meta.get('doc_id')
        pages = allowlist_meta.get('pages')
        guard = f"[ALLOWLIST_DOCS] doc_id={doc_id}; pages={pages}\nOnly use text from chunks where doc_id={doc_id} and page_label in {pages}. If none present reply exactly: Not found in sources for page {pages} of {doc_id}. Do not use any other documents.\n"
    parts.append(f"{guard}Context Sources:\n{context_block}")
    user_prompt = "\n".join(parts)

    return [
    {"role": "system", "content": template["system"]},
        {"role": "user", "content": user_prompt}
    ]


# --- Generation orchestrator ---
def generate_policy_answer(
    api_key: str,
    model_name: str,
    query: str,
    top_docs: List[Union[Dict[str, Any], Any]],
    sentiment_rollup: Dict[str, List[str]],
    coherence_report: str = "",
    task_mode: str = "verbatim_sentiment",
    temperature: float = 0.2,
    max_tokens: int = 2000
) -> str:
    if not top_docs:
        return "No documents available to answer."

    messages = build_messages(
        query=query,
        top_docs=top_docs,
        task_mode=task_mode,
        sentiment_rollup=sentiment_rollup,
        coherence_report=coherence_report
    )
    resp = get_do_completion(api_key, model_name, messages, temperature=temperature, max_tokens=max_tokens)
    if resp is None:
        return "Upstream model error. No response."
    try:
        return resp["choices"][0]["message"]["content"].strip()
    except Exception:
        return json.dumps(resp, indent=2)