app.py

import os
import json
import numpy as np
import gradio as gr
from huggingface_hub import hf_hub_download
from sentence_transformers import SentenceTransformer
from llama_cpp import Llama

# -----------------------------
# 1. LLM NorwAI Mistral 7B via GGUF
# -----------------------------
REPO_ID = "NorwAI/NorwAI-Mistral-7B-instruct"
GGUF_FILENAME = "norwai-mistral-7b-instruct-q4_k_m.gguf"  

HF_TOKEN = os.environ.get("HF_TOKEN")
if HF_TOKEN is None:
    raise RuntimeError(
        "HF_TOKEN environment variable is not set. "
    )

model_path = hf_hub_download(
    repo_id=REPO_ID,
    filename=GGUF_FILENAME,
    token=HF_TOKEN,
)

llm = Llama(
    model_path=model_path,
    n_ctx=4096,
    n_threads=2, #12   
    verbose=False,
)

# -----------------------------
# 2. Embedding model + knowledge base
# -----------------------------
EMBED_MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
embedder = SentenceTransformer(EMBED_MODEL_NAME)

KB_PATH = "kb_camhs_no.jsonl"
if not os.path.exists(KB_PATH):
    raise RuntimeError(f"Knowledge base file {KB_PATH} not found in repository.")

kb_docs = []
with open(KB_PATH, encoding="utf-8") as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        kb_docs.append(json.loads(line))

kb_texts = [d["text"] for d in kb_docs]
kb_embeddings = embedder.encode(
    kb_texts,
    convert_to_numpy=True,
    normalize_embeddings=True,
)

# -----------------------------
# 3. Retrive and Output formatting 
# -----------------------------
def retrieve_relevant_chunks(query: str, k: int = 4):
    """Return top-k relevant KB entries for a Norwegian clinical query."""
    q_emb = embedder.encode(
        [query],
        convert_to_numpy=True,
        normalize_embeddings=True,
    )[0]
    scores = kb_embeddings @ q_emb  # cosine similarity (unit vectors)
    top_idx = np.argsort(-scores)[:k]
    return [kb_docs[i] for i in top_idx]


def format_context(chunks):
    """Format retrieved chunks as numbered knowledge base blocks (for the prompt)."""
    lines = []
    for i, ch in enumerate(chunks, start=1):
        src = ch.get("source_short", "Ukjent kilde")
        cit = ch.get("citation", "")
        url = ch.get("url", "")
        text = ch.get("text", "")

        block = f"[{i}] Kilde: {src}\n"
        if cit:
            block += f"    Referanse: {cit}\n"
        if url:
            block += f"    URL: {url}\n"
        block += f"    Utdrag: {text}\n"
        lines.append(block)
    return "\n\n".join(lines)


def format_references_for_display(chunks):
    """
    Build a small, light-grey HTML block with the references that the user sees
    under the main answer.
    """
    if not chunks:
        return ""

    ref_lines = []
    for i, ch in enumerate(chunks, start=1):
        src = ch.get("source_short", "Ukjent kilde")
        cit = ch.get("citation", "")
        url = ch.get("url", "")

        # Keep it fairly short in the UI
        line = f"[{i}] {src}"
        if cit:
            line += f" – {cit}"
        if url:
            line += f" ({url})"
        ref_lines.append(line)

    refs_html = "<br>".join(ref_lines)
    styled_block = (
        "<br><span style='font-size: 0.8em; color: #888;'>"
        "<b>Kilder brukt i dette svaret:</b><br>"
        f"{refs_html}"
        "</span>"
    )
    return styled_block


SYSTEM_PROMPT = (
    "Du er en norsk faglig assistent for klinikere i barne- og ungdomspsykiatrien (BUP) i Norge. "
    "Du skal gi korte, presise og faglig korrekte svar basert på gjeldende norske retningslinjer, "
    "veiledere, pakkeforløp og lovverk (for eksempel Helsedirektoratet, Helseregisterloven, "
    "Pasientjournalloven). Internasjonale kilder som WHO og NICE og forskningsartikler (for eksempel "
    "IDDEAS prosjekt og CAMHS-studier) kan brukes som sekundær støtte når norske retningslinjer ikke er dekkende, "
    "men skal da omtales som henholdsvis internasjonale anbefalinger eller forskningsfunn.\n\n"
    "Viktige prinsipper:\n"
    "- Svar alltid på norsk.\n"
    "- Ikke still diagnose, ikke foreslå konkrete medikamentdoser, og ikke gi individuell behandlingsplan basert på en anonym tekst.\n"
    "- Beskriv heller hva retningslinjer, pakkeforløp, lovverk og forskning sier om utredning, vurdering, samarbeid og ansvar.\n"
    "- Vær tydelig på at svaret er generell informasjon og ikke erstatter klinisk vurdering eller lokale prosedyrer.\n"
    "- Bruk kun kunnskapsgrunnlaget du har fått utdrag fra; hvis informasjon mangler, si at du er usikker eller at det ikke er tydelig beskrevet.\n"
    "- Når du bruker et utdrag, henvis til kilden med [1], [2] osv. i brødteksten, der [n] samsvarer med nummeret i kunnskapsgrunnlaget.\n"
    "- Skriv svaret som sammenhengende tekst uten egne overskrifter som 'Forklaring:' eller 'Svar:'.\n"
    "- Ikke gjenta samme setning flere ganger; skriv budskapet én gang, tydelig og kort.\n"
)


def build_prompt(user_message: str, context_blocks):
    context_text = format_context(context_blocks)
    prompt = (
        f"{SYSTEM_PROMPT}\n\n"
        "Nedenfor får du relevante utdrag fra norske og internasjonale kilder. Bruk dem aktivt i svaret ditt, "
        "og henvis i teksten med [1], [2] osv. der det passer.\n\n"
        "Kunnskapsgrunnlag:\n"
        f"{context_text}\n\n"
        "Oppgave:\n"
        f"Bruker: {user_message}\n"
        "Ikke gjenta samme setning to ganger.\n"
        "Assistent (Svar på norsk, rettet mot klinikere i Norge, svar direkte i løpende tekst, uten egne overskrifter som 'Forklaring:' eller 'Svar:'.)"
    )
    return prompt


def chat_fn(message, history):
    # 1) Retrieve relevant guideline and research papers
    chunks = retrieve_relevant_chunks(message, k=4)

    # 2) Build grounded prompt to the model
    prompt = build_prompt(message, chunks)

    # 3) LLM
    out = llm(
        prompt,
        max_tokens=256,
        temperature=0.4,   
        top_p=0.9,
        stop=["Bruker:", "User:"],
    )
    main_reply = out["choices"][0]["text"].strip()

    #4) Build small, light gray reference block below the answer
    refs_html = format_references_for_display(chunks)

    # Return main text + small gray reference part
    full_reply = main_reply + refs_html
    return full_reply


demo = gr.ChatInterface(
    fn=chat_fn,
    title="Spør Datadrevet beslutningsstøtte for CAMHS Norway",
    description=(
        "Eksperimentell faglig assistent for barne- og ungdomspsykiatrien (BUP) i Norge. "
        "Svarene er generelle og basert på norske retningslinjer, pakkeforløp, lovverk og publisert forskning "
        "(inkludert IDDEAS prosjekt og norske CAMHS-studier), og kan ikke erstatte klinisk vurdering eller lokale prosedyrer."
        "\n\nGrunnleggende modell: **norwai-mistral-7b-instruct-q4_k_m.gguf**. fra NorwAI/NorwAI-Mistral-7B-instruct"
    ),
)

if __name__ == "__main__":
    demo.launch()