open-standards/eoc/eoc-rs/crates/eoc-vendor-api/src/mistral.rs

//! Mistral chat-completions backend (OpenAI-compatible schema).

use async_trait::async_trait;
use eoc_core::{JouleCost, JouleSource, Query, Response, Stage};
use eoc_neural::NeuralBackend;
use tracing::field;

use crate::auth::Auth;
use crate::config::VendorConfig;
use crate::error::VendorResult;
use crate::openai_compat;

/// Default Mistral endpoint.
pub const DEFAULT_ENDPOINT: &str = "https://api.mistral.ai/v1/chat/completions";

/// Mistral chat backend.
pub struct MistralBackend {
    client: reqwest::Client,
    auth: Auth,
    model: String,
    stream: bool,
    config: VendorConfig,
}

impl MistralBackend {
    /// Construct with an API key and target model.
    pub fn new(api_key: impl Into<String>, model: impl Into<String>) -> Self {
        Self {
            client: reqwest::Client::new(),
            auth: Auth::Bearer(api_key.into()),
            model: model.into(),
            stream: true,
            config: VendorConfig::new(),
        }
    }

    /// Override the [`VendorConfig`].
    pub fn with_config(mut self, config: VendorConfig) -> Self {
        self.config = config;
        self
    }

    /// Disable SSE streaming.
    pub fn without_stream(mut self) -> Self {
        self.stream = false;
        self
    }

    fn endpoint(&self) -> &str {
        self.config
            .endpoint
            .as_deref()
            .unwrap_or(DEFAULT_ENDPOINT)
    }

    async fn try_infer(&self, q: &Query) -> VendorResult<Response> {
        tracing::debug!(
            target: "mistral.infer",
            model = %self.model,
            api_key = field::Empty,
            "dispatching mistral inference"
        );
        openai_compat::execute(
            &self.client,
            self.endpoint(),
            &self.auth,
            &self.model,
            q,
            self.stream,
            &self.config,
        )
        .await
    }
}

#[async_trait]
impl NeuralBackend for MistralBackend {
    async fn infer(&self, q: &Query) -> Response {
        match self.try_infer(q).await {
            Ok(r) => r,
            Err(e) => Response::new(
                q.id,
                format!("[mistral-error: {e}]"),
                Stage::Neural,
                JouleCost { microjoules: 0, source: JouleSource::Estimated },
            ),
        }
    }
}