newt-core 0.8.0

use crate::model_card::{ChatCompletionsCapability, ReasoningReplayScope};
use crate::role_profile::Cognition;

/// The cognition controls this backend can receive. Semantic intent remains in
/// the captured runtime even when this projection is absent.
#[must_use]
pub fn projected_cognition(
    cognition: Option<Cognition>,
    kind: crate::BackendKind,
    api: crate::OpenAiApi,
    chat_capability: ChatCompletionsCapability,
) -> Option<Cognition> {
    match (kind, api) {
        (crate::BackendKind::Openai, crate::OpenAiApi::Responses) => cognition,
        (crate::BackendKind::Openai, crate::OpenAiApi::ChatCompletions)
            if chat_capability.cognition == Some(true) =>
        {
            cognition
        }
        _ => None,
    }
}

/// The ONE output-allowance resolver (#2312). Precedence: an explicit operator
/// allowance (`[[model_tuning]] output_allowance`) wins; otherwise the
/// per-cognition table supplies it; otherwise `None`, and a wire that REQUIRES
/// a cap applies its own default (Anthropic: `NEWT_ANTHROPIC_MAX_TOKENS`, else
/// 8192). The Chat Completions policy sends the result as `max_tokens` where the
/// endpoint declared cognition projection. At dispatch, the independent generation
/// guard caps both Chat and Responses by this allowance and remaining context. The explicit value
/// never changes cognition, thinking or sampling.
pub(crate) fn resolve_output_allowance(
    explicit: Option<u32>,
    cognition: Option<Cognition>,
) -> Option<u32> {
    explicit.or_else(|| {
        cognition.map(|c| match c {
            Cognition::Zen => 2_048,
            Cognition::Rational => 4_096,
            Cognition::Thoughtful => 10_000,
            Cognition::Meticulous => 16_000,
        })
    })
}

/// Refuse an explicit output allowance no request could honour (#2312): `0`
/// permits no output, and one at or above the declared window leaves no input
/// room. Only the operator's explicit value is checked — defaults are not their
/// choice — and without a declared window only the `0` check can apply.
pub fn validate_output_allowance(
    explicit: Option<u32>,
    num_ctx: Option<u32>,
) -> anyhow::Result<()> {
    match (explicit, num_ctx) {
        (Some(0), _) => {
            anyhow::bail!("output_allowance 0 permits no output; omit it to use the default")
        }
        (Some(allowance), Some(window)) if allowance >= window => anyhow::bail!(
            "output_allowance {allowance} leaves no input room in the declared {window}-token \
             context window; lower --output-allowance or [[model_tuning]] output_allowance, \
             or declare a larger window"
        ),
        _ => Ok(()),
    }
}

/// Generation policy and output upper bound before prompt admission — the public,
/// read-only view of [`GenerationPolicy::resolve`] for display (#2567). It is
/// computed by the same policy resolver the dispatch loop uses. The final wire
/// cap also depends on the prompt size and configured generation maximum.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct ChatGenerationPreview {
    /// The `enable_thinking` value sent in `chat_template_kwargs`; `None` when
    /// the request sends no thinking selection at all (server default applies).
    pub enable_thinking: Option<bool>,
    pub temperature: Option<f64>,
    pub top_p: Option<f64>,
    /// Output upper bound; prompt admission can lower it further.
    pub max_output_tokens: Option<u32>,
    /// The output allowance reserved locally, sent or not.
    pub output_allowance: Option<u32>,
}

impl ChatGenerationPreview {
    /// Apply the independent guard only on the OpenAI wires it governs.
    /// Other backends retain their own generation policy.
    #[must_use]
    pub fn with_generation_bound(mut self, kind: crate::BackendKind) -> Self {
        if kind == crate::BackendKind::Openai {
            self.max_output_tokens =
                super::generation_bounds::output_cap(self.output_allowance, None, 0).ok();
        }
        self
    }
}

#[must_use]
pub fn preview_chat_generation(
    cognition: Option<Cognition>,
    output_allowance: Option<u32>,
    capability: ChatCompletionsCapability,
    reasoning_replay_scope: ReasoningReplayScope,
) -> ChatGenerationPreview {
    let policy = GenerationPolicy::resolve(
        cognition,
        output_allowance,
        capability,
        reasoning_replay_scope,
    );
    ChatGenerationPreview {
        enable_thinking: policy.thinking.filter(|_| policy.chat_template_kwargs),
        temperature: policy.temperature,
        top_p: policy.top_p,
        max_output_tokens: policy.max_output_tokens,
        output_allowance: policy.output_allowance,
    }
}

/// Backend-neutral generation choices resolved once per Chat Completions turn.
/// Request serialization projects these values only onto fields the endpoint
/// explicitly declared it accepts.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub(crate) struct GenerationPolicy {
    pub(crate) thinking: Option<bool>,
    pub(crate) reasoning_effort: Option<crate::model_card::ReasoningEffort>,
    /// Wire cap: projected as `max_tokens` only where the endpoint declared it.
    pub(crate) max_output_tokens: Option<u32>,
    /// The resolved allowance local admission reserves, sent or not.
    pub(crate) output_allowance: Option<u32>,
    pub(crate) temperature: Option<f64>,
    pub(crate) top_p: Option<f64>,
    pub(crate) parallel_tool_calls: Option<bool>,
    pub(crate) reasoning_replay_scope: ReasoningReplayScope,
    pub(crate) chat_template_kwargs: bool,
    pub(crate) one_bounded_reasoning_continuation: bool,
}

impl GenerationPolicy {
    /// Resolve the operator's cognition dial and explicit output allowance
    /// against endpoint capability data. The numeric table is the initial
    /// local-agent policy from the Nemotron qualification plan; no model
    /// display name participates in selection. An explicit allowance is
    /// reserved locally everywhere but sent only where cognition projects.
    #[must_use]
    pub(crate) fn resolve(
        cognition: Option<Cognition>,
        output_allowance: Option<u32>,
        capability: ChatCompletionsCapability,
        reasoning_replay_scope: ReasoningReplayScope,
    ) -> Self {
        let mut policy = Self {
            parallel_tool_calls: capability.parallel_tool_calls,
            reasoning_replay_scope,
            chat_template_kwargs: capability.chat_template_kwargs == Some(true),
            one_bounded_reasoning_continuation: capability.bounded_reasoning_continuation
                == Some(true),
            output_allowance,
            ..Self::default()
        };

        if capability.cognition != Some(true) {
            return policy;
        }
        // A projecting endpoint accepts the cap field, so an explicit allowance
        // is sent even with no dial; thinking and sampling still need one.
        policy.max_output_tokens = output_allowance;
        let Some(cognition) = cognition else {
            return policy;
        };

        let (thinking, temperature, top_p) = match cognition {
            Cognition::Zen => (false, 0.0, 1.0),
            Cognition::Rational => (true, 0.6, 0.95),
            Cognition::Thoughtful => (true, 0.6, 0.95),
            Cognition::Meticulous => (true, 0.6, 0.95),
        };
        policy.thinking = Some(thinking);
        policy.max_output_tokens = resolve_output_allowance(output_allowance, Some(cognition));
        policy.output_allowance = policy.max_output_tokens;
        policy.temperature = Some(temperature);
        policy.top_p = Some(top_p);
        policy
    }

    /// Resolve Responses fields and local reservation from the same captured
    /// declaration. Unsupported explicit effort fails before any dispatch.
    pub(crate) fn resolve_responses(
        cognition: Option<Cognition>,
        output_allowance: Option<u32>,
        capability: &crate::model_card::ResponsesCapability,
    ) -> anyhow::Result<Self> {
        let reasoning_effort = capability.resolve(cognition).map_err(anyhow::Error::msg)?;
        Ok(Self {
            reasoning_effort,
            output_allowance: resolve_output_allowance(output_allowance, cognition),
            ..Self::default()
        })
    }

    /// Add only explicitly resolved Chat Completions fields. An empty policy
    /// is a strict no-op for unknown compatible endpoints.
    pub(crate) fn apply_to_chat_completions_body(self, body: &mut serde_json::Value) {
        let Some(object) = body.as_object_mut() else {
            return;
        };
        if let Some(max_tokens) = self.max_output_tokens {
            object.insert("max_tokens".into(), serde_json::json!(max_tokens));
        }
        if let Some(temperature) = self.temperature {
            object.insert("temperature".into(), serde_json::json!(temperature));
        }
        if let Some(top_p) = self.top_p {
            object.insert("top_p".into(), serde_json::json!(top_p));
        }
        if object.contains_key("tools") {
            if let Some(parallel_tool_calls) = self.parallel_tool_calls {
                object.insert(
                    "parallel_tool_calls".into(),
                    serde_json::json!(parallel_tool_calls),
                );
            }
        }
        if self.chat_template_kwargs {
            if let Some(thinking) = self.thinking {
                object.insert(
                    "chat_template_kwargs".into(),
                    serde_json::json!({
                        "enable_thinking": thinking,
                        "truncate_history_thinking": true,
                    }),
                );
            }
        }
    }

    /// One automatic reasoning continuation is safe only when the endpoint
    /// opted in, current-turn reasoning can be replayed, and the ordinary
    /// round budget can pay for the extra request.
    #[must_use]
    pub(crate) fn allows_reasoning_continuation(
        self,
        already_attempted: bool,
        has_round_budget: bool,
    ) -> bool {
        !already_attempted
            && self.one_bounded_reasoning_continuation
            && self.reasoning_replay_scope != ReasoningReplayScope::Never
            && has_round_budget
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model_card::{ChatCompletionsCapability, ReasoningReplayScope};
    use crate::role_profile::Cognition;

    fn local_capability() -> ChatCompletionsCapability {
        ChatCompletionsCapability {
            cognition: Some(true),
            chat_template_kwargs: Some(true),
            parallel_tool_calls: Some(false),
            bounded_reasoning_continuation: Some(true),
        }
    }

    /// #2566/#2567: the preview reports exactly what is sent, including the
    /// two silent-off cases — no `cognition` projection, no cognition level.
    #[test]
    fn the_preview_reports_the_thinking_selection_the_request_sends() {
        let scope = ReasoningReplayScope::default();
        let on =
            preview_chat_generation(Some(Cognition::Thoughtful), None, local_capability(), scope);
        assert_eq!(on.enable_thinking, Some(true));
        assert_eq!((on.temperature, on.top_p), (Some(0.6), Some(0.95)));
        let zen = preview_chat_generation(Some(Cognition::Zen), None, local_capability(), scope);
        assert_eq!(zen.enable_thinking, Some(false));
        let no_level = preview_chat_generation(None, None, local_capability(), scope);
        assert_eq!(
            no_level.enable_thinking, None,
            "no cognition level: nothing sent"
        );
        let no_projection = ChatCompletionsCapability {
            cognition: None,
            ..local_capability()
        };
        let silent =
            preview_chat_generation(Some(Cognition::Thoughtful), None, no_projection, scope);
        assert_eq!(
            silent.enable_thinking, None,
            "kwargs flag alone sends nothing"
        );
        let no_kwargs = ChatCompletionsCapability {
            chat_template_kwargs: None,
            ..local_capability()
        };
        let sampled = preview_chat_generation(Some(Cognition::Thoughtful), None, no_kwargs, scope);
        assert_eq!(sampled.enable_thinking, None);
        assert_eq!(sampled.temperature, Some(0.6), "sampling still projects");
    }

    /// #2782: the preview cannot claim unlimited output without cognition.
    #[test]
    fn bounds_2782_preview_includes_the_independent_output_maximum() {
        let preview = preview_chat_generation(None, None, Default::default(), Default::default());
        assert_eq!(
            preview
                .with_generation_bound(crate::BackendKind::Ollama)
                .max_output_tokens,
            None
        );
        assert_eq!(
            preview
                .with_generation_bound(crate::BackendKind::Openai)
                .max_output_tokens,
            Some(16_384)
        );
        let limited =
            preview_chat_generation(None, Some(3000), Default::default(), Default::default());
        assert_eq!(
            limited
                .with_generation_bound(crate::BackendKind::Openai)
                .max_output_tokens,
            Some(3000)
        );
    }

    #[test]
    fn output_allowance_validation_refuses_only_what_no_request_could_honour() {
        assert!(validate_output_allowance(None, None).is_ok());
        assert!(validate_output_allowance(None, Some(32_768)).is_ok());
        assert!(
            validate_output_allowance(Some(1), None).is_ok(),
            "no window, no room check"
        );
        assert!(validate_output_allowance(Some(32_767), Some(32_768)).is_ok());
        assert!(validate_output_allowance(Some(0), None).is_err());
        assert!(validate_output_allowance(Some(32_768), Some(32_768)).is_err());
    }

    #[test]
    fn output_allowance_prefers_explicit_over_the_cognition_table() {
        // Default: the cognition table, the one number both wires use (Chat
        // sends it, Responses reserves it locally).
        for (cognition, allowance) in [
            (Cognition::Zen, 2_048),
            (Cognition::Rational, 4_096),
            (Cognition::Thoughtful, 10_000),
            (Cognition::Meticulous, 16_000),
        ] {
            assert_eq!(
                resolve_output_allowance(None, Some(cognition)),
                Some(allowance)
            );
        }
        // Opt-in: no cognition dial and no override reserves nothing.
        assert_eq!(resolve_output_allowance(None, None), None);
        // An explicit allowance wins, with or without a cognition dial.
        assert_eq!(
            resolve_output_allowance(Some(3_000), Some(Cognition::Meticulous)),
            Some(3_000)
        );
        assert_eq!(resolve_output_allowance(Some(3_000), None), Some(3_000));

        // `resolve` draws its cap from the same resolver, and the override
        // moves ONLY the allowance: thinking and sampling are unchanged.
        let resolve = |explicit| {
            GenerationPolicy::resolve(
                Some(Cognition::Meticulous),
                explicit,
                local_capability(),
                ReasoningReplayScope::Never,
            )
        };
        let overridden = resolve(Some(3_000));
        assert_eq!(overridden.max_output_tokens, Some(3_000));
        assert_eq!(overridden.output_allowance, Some(3_000));
        assert_eq!(
            GenerationPolicy {
                max_output_tokens: Some(16_000),
                output_allowance: Some(16_000),
                ..overridden
            },
            resolve(None)
        );
    }

    #[test]
    fn projecting_endpoint_sends_an_explicit_allowance_without_a_dial() {
        let policy = GenerationPolicy::resolve(
            None,
            Some(3_000),
            local_capability(),
            ReasoningReplayScope::Never,
        );
        assert_eq!(policy.max_output_tokens, Some(3_000));
        assert_eq!(policy.output_allowance, Some(3_000));
        assert_eq!(policy.thinking, None);
        assert_eq!(policy.temperature, None);
        assert_eq!(policy.top_p, None);
        // No dial and no override: nothing to send (defaults unchanged).
        let unset =
            GenerationPolicy::resolve(None, None, local_capability(), ReasoningReplayScope::Never);
        assert_eq!(unset.max_output_tokens, None);
    }

    #[test]
    fn unknown_endpoint_keeps_the_chat_request_policy_empty() {
        let policy = GenerationPolicy::resolve(
            Some(Cognition::Thoughtful),
            None,
            ChatCompletionsCapability::default(),
            ReasoningReplayScope::Never,
        );

        assert_eq!(policy, GenerationPolicy::default());
    }

    #[test]
    fn cognition_levels_resolve_to_the_initial_local_generation_table() {
        let cases = [
            (Cognition::Zen, false, 2_048, 0.0, 1.0),
            (Cognition::Rational, true, 4_096, 0.6, 0.95),
            (Cognition::Thoughtful, true, 10_000, 0.6, 0.95),
            (Cognition::Meticulous, true, 16_000, 0.6, 0.95),
        ];

        for (cognition, thinking, max_tokens, temperature, top_p) in cases {
            let policy = GenerationPolicy::resolve(
                Some(cognition),
                None,
                local_capability(),
                ReasoningReplayScope::CurrentUserTurn,
            );
            assert_eq!(policy.thinking, Some(thinking), "{cognition}");
            assert_eq!(policy.max_output_tokens, Some(max_tokens), "{cognition}");
            assert_eq!(policy.temperature, Some(temperature), "{cognition}");
            assert_eq!(policy.top_p, Some(top_p), "{cognition}");
            assert_eq!(policy.parallel_tool_calls, Some(false), "{cognition}");
            assert!(policy.chat_template_kwargs, "{cognition}");
            assert!(policy.one_bounded_reasoning_continuation, "{cognition}");
            assert_eq!(
                policy.reasoning_replay_scope,
                ReasoningReplayScope::CurrentUserTurn,
                "{cognition}"
            );
        }
    }

    #[test]
    fn endpoint_extensions_can_opt_in_without_enabling_cognition_projection() {
        // #2312: an explicit allowance is reserved locally, but this endpoint
        // declared no cap projection, so no `max_tokens` is sent.
        let policy = GenerationPolicy::resolve(
            Some(Cognition::Meticulous),
            Some(12_000),
            ChatCompletionsCapability {
                cognition: Some(false),
                chat_template_kwargs: Some(true),
                parallel_tool_calls: Some(false),
                bounded_reasoning_continuation: Some(true),
            },
            ReasoningReplayScope::CurrentUserTurn,
        );

        assert_eq!(policy.thinking, None);
        assert_eq!(policy.max_output_tokens, None);
        assert_eq!(policy.output_allowance, Some(12_000));
        assert_eq!(policy.temperature, None);
        assert_eq!(policy.top_p, None);
        assert_eq!(policy.parallel_tool_calls, Some(false));
        assert!(policy.chat_template_kwargs);
        assert!(policy.one_bounded_reasoning_continuation);
    }

    #[test]
    fn opted_in_policy_projects_to_chat_completions_fields() {
        let policy = GenerationPolicy::resolve(
            Some(Cognition::Thoughtful),
            None,
            local_capability(),
            ReasoningReplayScope::CurrentUserTurn,
        );
        let mut body = serde_json::json!({
            "model": "served-model",
            "messages": [],
            "tools": [],
            "tool_choice": "auto",
            "stream": false
        });

        policy.apply_to_chat_completions_body(&mut body);

        assert_eq!(body["max_tokens"], 10_000);
        assert_eq!(body["temperature"], 0.6);
        assert_eq!(body["top_p"], 0.95);
        assert_eq!(body["parallel_tool_calls"], false);
        assert_eq!(body["chat_template_kwargs"]["enable_thinking"], true);
        assert_eq!(
            body["chat_template_kwargs"]["truncate_history_thinking"],
            true
        );
    }

    #[test]
    fn default_policy_leaves_the_chat_completions_body_byte_for_byte_equal() {
        let mut body = serde_json::json!({
            "model": "strict-endpoint",
            "messages": [{"role": "user", "content": "hello"}],
            "stream": false
        });
        let before = body.clone();

        GenerationPolicy::default().apply_to_chat_completions_body(&mut body);

        assert_eq!(body, before);
    }

    #[test]
    fn tools_disabled_completion_omits_parallel_tool_calls() {
        let policy = GenerationPolicy::resolve(
            Some(Cognition::Rational),
            None,
            local_capability(),
            ReasoningReplayScope::CurrentUserTurn,
        );
        let mut body = serde_json::json!({
            "model": "served-model",
            "messages": [],
            "stream": false
        });

        policy.apply_to_chat_completions_body(&mut body);

        assert!(body.get("parallel_tool_calls").is_none());
        assert_eq!(body["max_tokens"], 4_096);
        assert_eq!(body["chat_template_kwargs"]["enable_thinking"], true);
    }

    #[test]
    fn bounded_reasoning_continuation_requires_capability_replay_and_round_budget() {
        let enabled = GenerationPolicy::resolve(
            None,
            None,
            local_capability(),
            ReasoningReplayScope::CurrentUserTurn,
        );
        assert!(enabled.allows_reasoning_continuation(false, true));
        assert!(!enabled.allows_reasoning_continuation(true, true));
        assert!(!enabled.allows_reasoning_continuation(false, false));

        let no_replay =
            GenerationPolicy::resolve(None, None, local_capability(), ReasoningReplayScope::Never);
        assert!(!no_replay.allows_reasoning_continuation(false, true));

        let no_capability = GenerationPolicy::resolve(
            None,
            None,
            ChatCompletionsCapability::default(),
            ReasoningReplayScope::CurrentUserTurn,
        );
        assert!(!no_capability.allows_reasoning_continuation(false, true));
    }
}