doctate/server/src/analyze/backend.rs

//! LLM backend catalog.
//!
//! A "backend" is a fully self-describing profile: model + sampling params +
//! system prompt + endpoint + auth. Multiple profiles per model are allowed
//! (e.g. several Llama-3.1 variants with different temperatures).
//!
//! Provider constructors (`ionos_backend`, future `ollama_backend`) own the
//! "auth + endpoint" knowledge for their family; only globally-shared values
//! (temperature, max_completion_tokens, timeout, use_json_schema, default
//! prompt) live as module-level `const`s. Adding a new profile within an
//! existing provider is a one-line `vec!` entry; adding a new provider is a
//! new constructor function.

use std::sync::OnceLock;

const DEFAULT_TEMPERATURE: f32 = 0.5;
const DEFAULT_MAX_COMPLETION_TOKENS: u32 = 4096;
const DEFAULT_TIMEOUT_SECONDS: u64 = 180;
const DEFAULT_USE_JSON_SCHEMA: bool = true;

/// Llama 3.1 405B FP8 needs all three of:
///   1. sampling (`temperature: 0.6`, `top_p: 0.9`) — without it the model
///      enters infinite repetition loops on prompts with conflicting
///      constraints (HuggingFace meta-llama/Llama-3.1-8B-Instruct disc. #32,
///      Ionos docs for the 405B model).
///   2. `response_format: json_schema` — provides a hard server-side stop via
///      vLLM grammar-constrained decoding.
///   3. an explicit format instruction telling the model what to put inside
///      `document` — without it the model emits `{"document": ""}` (silent
///      deletion, observed in test T4 on 2026-05-03).
const LLAMA_TEMPERATURE: f32 = 0.6;
const LLAMA_TOP_P: f32 = 0.9;

/// Format instruction injected as a separate `system` message for backends
/// that aren't natively schema-aware. Kept as a Rust constant rather than a
/// prompt-file edit so the medical prompt stays purely about clinical rules.
const LLAMA_FORMAT_INSTRUCTION: &str = "AUSGABEFORMAT: Antworte ausschließlich \
    als JSON-Objekt der Form {\"document\": \"<konsolidierter Fließtext>\"}. \
    Der Wert von \"document\" ist der Fließtext nach allen obigen Regeln. \
    Keine zusätzlichen Felder, kein Text außerhalb des JSON.";

/// Default system prompt for the consolidation LLM. Used by all backends
/// unless overridden via `with_system_prompt`.
///
/// Sandbox-validated against case `c414cf52` (3 runs each, fair pipeline
/// with pre- and post-LLM gazetteer pass against the current vocab):
/// - eliminates two hallucination classes the previous prompt produced —
///   unmarked "Hypotonie" when the dictation said "Hypertonie" (0/3 vs 2/3),
///   and inventing a unit ("35 ng/l") for a dictated value without unit
///   (0/3 vs 2/3)
/// - keeps Latin terms intact ("Punctum Maximum", "Apex Cordis", "Axilla"
///   in 3/3 vs 2/3)
/// - relies on the gazetteer for drug-name correction (Enoxaparin etc.) —
///   the prompt no longer needs a separate "AKTIVE KORREKTUR" hammer-block
///   because the gazetteer's pre-LLM pass repairs typical typos before the
///   LLM ever sees them
/// - block layout: AUFGABE / QUELLE / KORREKTUR-POLITIK / TREUE / CHRONOLOGIE
///   / DOSIERUNGSSCHEMA / MARKIERUNGEN — each rule appears exactly once
const DEFAULT_SYSTEM_PROMPT: &str = include_str!("../../prompts/default_system_prompt.md");

/// Llama-specific system prompt. Currently identical to the default; kept
/// as a separate file so adjustments for Llama-3.1's behaviour (e.g.
/// stricter "do not hallucinate" wording) can be tuned independently of
/// gpt-oss without forking the default for everyone.
const LLAMA_SYSTEM_PROMPT: &str = include_str!("../../prompts/llama_system_prompt.md");

#[derive(Debug)]
pub struct LlmBackend {
    pub id: String,
    pub label: String,
    pub url: String,
    pub api_key: String,
    pub requires_api_key: bool,
    pub model_id: String,
    pub temperature: f32,
    pub top_p: Option<f32>,
    pub max_completion_tokens: u32,
    pub reasoning_effort: Option<String>,
    pub use_json_schema: bool,
    pub system_prompt: String,
    /// Optional second `system` message appended after `system_prompt`. Used
    /// to tell non-schema-aware models how to fill the schema's fields. `None`
    /// means the body is sent unchanged (gpt-oss-style backends).
    pub format_instruction: Option<String>,
    pub timeout_seconds: u64,
}

impl LlmBackend {
    pub fn is_available(&self) -> bool {
        !self.requires_api_key || !self.api_key.is_empty()
    }

    fn with_reasoning_effort(mut self, effort: &str) -> Self {
        self.reasoning_effort = Some(effort.into());
        self
    }

    fn with_use_json_schema(mut self, use_schema: bool) -> Self {
        self.use_json_schema = use_schema;
        self
    }

    fn with_system_prompt(mut self, prompt: &str) -> Self {
        self.system_prompt = prompt.into();
        self
    }

    fn with_temperature(mut self, temperature: f32) -> Self {
        self.temperature = temperature;
        self
    }

    fn with_top_p(mut self, top_p: f32) -> Self {
        self.top_p = Some(top_p);
        self
    }

    fn with_format_instruction(mut self, instruction: &str) -> Self {
        self.format_instruction = Some(instruction.into());
        self
    }

    fn with_max_completion_tokens(mut self, tokens: u32) -> Self {
        self.max_completion_tokens = tokens;
        self
    }
}

/// Ionos provider: shared endpoint, `IONOS_API_KEY` env var as bearer token.
fn ionos_backend(id: &str, label: &str, model_id: &str) -> LlmBackend {
    LlmBackend {
        id: id.into(),
        label: label.into(),
        url: "https://openai.inference.de-txl.ionos.com".into(),
        api_key: std::env::var("IONOS_API_KEY").unwrap_or_default(),
        requires_api_key: true,
        model_id: model_id.into(),
        temperature: DEFAULT_TEMPERATURE,
        top_p: None,
        max_completion_tokens: DEFAULT_MAX_COMPLETION_TOKENS,
        reasoning_effort: None,
        use_json_schema: DEFAULT_USE_JSON_SCHEMA,
        system_prompt: DEFAULT_SYSTEM_PROMPT.into(),
        format_instruction: None,
        timeout_seconds: DEFAULT_TIMEOUT_SECONDS,
    }
}

/// Returns the curated catalog of backends. Lazy-initialized once on first
/// access; the env-var read for `api_key` happens here, so missing
/// `IONOS_API_KEY` at startup yields backends with empty keys (filtered out
/// by `available_backends()`).
pub fn backends() -> &'static [LlmBackend] {
    static B: OnceLock<Vec<LlmBackend>> = OnceLock::new();
    B.get_or_init(|| {
        vec![
            // gpt-oss-120b: schema OFF on purpose. Reasoning tokens count
            // against `max_completion_tokens`; with `high` reasoning the
            // CoT alone can eat several thousand tokens, so the budget is
            // doubled to 8192 to leave room for the actual answer. Without
            // headroom the model returns `choices[0].message.content == null`
            // — the LlmError::Parse path. Pre-multi-backend (commit bf6464d)
            // gpt-oss ran without schema and worked fine; this restores
            // that body byte-for-byte except for the larger budget.
            ionos_backend("gpt_oss_120b", "GPT OSS 120b", "openai/gpt-oss-120b")
                .with_reasoning_effort("high")
                .with_max_completion_tokens(8192)
                .with_use_json_schema(false),
            ionos_backend(
                "llama_3_1_405b",
                "Llama 3.1 405B",
                "meta-llama/Meta-Llama-3.1-405B-Instruct-FP8",
            )
            .with_system_prompt(LLAMA_SYSTEM_PROMPT)
            .with_temperature(LLAMA_TEMPERATURE)
            .with_top_p(LLAMA_TOP_P)
            .with_format_instruction(LLAMA_FORMAT_INSTRUCTION),
        ]
    })
}

pub fn find_backend(id: &str) -> Option<&'static LlmBackend> {
    backends().iter().find(|b| b.id == id)
}

pub fn default_backend() -> &'static LlmBackend {
    &backends()[0]
}

pub fn available_backends() -> Vec<&'static LlmBackend> {
    backends().iter().filter(|b| b.is_available()).collect()
}

pub fn any_backend_available() -> bool {
    backends().iter().any(|b| b.is_available())
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashSet;

    /// The default backend is the first entry. Anything else would silently
    /// shift production traffic when entries are reordered.
    #[test]
    fn default_backend_is_first_in_catalog() {
        assert_eq!(default_backend().id, backends()[0].id);
    }

    #[test]
    fn default_backend_is_gpt_oss_120b() {
        assert_eq!(default_backend().id, "gpt_oss_120b");
        assert_eq!(default_backend().model_id, "openai/gpt-oss-120b");
    }

    #[test]
    fn find_backend_returns_known_id() {
        assert!(find_backend("gpt_oss_120b").is_some());
        assert!(find_backend("llama_3_1_405b").is_some());
    }

    #[test]
    fn find_backend_returns_none_for_unknown_id() {
        assert!(find_backend("does-not-exist").is_none());
        assert!(find_backend("").is_none());
    }

    /// Catalog ids are routed via stringly-typed lookup; duplicates would
    /// turn `find_backend` non-deterministic in the face of reorderings.
    #[test]
    fn catalog_ids_are_unique() {
        let mut seen: HashSet<&str> = HashSet::new();
        for b in backends() {
            assert!(seen.insert(b.id.as_str()), "duplicate id: {}", b.id);
        }
    }

    /// gpt-oss is a reasoning model, llama 3.1 is not. Asserting both keeps
    /// the conditional `reasoning_effort` field in `chat_once` correct.
    #[test]
    fn reasoning_effort_is_set_per_model_family() {
        assert_eq!(
            find_backend("gpt_oss_120b").unwrap().reasoning_effort,
            Some("high".to_string())
        );
        assert!(
            find_backend("llama_3_1_405b")
                .unwrap()
                .reasoning_effort
                .is_none()
        );
    }

    #[test]
    fn ionos_backends_share_endpoint_and_require_api_key() {
        for id in ["gpt_oss_120b", "llama_3_1_405b"] {
            let b = find_backend(id).unwrap();
            assert_eq!(b.url, "https://openai.inference.de-txl.ionos.com");
            assert!(b.requires_api_key, "{id} should require api key");
        }
    }

    #[test]
    fn is_available_requires_filled_api_key_when_required() {
        let with_key = LlmBackend {
            id: "x".into(),
            label: "x".into(),
            url: "https://x".into(),
            api_key: "k".into(),
            requires_api_key: true,
            model_id: "m".into(),
            temperature: 0.0,
            top_p: None,
            max_completion_tokens: 1,
            reasoning_effort: None,
            use_json_schema: false,
            system_prompt: String::new(),
            format_instruction: None,
            timeout_seconds: 1,
        };
        assert!(with_key.is_available());

        let no_key = LlmBackend {
            api_key: String::new(),
            ..with_key
        };
        assert!(!no_key.is_available());

        let no_auth_needed = LlmBackend {
            api_key: String::new(),
            requires_api_key: false,
            ..no_key
        };
        assert!(no_auth_needed.is_available());
    }

    /// Llama 3.1 405B FP8 needs sampling + a format hint to be usable on
    /// Ionos — verified empirically on 2026-05-03 (tests T1-T6 against the
    /// real endpoint). Pin those values so a future refactor cannot silently
    /// regress to the timeout-prone "greedy + no hint" config.
    #[test]
    fn llama_uses_sampling_and_format_instruction() {
        let b = find_backend("llama_3_1_405b").unwrap();
        assert!((b.temperature - 0.6).abs() < f32::EPSILON);
        assert_eq!(b.top_p, Some(0.9));
        assert!(
            b.format_instruction.is_some(),
            "Llama needs the format hint or it returns empty documents"
        );
        assert!(b.use_json_schema, "Llama needs schema-enforced stop");
    }

    /// gpt-oss-120b runs in plain-text mode (no `response_format`, no
    /// `top_p`, no second system message). Reason: schema-constrained
    /// decoding plus any non-zero `reasoning_effort` can exhaust
    /// `max_completion_tokens` on large bodies and return `content: null`.
    /// This pin guarantees the body sent to Ionos stays byte-identical to
    /// the pre-multi-backend path (commit bf6464d).
    #[test]
    fn gpt_oss_runs_in_plain_text_mode_no_schema_no_top_p_no_format_hint() {
        let b = find_backend("gpt_oss_120b").unwrap();
        assert!((b.temperature - 0.5).abs() < f32::EPSILON);
        assert!(b.top_p.is_none(), "gpt-oss must not send top_p");
        assert!(
            !b.use_json_schema,
            "gpt-oss must NOT use json_schema — collides with reasoning_effort budget"
        );
        assert!(
            b.format_instruction.is_none(),
            "gpt-oss must not get a second system message — body byte-stable"
        );
    }

    #[test]
    fn system_prompts_are_loaded() {
        for b in backends() {
            assert!(
                !b.system_prompt.trim().is_empty(),
                "{} has an empty system prompt",
                b.id
            );
        }
    }
}