newt-inference 0.8.0

//! Local inference backends — the only backends compiled into the default
//! Newt binary. Cloud APIs live behind opt-in `ProviderPluginBackend`.
//!
//! Local requests default to a 120-second send/header deadline and maximum
//! idle gap between response chunks, plus a total read deadline
//! ([`STREAM_TOTAL_TIMEOUT_MULTIPLIER`] times the idle gap) so a stream that
//! keeps trickling bytes without ever finishing is still bounded. Callers
//! can override the send/idle timeout with `with_timeout` without replacing
//! an injected HTTP client. OpenAI-compatible and native Ollama completion
//! requests stream.

use std::future::Future;
use std::time::Duration;
use tokio::time::Instant;

use async_trait::async_trait;
use newt_core::router::Tier;

use crate::backend::{ChatReply, ChatRequest, InferenceBackend};
use crate::retry::{with_backoff, RetryPolicy};

const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);

/// A stream whose idle gap between chunks never exceeds `idle_timeout` can
/// still run forever: a model-swap proxy that drip-feeds one byte just under
/// the idle window (or a cold-loaded model producing a token every ~90s)
/// resets the per-chunk timer on every arrival and is never caught. The total
/// read deadline is `idle_timeout * this multiplier`, wide enough to cover a
/// slow-but-progressing generation while still bounding a stream that never
/// finishes. Diagnosed from a `newt` session stuck 5689s past its stated
/// "120s idle timeout" against a DGX model-swap backend.
const STREAM_TOTAL_TIMEOUT_MULTIPLIER: u32 = 5;

#[derive(Debug)]
struct ObservedProviderError(String);

impl std::fmt::Display for ObservedProviderError {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl std::error::Error for ObservedProviderError {}

fn is_observed_provider_error(error: &anyhow::Error) -> bool {
    newt_core::agentic::openai_sse::is_provider_error(error)
        || error
            .chain()
            .any(|cause| cause.is::<ObservedProviderError>())
}

fn local_client(read_timeout: Duration) -> reqwest::Client {
    reqwest::Client::builder()
        .read_timeout(read_timeout)
        .build()
        .expect("build local inference client")
}

// An outer deadline preserves a caller-owned client's stricter timeouts,
// connection pool, headers, proxy, and TLS configuration. A RequestBuilder
// timeout would instead override its client's existing total deadline.
async fn bounded_request<T>(
    deadline: Instant,
    request: impl Future<Output = anyhow::Result<T>>,
) -> anyhow::Result<T> {
    tokio::time::timeout_at(deadline, request)
        .await
        .map_err(|_| anyhow::anyhow!("local inference request failed: deadline elapsed"))?
}

#[cfg(test)]
tokio::task_local! {
    /// Test synchronization after client consumption, never after a server write.
    static RESPONSE_BYTES_READ: std::sync::Arc<std::sync::atomic::AtomicUsize>;
}

async fn decode_http_reply(
    mut response: reqwest::Response,
    idle_timeout: Duration,
    total_deadline: Instant,
    backend: &str,
    decode: impl FnOnce(&[u8]) -> anyhow::Result<serde_json::Value>,
) -> anyhow::Result<serde_json::Value> {
    let status = response.status();
    // Own the evidence outside each cancellable read. An idle timeout may
    // close the response after an error frame arrived but before its EOF.
    let mut bytes = Vec::new();
    let read_error = loop {
        let remaining = total_deadline.saturating_duration_since(Instant::now());
        if remaining.is_zero() {
            break Some(anyhow::anyhow!(
                "total read deadline elapsed (each chunk arrived within the \
                 {idle_timeout:?} idle window, but the stream never finished)"
            ));
        }
        match tokio::time::timeout(idle_timeout.min(remaining), response.chunk()).await {
            Ok(Ok(Some(chunk))) => {
                bytes.extend_from_slice(&chunk);
                #[cfg(test)]
                let _ = RESPONSE_BYTES_READ.try_with(|count| {
                    count.fetch_add(chunk.len(), std::sync::atomic::Ordering::SeqCst);
                });
            }
            Ok(Ok(None)) => break None,
            Ok(Err(error)) => break Some(anyhow::Error::from(error)),
            Err(error) => break Some(anyhow::Error::from(error)),
        }
    };
    if !status.is_success() {
        let detail = String::from_utf8_lossy(&bytes);
        if let Some(error) = read_error {
            anyhow::bail!(
                "{backend} returned {status}: {detail}; response body read failure: {error}"
            );
        }
        anyhow::bail!("{backend} returned {status}: {detail}");
    }
    let decoded = decode(&bytes);
    if let Some(read_error) = read_error {
        if let Err(error) = decoded {
            if is_observed_provider_error(&error) {
                let message = format!(
                    "{backend} response: {error}; response body read failure: {read_error}"
                );
                return Err(error.context(message));
            }
        }
        anyhow::bail!("{backend} request failed reading response: {read_error}");
    }
    decoded
}

fn decode_ollama_stream(bytes: &[u8]) -> anyhow::Result<serde_json::Value> {
    let body = std::str::from_utf8(bytes)
        .map_err(|error| anyhow::anyhow!("invalid Ollama stream UTF-8: {error}"))?;
    let mut content = String::new();
    let mut terminal = None;

    for (index, line) in body.lines().enumerate() {
        if line.trim().is_empty() {
            continue;
        }
        let frame: serde_json::Value = serde_json::from_str(line).map_err(|error| {
            anyhow::anyhow!("invalid Ollama stream record {}: {error}", index + 1)
        })?;
        anyhow::ensure!(
            terminal.is_none(),
            "Ollama stream continued after its terminal frame"
        );
        anyhow::ensure!(frame.is_object(), "Ollama stream frame is not an object");
        if let Some(error) = frame.get("error").filter(|error| !error.is_null()) {
            let detail = error
                .as_str()
                .map(str::to_owned)
                .unwrap_or_else(|| error.to_string());
            return Err(ObservedProviderError(format!("Ollama response error: {detail}")).into());
        }
        let message = match frame.get("message") {
            Some(serde_json::Value::Object(message)) => Some(message),
            None | Some(serde_json::Value::Null) => None,
            Some(_) => anyhow::bail!("Ollama stream message is not an object"),
        };
        match message.and_then(|message| message.get("content")) {
            Some(serde_json::Value::String(part)) => content.push_str(part),
            None | Some(serde_json::Value::Null) => {}
            Some(_) => anyhow::bail!("Ollama stream content is not a string"),
        }
        match frame.get("done").and_then(serde_json::Value::as_bool) {
            Some(true) => terminal = Some(frame),
            Some(false) => {}
            None => anyhow::bail!("Ollama stream frame has no boolean done marker"),
        }
    }

    let mut terminal =
        terminal.ok_or_else(|| anyhow::anyhow!("Ollama stream ended before done"))?;
    let object = terminal
        .as_object_mut()
        .expect("terminal stream frame was checked as an object");
    let message = object
        .entry("message")
        .or_insert_with(|| serde_json::json!({}))
        .as_object_mut()
        .ok_or_else(|| anyhow::anyhow!("Ollama stream message is not an object"))?;
    message.insert("content".to_owned(), serde_json::Value::String(content));
    Ok(terminal)
}

#[derive(Debug)]
pub struct LocalOllamaBackend {
    endpoint: String,
    model: String,
    client: reqwest::Client,
    managed_client: bool,
    timeout: Duration,
    /// Optional bearer token sent as `Authorization: Bearer <token>` —
    /// Ollama Cloud (`https://ollama.com`) requires one; LAN Ollama needs
    /// none (the default) and ignores an unexpected header.
    api_key: Option<String>,
    retry: RetryPolicy,
}

impl LocalOllamaBackend {
    pub fn new(endpoint: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            endpoint: endpoint.into(),
            model: model.into(),
            client: local_client(DEFAULT_REQUEST_TIMEOUT),
            managed_client: true,
            timeout: DEFAULT_REQUEST_TIMEOUT,
            api_key: None,
            retry: RetryPolicy::from_env(),
        }
    }

    /// Use a caller-owned HTTP client, preserving its connection pool across
    /// backend values constructed for separate requests.
    pub fn with_client(mut self, client: reqwest::Client) -> Self {
        self.client = client;
        self.managed_client = false;
        self
    }

    /// Override the retry/backoff policy (defaults to [`RetryPolicy::from_env`]).
    /// Used by tests to inject a zero-delay policy; production callers can tune
    /// it via the `NEWT_HTTP_*` env vars instead.
    pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
        self.retry = policy;
        self
    }

    /// Return the configured endpoint URL.
    pub fn endpoint(&self) -> &str {
        &self.endpoint
    }

    /// Attach a bearer token, sent as `Authorization: Bearer <token>` on
    /// every request (same contract as [`LocalVllmBackend::with_api_key`]).
    /// `None`/empty leaves the backend unauthenticated.
    pub fn with_api_key(mut self, api_key: impl Into<Option<String>>) -> Self {
        self.api_key = api_key.into().filter(|k| !k.is_empty());
        self
    }

    /// Build from a [`BackendConfig`](newt_core::BackendConfig): endpoint,
    /// model, and bearer auth resolved from `api_key_env` / `api_key_file`
    /// (parity with [`LocalVllmBackend::from_config`]).
    pub fn from_config(cfg: &newt_core::BackendConfig) -> Self {
        let model = cfg.effective_model().unwrap_or_default().to_string();
        Self::new(cfg.endpoint.clone(), model).with_api_key(cfg.resolve_api_key())
    }

    /// Apply bearer auth to a request builder when a token is configured.
    fn authed(&self, rb: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
        match &self.api_key {
            Some(key) => rb.bearer_auth(key),
            None => rb,
        }
    }

    /// Override the send/header deadline and idle read timeout per attempt
    /// (default: 120 seconds). Retains an injected client and any stricter
    /// timeout that client carries.
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        if self.managed_client {
            self.client = local_client(timeout);
        }
        self
    }

    /// Try endpoints in order, return the first reachable one.
    /// Reachability = GET /api/tags returns 2xx within 500ms.
    /// Checks `OLLAMA_HOST` env var first, then the default endpoint list.
    pub async fn discover(model: &str) -> anyhow::Result<Self> {
        let env_host = std::env::var("OLLAMA_HOST").ok();
        Self::discover_inner(model, env_host.as_deref(), &Self::default_endpoints()).await
    }

    /// Like [`discover`](Self::discover), but with a caller-supplied candidate
    /// list instead of the built-in defaults. `OLLAMA_HOST` is still checked
    /// first.
    pub async fn discover_with_candidates(
        model: &str,
        candidates: &[String],
    ) -> anyhow::Result<Self> {
        let env_host = std::env::var("OLLAMA_HOST").ok();
        Self::discover_inner(model, env_host.as_deref(), candidates).await
    }

    /// Like [`discover`](Self::discover), but with an explicit env-host
    /// override and candidate list. Useful for testing without mutating
    /// process-global environment variables.
    pub async fn discover_with_env(
        model: &str,
        env_host: Option<&str>,
        candidates: &[String],
    ) -> anyhow::Result<Self> {
        Self::discover_inner(model, env_host, candidates).await
    }

    async fn discover_inner(
        model: &str,
        env_host: Option<&str>,
        candidates: &[String],
    ) -> anyhow::Result<Self> {
        // If OLLAMA_HOST is explicitly set, use it VERBATIM — no probe,
        // no fallback. User intent overrides discovery. This eliminates
        // the silent-fallthrough foot-gun where a stale env var (or a
        // mocked endpoint missing /api/tags) causes discover() to hit
        // a different Ollama than the user asked for.
        //
        // Use `discover_strict` if you want the env host to be probed.
        if let Some(host) = env_host {
            tracing::info!(
                endpoint = %host,
                "Ollama endpoint chosen via OLLAMA_HOST (verbatim, not probed)"
            );
            return Ok(Self::new(host, model));
        }

        let probe_client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_millis(500))
            .build()?;

        for endpoint in candidates {
            if Self::probe(&probe_client, endpoint).await {
                tracing::info!(endpoint = %endpoint, "Ollama endpoint chosen by discovery probe");
                return Ok(Self::new(endpoint, model));
            }
        }

        anyhow::bail!(
            "no reachable Ollama endpoint found (tried {} candidates)",
            candidates.len()
        )
    }

    /// Like [`discover`](Self::discover) but requires successful
    /// probing — no silent fallthrough. Even if `OLLAMA_HOST` is
    /// set, the probe must succeed. Use this in tests and CI to
    /// assert a specific endpoint is reachable rather than just
    /// trusted by env-var contract.
    pub async fn discover_strict(model: &str) -> anyhow::Result<Self> {
        let env_host = std::env::var("OLLAMA_HOST").ok();
        Self::discover_strict_with_env(model, env_host.as_deref(), &Self::default_endpoints()).await
    }

    /// Like [`discover_strict`](Self::discover_strict), but with an
    /// explicit env-host override and candidate list. Useful for
    /// testing without mutating process-global environment variables.
    pub async fn discover_strict_with_env(
        model: &str,
        env_host: Option<&str>,
        candidates: &[String],
    ) -> anyhow::Result<Self> {
        // Build the full candidate list with the env-host (if any) at
        // the front. We probe every candidate, including the env-host.
        let all_candidates: Vec<&str> = env_host
            .into_iter()
            .chain(candidates.iter().map(String::as_str))
            .collect();

        let probe_client = reqwest::Client::builder()
            .timeout(std::time::Duration::from_millis(500))
            .build()?;

        for endpoint in &all_candidates {
            if Self::probe(&probe_client, endpoint).await {
                tracing::info!(endpoint = %endpoint, "Ollama endpoint chosen (strict)");
                return Ok(Self::new(*endpoint, model));
            }
        }

        anyhow::bail!(
            "discover_strict: no reachable Ollama endpoint (tried {} candidates)",
            all_candidates.len()
        )
    }

    /// The built-in fallback endpoint list for [`discover`](Self::discover).
    pub fn default_endpoints() -> Vec<String> {
        vec![
            "http://ollama-proxy.inference.svc.cluster.local:11434".to_string(),
            "http://REDACTED-HOST:11434".to_string(),
            "http://REDACTED-HOST:11434".to_string(),
            "http://REDACTED-HOST:11434".to_string(),
            "http://127.0.0.1:11434".to_string(),
        ]
    }

    async fn probe(client: &reqwest::Client, endpoint: &str) -> bool {
        let url = format!("{}/api/tags", endpoint.trim_end_matches('/'));
        match client.get(&url).send().await {
            Ok(resp) => resp.status().is_success(),
            Err(_) => false,
        }
    }

    /// Single HTTP attempt — no retries. Returns a structured error that
    /// [`crate::retry::classify`] can classify for the backoff loop.
    async fn try_complete(&self, req: &ChatRequest) -> anyhow::Result<ChatReply> {
        let deadline = Instant::now() + self.timeout;
        let body = serde_json::json!({
            "model": self.model,
            "messages": req.messages.iter().map(|m| {
                serde_json::json!({ "role": &m.role, "content": &m.content })
            }).collect::<Vec<_>>(),
            "stream": true,
            "options": req.max_tokens.map(|t| serde_json::json!({ "num_predict": t })),
        });

        let url = format!("{}/api/chat", self.endpoint.trim_end_matches('/'));
        let resp = bounded_request(deadline, async {
            self.authed(self.client.post(&url).json(&body))
                .send()
                .await
                .map_err(|e| anyhow::anyhow!("Ollama request failed: {e}"))
        })
        .await?;

        let total_deadline = Instant::now() + self.timeout * STREAM_TOTAL_TIMEOUT_MULTIPLIER;
        let json = decode_http_reply(
            resp,
            self.timeout,
            total_deadline,
            "Ollama",
            decode_ollama_stream,
        )
        .await?;
        // #385: strip inline <think>…</think> reasoning from the content.
        let (content, _reasoning) =
            newt_core::split_reasoning(json["message"]["content"].as_str().unwrap_or(""));

        // Extract token counts from Ollama's response if present.
        let usage = {
            let input = json["prompt_eval_count"].as_u64().map(|n| n as u32);
            let output = json["eval_count"].as_u64().map(|n| n as u32);
            input.zip(output).map(|(i, o)| newt_core::TokenUsage {
                input_tokens: i,
                output_tokens: o,
            })
        };

        Ok(ChatReply {
            content,
            model_id: self.model.clone(),
            usage,
        })
    }
}

#[async_trait]
impl InferenceBackend for LocalOllamaBackend {
    fn name(&self) -> &str {
        "ollama-local"
    }

    fn model_id(&self) -> &str {
        &self.model
    }

    fn supports_tier(&self, _tier: Tier) -> bool {
        true
    }

    fn endpoint(&self) -> Option<&str> {
        Some(&self.endpoint)
    }

    async fn complete(&self, req: ChatRequest) -> anyhow::Result<ChatReply> {
        with_backoff(&self.retry, || self.try_complete(&req)).await
    }
}

/// Metadata for a single model exposed by a vLLM server's `/v1/models`
/// endpoint. The shape mirrors the OpenAI-compatible response — we
/// only surface the fields the rest of Newt cares about today.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelInfo {
    pub id: String,
}

/// A backend that speaks the OpenAI-compatible HTTP API exposed by a
/// local vLLM server (`POST /v1/chat/completions`, `GET /v1/models`).
///
/// vLLM endpoints are explicit — unlike Ollama, vLLM has no canonical
/// default port, so we deliberately skip endpoint auto-discovery here.
/// Callers must supply the endpoint via config or CLI flag.
#[derive(Debug)]
pub struct LocalVllmBackend {
    endpoint: String,
    model: String,
    client: reqwest::Client,
    managed_client: bool,
    timeout: Duration,
    /// Optional bearer token sent as `Authorization: Bearer <token>`.
    /// `None` for unauthenticated local servers (the default); `Some`
    /// for hosted OpenAI-compatible endpoints that require an API key.
    api_key: Option<String>,
    retry: RetryPolicy,
}

impl LocalVllmBackend {
    pub fn new(endpoint: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            endpoint: endpoint.into(),
            model: model.into(),
            client: local_client(DEFAULT_REQUEST_TIMEOUT),
            managed_client: true,
            timeout: DEFAULT_REQUEST_TIMEOUT,
            api_key: None,
            retry: RetryPolicy::from_env(),
        }
    }

    /// Use a caller-owned HTTP client, preserving its connection pool across
    /// backend values constructed for separate requests.
    pub fn with_client(mut self, client: reqwest::Client) -> Self {
        self.client = client;
        self.managed_client = false;
        self
    }

    /// Override the retry/backoff policy (defaults to [`RetryPolicy::from_env`]).
    /// Used by tests to inject a zero-delay policy; production callers can tune
    /// it via the `NEWT_HTTP_*` env vars instead.
    pub fn with_retry_policy(mut self, policy: RetryPolicy) -> Self {
        self.retry = policy;
        self
    }

    /// Return the configured endpoint URL.
    pub fn endpoint(&self) -> &str {
        &self.endpoint
    }

    /// Attach a bearer token, sent as `Authorization: Bearer <token>` on
    /// every request. A `None` argument (or an empty token) leaves the
    /// backend unauthenticated, so callers can pass a resolved
    /// `Option<String>` straight through.
    pub fn with_api_key(mut self, api_key: impl Into<Option<String>>) -> Self {
        self.api_key = api_key.into().filter(|k| !k.is_empty());
        self
    }

    /// Build from a [`BackendConfig`](newt_core::BackendConfig), wiring up
    /// the endpoint, model, and bearer auth resolved from the config's
    /// `api_key_env` / `api_key_file`. Used by the worker to construct an
    /// authenticated OpenAI-compatible backend from `~/.newt/config.toml`.
    pub fn from_config(cfg: &newt_core::BackendConfig) -> Self {
        // A model-less backend (server dictates, #1128) yields "" here until
        // Phase B's adopt() supplies the served model; callers today always
        // pass configs with a declared model.
        let model = cfg.effective_model().unwrap_or_default().to_string();
        Self::new(cfg.endpoint.clone(), model).with_api_key(cfg.resolve_api_key())
    }

    /// Apply bearer auth to a request builder when a token is configured.
    fn authed(&self, rb: reqwest::RequestBuilder) -> reqwest::RequestBuilder {
        match &self.api_key {
            Some(key) => rb.bearer_auth(key),
            None => rb,
        }
    }

    /// Override the send/header deadline and idle read timeout per attempt
    /// (default: 120 seconds). Retains an injected client and any stricter
    /// timeout that client carries.
    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = timeout;
        if self.managed_client {
            self.client = local_client(timeout);
        }
        self
    }

    /// Single HTTP attempt — no retries. Returns a structured error that
    /// [`crate::retry::classify`] can classify for the backoff loop.
    async fn try_complete(&self, req: &ChatRequest) -> anyhow::Result<ChatReply> {
        let deadline = Instant::now() + self.timeout;
        let mut body = serde_json::json!({
            "model": self.model,
            "messages": req.messages.iter().map(|m| {
                serde_json::json!({ "role": &m.role, "content": &m.content })
            }).collect::<Vec<_>>(),
            "stream": true,
            "stream_options": {"include_usage": true},
        });
        if let Some(max) = req.max_tokens {
            body["max_tokens"] = serde_json::json!(max);
        }

        let url = format!(
            "{}/v1/chat/completions",
            self.endpoint.trim_end_matches('/')
        );
        let resp = bounded_request(deadline, async {
            self.authed(self.client.post(&url).json(&body))
                .send()
                .await
                .map_err(|e| anyhow::anyhow!("vLLM request failed: {e}"))
        })
        .await?;

        let total_deadline = Instant::now() + self.timeout * STREAM_TOTAL_TIMEOUT_MULTIPLIER;
        let json = decode_http_reply(
            resp,
            self.timeout,
            total_deadline,
            "vLLM",
            newt_core::agentic::openai_sse::decode_response,
        )
        .await?;
        // OpenAI-compatible: choices[0].message.content
        // #385: strip inline <think>…</think> reasoning from the content.
        let (content, _reasoning) = newt_core::split_reasoning(
            json["choices"][0]["message"]["content"]
                .as_str()
                .unwrap_or(""),
        );
        // Prefer the model echoed back by the server (helps callers
        // distinguish aliases) but fall back to the configured id.
        let model_id = json["model"]
            .as_str()
            .map(|s| s.to_string())
            .unwrap_or_else(|| self.model.clone());

        // Extract exact token counts from the OpenAI-compatible `usage`
        // object when present. Preferring the server's reported counts
        // over heuristic estimation directly addresses #247 (token
        // estimator error) on the vLLM path.
        let usage = {
            let input = json["usage"]["prompt_tokens"].as_u64().map(|n| n as u32);
            let output = json["usage"]["completion_tokens"]
                .as_u64()
                .map(|n| n as u32);
            input.zip(output).map(|(i, o)| newt_core::TokenUsage {
                input_tokens: i,
                output_tokens: o,
            })
        };

        Ok(ChatReply {
            content,
            model_id,
            usage,
        })
    }

    /// List the models the vLLM server is currently serving.
    ///
    /// Mirrors the OpenAI `GET /v1/models` response shape:
    ///
    /// ```json
    /// { "data": [{ "id": "llama3.1:8b", "object": "model" }, ...] }
    /// ```
    ///
    /// Used by `newt doctor` (follow-up) to probe vLLM endpoints.
    pub async fn list_models(&self) -> anyhow::Result<Vec<ModelInfo>> {
        bounded_request(Instant::now() + self.timeout, self.try_list_models()).await
    }

    async fn try_list_models(&self) -> anyhow::Result<Vec<ModelInfo>> {
        let url = format!("{}/v1/models", self.endpoint.trim_end_matches('/'));
        let resp = self
            .authed(self.client.get(&url))
            .send()
            .await
            .map_err(|e| anyhow::anyhow!("vLLM list_models request failed: {e}"))?;

        if !resp.status().is_success() {
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            anyhow::bail!("vLLM list_models returned {status}: {text}");
        }

        let json: serde_json::Value = resp.json().await?;
        let data = json["data"]
            .as_array()
            .ok_or_else(|| anyhow::anyhow!("vLLM list_models response missing 'data' array"))?;

        let models = data
            .iter()
            .filter_map(|entry| {
                entry["id"]
                    .as_str()
                    .map(|id| ModelInfo { id: id.to_string() })
            })
            .collect();

        Ok(models)
    }
}

#[async_trait]
impl InferenceBackend for LocalVllmBackend {
    fn name(&self) -> &str {
        "vllm-local"
    }

    fn model_id(&self) -> &str {
        &self.model
    }

    fn supports_tier(&self, _tier: Tier) -> bool {
        true
    }

    fn endpoint(&self) -> Option<&str> {
        Some(&self.endpoint)
    }

    async fn complete(&self, req: ChatRequest) -> anyhow::Result<ChatReply> {
        with_backoff(&self.retry, || self.try_complete(&req)).await
    }
}

#[cfg(test)]
#[path = "local_timing_tests.rs"]
mod timing_tests;