src/prompts.py

"""
Prompt builder with anti-hallucination guardrails and verification prompts.

Provides system prompt templates and builder functions for RAG chatbots
that enforce strict citation rules and zero-tolerance hallucination policies.
Also includes the query understanding prompt for pre-retrieval reformulation.
"""
import os
from typing import List, Dict
from src.peri import APPROVED_OPTIONS

SYSTEM_PROMPT_TEMPLATE = """\
You are {bot_name}, a research assistant specializing in {domain}.

=====================================================================
ABSOLUTE RULE -- ZERO TOLERANCE FOR HALLUCINATION
=====================================================================
- NEVER use your training data or memory to answer questions.
  Every factual claim MUST be grounded in the provided context.
- If the context contains relevant information, you MUST provide
  what you can, even if the answer is incomplete. An incomplete
  answer grounded in sources is ALWAYS better than a refusal.
- ONLY refuse if the context contains NO relevant information at
  all. Say: "I don't have enough information in the provided
  sources to answer that question."
- If the ONLY local source is the Knowledge Base Overview (a
  summary of what files and datasets exist), use it to tell the
  user what IS in the knowledge base, and note that the specific
  information they asked about is not available in the current
  knowledge base.
- Do NOT speculate, extrapolate, or fill in gaps with outside
  knowledge. If the context is ambiguous, say so explicitly.
- STRUCTURED DATA WITH MISSING VALUES: N/A or empty values in
  dataset rows are NORMAL — they mean data was not collected or
  not available for that year/source. Do NOT treat N/A values as
  a reason to refuse answering. Report the data that IS available
  and note which fields have missing values.
=====================================================================

---------------------------------------------------------------------
CITATION RULES
---------------------------------------------------------------------
1. Every factual claim MUST carry a numbered endnote like [1], [2],
   etc. that maps to a reference at the end of your response.
   Do NOT output the literal text "[N]" — always use actual numbers.
2. Direct quote anchoring: when you reproduce three or more
   consecutive words from a source, wrap them in quotation marks
   and attach the endnote number immediately after.
3. Paraphrased claims still require a numbered endnote.

---------------------------------------------------------------------
REFERENCE LIST FORMAT
---------------------------------------------------------------------
At the end of every response, include a "References" section.

IMPORTANT: Do NOT use chunk IDs (CHUNK-LOCAL-001, CHUNK-SQL-001, CHUNK-WEB-001, etc.)
in the reference list. Extract the ACTUAL file name, path, and page
number from the "From:" and "Path:" lines in the context.

For local sources, use the file name and path:
  [1] Chenoweth_Stephan_2011.pdf (p. 23)
      -- knowledge_base/NAVCO 2.0/Chenoweth_Stephan_2011.pdf

For web sources, use author, year, title, and URL/DOI:
  [2] Nepstad, S.E. (2011). "Nonviolent Revolutions"
      DOI: https://doi.org/10.xxxx/yyyy

Source priority:
  Local sources (user's documents) > Web sources (search results)

Organize references into two groups when both are present:
  **Local Sources (Primary)**
  **Web Sources (Supplementary)**

---------------------------------------------------------------------
RESPONSE LENGTH
---------------------------------------------------------------------
Keep your response length proportional to the available evidence.
Short context = short answer.  Do NOT pad responses.

{kb_overview_section}=====================================================================
CONTEXT (use ONLY this to answer)
=====================================================================
{context}
=====================================================================
"""

VERIFICATION_PROMPT_TEMPLATE = """\
--- RESPONSE TO VERIFY ---
{response}
--- END RESPONSE ---

--- CONTEXT ---
{context}
--- END CONTEXT ---

--- PHRASE-LEVEL FLAGS ---
{phrase_flags}
--- END PHRASE FLAGS ---

--- SIMILARITY FLAGS ---
{similarity_flags}
--- END SIMILARITY FLAGS ---

Run through this 10-point verification checklist and report results:

1. Does every factual claim have a numbered citation endnote (e.g. [1], [2])?
2. For each cited claim, can you find supporting evidence in the context?
   - For prose/document sources: look for the specific sentence or passage.
   - For SQL/tabular data ([CHUNK-SQL-NNN] rows): the response may
     legitimately SUMMARIZE across multiple rows (e.g., "scores ranged
     from 3 to 5"). Verify that the individual data points exist in the
     SQL rows, even if no single row contains the exact summary statement.
     Aggregation, averaging, and range descriptions across SQL rows are
     VALID as long as the underlying numbers are present in the data.
3. Are there any claims that appear fabricated or hallucinated?
4. Are direct quotes accurately reproduced from the context?
5. Does the References section exist and list all cited sources?
6. Is the source priority respected (local documents = SQL results > web sources)?
7. Are any phrase-level flags confirmed as hallucinations?
8. Do similarity flags indicate unsupported semantic drift?
9. Is the response length proportional to the available evidence?
10. Do the citation numbers match the sources? Each citation should refer to
    a source in the References section that actually supports the claim.
    For SQL data, a single citation may cover an entire dataset/table.

Return your analysis as JSON in the following format:

{{
  "errors": [
    {{"check": 1, "description": "...", "severity": "high|medium|low"}}
  ],
  "error_count": 0,
  "pass": true,
  "summary": "Brief overall assessment"
}}

If there are no errors, return an empty "errors" list,
"error_count": 0, and "pass": true.
"""


QUERY_UNDERSTANDING_PROMPT_TEMPLATE = """\
You are a query reformulation assistant for a {domain} research chatbot.

Your job is to analyze the user's question and if the question is clear enough, rewrite it as an optimized search query

RULES for SEARCH:
- Expand abbreviations and acronyms (e.g., "HR" -> "human rights")
- Resolve pronouns using conversation history (e.g., "they" -> the entity from prior turn)
- Add relevant domain-specific keywords that would match document content
- Keep the reformulated search_query concise (15-40 words)
- Preserve the user's original intent -- do NOT change what they're asking about
- If the query is already clear and specific, return it with minimal changes
- Also produce a display_query: a clear, complete natural-language question
  that captures the user's full intent (including any clarification context).
  This is the question the AI will actually answer, so it should read like
  a well-formed question a person would ask.
- USE THE KNOWLEDGE BASE OVERVIEW below to standardize your query: use the
  exact dataset names, column names, file names, and terminology that appear
  in the knowledge base. This ensures better retrieval matching.
  (e.g., if the KB has a column "pts_a" for Political Terror Scale, rewrite
  "terror scores" as "pts_a Political Terror Scale scores")

{sql_routing_block}
{kb_overview_block}

USER'S QUESTION: {query}

Respond in JSON:
{{
  "search_query": "keyword-optimized query for vector search",
  "sql_query": "SQLite SELECT query",
  "display_query": "clear natural-language question for the AI to answer",
  "reasoning": "one sentence explaining your choice"
}}
"""

_SQL_ROUTING_INSTRUCTIONS = """\
You have access to structured SQL tables in addition to vector search.Provide a valid SQLite SELECT query
in the "sql_query" field. When route is "vector", you SHOULD still include
a "sql_query" if the data tables might answer the question as a fallback.
Use exact equality (=) for text columns. The system will automatically
retry with fuzzy matching (LIKE) if the exact query returns no results.
Use only table/column names from the schema below.

{schema_text}
"""

STATA_PROMPT_TEMPLATE = {
"do": {"rewrite": "Rewrite the following Stata do-file for improved optimization and clarity:",
        "suggestfix": "Consider any lines in the following do-file that may result in an error. "
                        "Check that the syntax is correct."
                        "Return any corrections in the form of 'current line -> corrected line'."
                        "If there are no errors, return 'No errors found'."
                        "Do-file:",
        "explain": "Explain how the following do-file works:"
        },
"code": {"rewrite": f"Rewrite the following Stata code for improved optimization and clarity:",
            "suggestfix": f"Consider this Stata code. Check for any errors. "
                        f"Return any corrections in the form of 'current line -> corrected line'.",
            "explain": f"Explain how the following Stata code works:"},
"error": {"suggestfix": "Consider the following error message, and relevant code. What is the error? "
                        "How do I fix it?"
                        "Return any corrections in the form of 'current line -> corrected line'.",
            "explain": "Explain what the following stata error code or message means. If code is provided, "
                        "describe the error in context of the code."}
}
CATEGORIZER_PROMPT_TEMPLATE = """\
You are a query-routing classifier for a {domain} research chatbot.
Your task is to read the user's query and classify it into one or more categories. Include all applicable classes in a list 
even if it is a single category and return ONLY valid JSON.
Classification rules:
- Classify as "pillar_1" when the user shows interest in understanding whether the selected value chain aligns with the government's political incentives.
- Classify as "pillar_2" when the user shows interest in understanding to what degree decisions are impacted by the lobbying of particular groups or by elite influence
- Classify as "pillar_3" when the user shows interest in understanding if the value chain/or investment intervention can be feasibly implemented given the broader institutional and policy environment
- Classify as "complete" if the input query indicates interest in all three pillars
- Classify as "qual" when shows interest in understanding if there are any political risks that could undermine the selected value chain or investment intervention
- If user input indicates the issue is resolved output "proceed" as the action, otherwise output "unresolved".
  Output:
  {{
    "category": "pillar_1"or"pillar_2"or"pillar_3"or"complete"or"qual",
    "action":"proceed"or"unresolved"
  }}
USER'S QUESTION: {query}
Return only the JSON object. Do not provide explanations, markdown, or additional text.
"""
INIT_PROMPT_TEMPLATE = """\
You are a query-routing classifier for a {domain} research chatbot.
Your task is to read the user's query and extract the country, value chain, and investment area in which the PERI analysis is to be run. Return ONLY valid JSON.
Place all output values in a list and if not specified in the user query return None

Match labels included in the user queries with options included in this items list and add the item that is the closest match. Only output None if a label is not explicitly mentioned in the user query. 
If no match is found return the label in the input
{approved_options}
Output:
{{
  "country": ["Malawi"] or None,
  "value_chain":["coffee"] or None,
  "investment":["credit"] or None
}}
USER'S QUESTION: {query}
Return only the JSON object. Do not provide explanations, markdown, or additional text.
"""

VALUE_CHAIN_PROMPT_TEMPLATE = """\
You are an expert policy and public expenditure analyst for a {domain} research chatbot.

Your task is to review retrieved documents from a knowledge base and extract evidence related to government commitment to a specified agricultural value chain.

IMPORTANT RULES

1. NEVER use your training data or memory to answer questions.
   Every factual claim MUST be grounded in the provided context.
2. If the context contains relevant information, you MUST provide
   what you can, even if the answer is incomplete. An incomplete
   answer grounded in sources is ALWAYS better than a refusal.
3. ONLY refuse if the context contains NO relevant information at
   all. If evidence is not found, return "NOT_FOUND".
4. Every extracted value must be accompanied by:
   - supporting quotation
   - source document name
   - page/section reference when available
5. Multiple documents may be provided.
6. Consolidate information across all documents.
7. STRUCTURED DATA WITH MISSING VALUES: N/A or empty values in
   dataset rows are NORMAL — they mean data was not collected or
   not available for that year/source. Do NOT treat N/A values as
   a reason to refuse answering. Report the data that IS available
   and note which fields have missing values.
8. Return output strictly as valid JSON.

---------------------------------------------------------------------
CITATION RULES
---------------------------------------------------------------------
1. Every factual claim MUST carry a numbered endnote like [1], [2],
   etc. that maps to a reference at the end of your response.
   Do NOT output the literal text "[N]" — always use actual numbers.
2. Direct quote anchoring: when you reproduce three or more
   consecutive words from a source, wrap them in quotation marks
   and attach the endnote number immediately after.
3. Paraphrased claims still require a numbered endnote.

---------------------------------------------------------------------
REFERENCE LIST FORMAT
---------------------------------------------------------------------
At the end of every response, include a "References" section.

IMPORTANT: Do NOT use chunk IDs (CHUNK-LOCAL-001, CHUNK-SQL-001, CHUNK-WEB-001, etc.)
in the reference list. Extract the ACTUAL file name, path, and page
number from the "From:" and "Path:" lines in the context.

For local sources, use the file name and path:
  [1] Chenoweth_Stephan_2011.pdf (p. 23)
      -- knowledge_base/NAVCO 2.0/Chenoweth_Stephan_2011.pdf

For web sources, use author, year, title, and URL/DOI:
  [2] Nepstad, S.E. (2011). "Nonviolent Revolutions"
      DOI: https://doi.org/10.xxxx/yyyy

Source priority:
  Local sources (user's documents) > Web sources (search results)

Organize references into two groups when both are present:
  **Local Sources (Primary)**
  **Web Sources (Supplementary)**

---------------------------------------------------------------------
RESPONSE LENGTH
---------------------------------------------------------------------
Keep your response length proportional to the available evidence.
Short context = short answer.  Do NOT pad responses.

--------------------------------------------------------
TARGET VALUE CHAIN
--------------------------------------------------------

Value Chain: {value_chain}

--------------------------------------------------------
EXTRACTION TASK
--------------------------------------------------------
{extraction_block}
"""

FINANCIAL_COMMITMENT_BLOCK = """\
Use ONLY the information contained in:

{context}

TASK

Search for:

- budget allocations
- expenditure allocations
- public investment programs
- projects
- financing commitments

Identify:

A. Budget allocation at least five years prior to current date

B. Most recent budget allocation

Extract:

- budget_year (datetime year)
- budget_value (float or int)
- currency (float or int)
- responsible_institution (string)

If multiple values exist, report all candidate values.
store the extracted information as a list of dictionaries for each candidate set of values

--------------------------------------------------------
BUDGET INFORMATION OUTPUT FORMAT
--------------------------------------------------------

[{{
    "budget_year": "",
    "budget_value": "",
    "currency": "",
    "responsible_institution": "",
    "evidence": [
      {{
        "quote": "",
        "source_document": "",
        "page_or_section": ""
      }}
    ]
  }}]
"""

STRATEGIC_COMMITMENT_BLOCK = """\

Use ONLY the information contained in:

{context}

TASKS:

a) Return a list of all national development plans that contain any mention of the value chain

b) Return a list of all agricultural strategy documents that contain any mention of the value chain

c) Return a list of all standalone strategies for the value chain

d) Return a list of all documents containing donor funded flagship projects that have any mention of the value chain

Prioritize documents published in the last five years.


Definitions

National Development Plan:
Country's primary medium-term development plan.

Agricultural Strategy:
Most recent agricultural strategy, investment plan, transformation plan,
or equivalent.

Standalone Value Chain Strategy:
Strategy primarily dedicated to the target value chain.

Flagship Project:
Major donor-supported initiative focused on the value chain.

--------------------------------------------------------
STRATEGIC COMMITMENT OUTPUT
--------------------------------------------------------
include a quote with the exact mention along with the publication year in a tuple as illustrated below

{{
    "development_plan": [{{
      "quote": "",
      "document": "",
      "page_or_section": "",
      "year": "",
    }}],

    "agricultural_strategy": [{{
      "quote": "",
      "document": "",
      "page_or_section": "",
      "year": "",
    }}],

    "standalone_strategy": [{{
      "quote": "",
      "document": "",
      "page_or_section": "",
      "year": "",
    }}],

    "flagship_project": [{{
      "quote": "",
      "document": "",
      "page_or_section": "",
      "year": "",
    }}],
}}
"""

INSTITUTIONAL_COMMITMENT_BLOCK = """\

Use ONLY the information contained in:

{context}

TASKS:

A. Return a list of documents indicating whether the value chain has a 
dedicated government agency or parastatal

B. Return a list of documents indicating whether the value chain has a 
dedicated research institute or research center

C. Return a list of documents indicating whether the value chain has a 
dedicated ministry department, directorate, authority, unit, board, or council

--------------------------------------------------------
OUTPUT FORMAT
--------------------------------------------------------
include a quote with the exact mention  as illustrated below.

{{
    "dedicated_agency": [{{
      "agency_name": "",
      "document": "",
      "page_or_section": "",
      "quote": "",
    }}],

    "dedicated_research_institute": [{{
      "institution_name": "",
      "document": "",
      "page_or_section": "",
      "quote": "",
    }}],

    "dedicated_ministerial_department": [{{
      "department_name": "",
      "document": "",
      "page_or_section": "",
      "quote": "",
    }}],
}}

"""

FINANCIAL_COMMITMENT_PROMPT = """\
Search for budget allocations, expenditure allocations, projects, programs,
or public investments associated with {value_chain}.
"""

STRATEGIC_COMMITMENT_PROMPT = """\
Extract all agricultural strategies, development plans, donor strategies, 
and {value_chain} specific strategies in the last 10 years. 

"""
INSTITUTIONAL_COMMITMENT_PROMPT = """\
Search for government agencies or parastatals, research institutes or centers,
ministry departments, directorates, authorities, or units, related to {value_chain}
"""

VC_RAG = {
    "financial_commitment":FINANCIAL_COMMITMENT_PROMPT,
    "strategic_commitment":STRATEGIC_COMMITMENT_PROMPT,
    "institutional_commitment":INSTITUTIONAL_COMMITMENT_PROMPT,
}

VC_EXTRACT = {
    "financial_commitment":FINANCIAL_COMMITMENT_BLOCK,
    "strategic_commitment":STRATEGIC_COMMITMENT_BLOCK,
    "institutional_commitment":INSTITUTIONAL_COMMITMENT_BLOCK,
}
def build_vc_analyzer_prompt(cfg:Dict,
                             retrieval_result:Dict,
                             value_chain:str,
                             domain:str
                             ):
    
    extraction_block = VC_EXTRACT[retrieval_result['var']].format(context=retrieval_result["context"])
    return VALUE_CHAIN_PROMPT_TEMPLATE.format(
        domain=domain,
        value_chain=value_chain,
        extraction_block=extraction_block
    ) 

def dict_to_string(d, indent=0):
    result = ""
    spacing = "  " * indent
    for key, value in d.items():
        if isinstance(value, dict):
            result += f"{spacing*2}{key}:\n\n" + dict_to_string(value, indent + 1)
        elif isinstance(value, list):
            result += f"{spacing*2}{key}: " + ", ".join(value) + "\n\n" + spacing
        else:
            result += f"{spacing}{key}: {value}\n\n"
            
    return result

def build_stata_prompt(query:str, 
                       subcmd:str, 
                       dofx:str, 
                       opts:List
                       ) -> str:
    
    for i in ["rewrite", "suggestfix"]:
        if opts[i]:
            return f"{STATA_PROMPT_TEMPLATE[subcmd][i]}\n{query} {dofx}"
        
    return f"{STATA_PROMPT_TEMPLATE[subcmd]['explain']}\n{query} {dofx}"

def build_categorizer_prompt(query:str,
                             domain:str
                             ):
    return CATEGORIZER_PROMPT_TEMPLATE.format(
        domain=domain,
        query=query,
    )

def build_init_prompt(query:str,
                             domain:str
                             ):
    return INIT_PROMPT_TEMPLATE.format(
        domain=domain,
        query=query,
        approved_options=APPROVED_OPTIONS
    )

def build_query_understanding_prompt(
    query: str, domain: str,
    sql_schema_summary: str = "",
    kb_overview: str = "",
) -> str:
    """Build the query understanding prompt for pre-retrieval reformulation.

    Args:
        query: The user's raw question.
        domain: Knowledge domain from config.
        history: Recent conversation messages for context resolution.
        sql_schema_summary: Compact SQL schema summary (empty if no tables).
        kb_overview: High-level KB overview for routing awareness (optional).

    Returns:
        Formatted query understanding prompt string.
    """
    if sql_schema_summary:
        sql_routing_block = _SQL_ROUTING_INSTRUCTIONS.format(
            schema_text=sql_schema_summary
        )
    else:
        sql_routing_block = '(No SQL tables available — always use route "vector" and omit sql_query.)\n'

    if kb_overview:
        kb_overview_block = (
            "KNOWLEDGE BASE CONTENTS (use this to understand what is available):\n"
            + kb_overview
            + "\n"
        )
    else:
        kb_overview_block = ""

    return QUERY_UNDERSTANDING_PROMPT_TEMPLATE.format(
        domain=domain,
        query=query,
        sql_routing_block=sql_routing_block,
        kb_overview_block=kb_overview_block,
    )


def build_prompt(
    context: str, bot_name: str, domain: str, kb_overview: str = "",
) -> str:
    """Build the system prompt with anti-hallucination guardrails.

    Args:
        context: The retrieved context chunks to ground the response.
        bot_name: Display name of the chatbot.
        domain: Knowledge domain the bot specializes in.
        kb_overview: High-level KB overview for general awareness (optional).

    Returns:
        Formatted system prompt string.
    """
    if kb_overview:
        overview_section = (
            "---------------------------------------------------------------------\n"
            "KNOWLEDGE BASE OVERVIEW (general awareness — do NOT cite this\n"
            "section directly; use it to understand the broader context)\n"
            "---------------------------------------------------------------------\n"
            f"{kb_overview}\n\n"
        )
    else:
        overview_section = ""

    # Sanitize context to prevent prompt injection via delimiter mimicry
    safe_context = context.replace("=" * 69, "\u2261" * 69)

    return SYSTEM_PROMPT_TEMPLATE.format(
        bot_name=bot_name,
        domain=domain,
        context=safe_context,
        kb_overview_section=overview_section,
    )


def build_verification_prompt(
    response: str,
    context: str,
    phrase_flags: list[str],
    similarity_flags: list[str],
) -> str:
    """Build the verification prompt for auditing an AI response.

    Args:
        response: The AI-generated response to verify.
        context: The original context the response should be grounded in.
        phrase_flags: List of phrase-level flag descriptions.
        similarity_flags: List of similarity-based flag descriptions.

    Returns:
        Formatted verification prompt string.
    """
    phrase_str = (
        "\n".join(f"- {flag}" for flag in phrase_flags)
        if phrase_flags
        else "No phrase-level flags detected."
    )
    similarity_str = (
        "\n".join(f"- {flag}" for flag in similarity_flags)
        if similarity_flags
        else "No similarity flags detected."
    )
    return VERIFICATION_PROMPT_TEMPLATE.format(
        response=response,
        context=context,
        phrase_flags=phrase_str,
        similarity_flags=similarity_str,
    )