car_inference/

lib.rs

//! # car-inference
//!
//! Local model inference for the Common Agent Runtime.
//!
//! Provides on-device inference using Candle with automatic hardware detection:
//! - **macOS**: Metal (Apple Silicon GPU)
//! - **Linux**: CUDA (NVIDIA GPU) or CPU fallback
//!
//! Ships with Qwen3 models downloaded on first use from HuggingFace.
//! Supports remote API models (OpenAI, Anthropic, Google) via the same schema.
//!
//! ## Architecture
//!
//! Models are first-class typed resources described by `ModelSchema` (analogous
//! to `ToolSchema`). The `UnifiedRegistry` holds local and remote models.
//! The `AdaptiveRouter` selects the best model using a three-phase strategy:
//! filter → score → explore. The `OutcomeTracker` learns from results to
//! improve routing over time.
//!
//! ## Dual purpose
//!
//! 1. **Internal** — powers skill learning/repair, semantic memory, policy evaluation
//! 2. **Service** — exposes `infer`, `embed`, `classify` as built-in CAR tools

pub mod adaptive_router;
pub mod backend;
pub mod hardware;
pub mod models;
pub mod outcome;
pub mod registry;
pub mod remote;
pub mod router;
pub mod schema;
pub mod service;
pub mod tasks;

use std::sync::Arc;
use std::time::Instant;

use thiserror::Error;
use tokio::sync::RwLock;
use tracing::debug;

// --- New types ---
pub use adaptive_router::{AdaptiveRouter, AdaptiveRoutingDecision, RoutingConfig, RoutingStrategy};
pub use outcome::{
    CodeOutcome, InferenceOutcome, InferenceTask, InferredOutcome, ModelProfile, OutcomeTracker,
};
pub use registry::{ModelFilter, ModelInfo, UnifiedRegistry};
pub use remote::RemoteBackend;
pub use schema::{ApiProtocol, CostModel, ModelCapability, ModelSchema, ModelSource, PerformanceEnvelope};

// --- Legacy re-exports (kept for backward compatibility) ---
pub use adaptive_router::TaskComplexity;
pub use backend::CandleBackend;
pub use backend::EmbeddingBackend;
pub use hardware::HardwareInfo;
pub use models::{ModelRegistry, ModelRole};
pub use router::{ModelRouter, RoutingDecision};
pub use tasks::{ClassifyRequest, ClassifyResult, EmbedRequest, GenerateParams, GenerateRequest};

#[derive(Error, Debug)]
pub enum InferenceError {
    #[error("model not found: {0}")]
    ModelNotFound(String),

    #[error("model download failed: {0}")]
    DownloadFailed(String),

    #[error("inference failed: {0}")]
    InferenceFailed(String),

    #[error("tokenization error: {0}")]
    TokenizationError(String),

    #[error("device error: {0}")]
    DeviceError(String),

    #[error("io error: {0}")]
    Io(#[from] std::io::Error),
}

/// Which device to run inference on.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Device {
    Cpu,
    Metal,
    Cuda(usize), // device ordinal
}

impl Device {
    /// Auto-detect the best available device for this platform.
    pub fn auto() -> Self {
        #[cfg(feature = "metal")]
        {
            return Device::Metal;
        }
        #[cfg(feature = "cuda")]
        {
            return Device::Cuda(0);
        }
        #[cfg(not(any(feature = "metal", feature = "cuda")))]
        {
            Device::Cpu
        }
    }
}

/// Configuration for the inference engine.
#[derive(Debug, Clone)]
pub struct InferenceConfig {
    /// Where to store downloaded models. Defaults to ~/.car/models/
    pub models_dir: std::path::PathBuf,
    /// Device override. None = auto-detect.
    pub device: Option<Device>,
    /// Default model for generation tasks.
    pub generation_model: String,
    /// Default model for embedding tasks.
    pub embedding_model: String,
    /// Default model for classification tasks.
    pub classification_model: String,
}

impl Default for InferenceConfig {
    fn default() -> Self {
        let models_dir = dirs_next()
            .unwrap_or_else(|| std::path::PathBuf::from("."))
            .join(".car")
            .join("models");

        let hw = HardwareInfo::detect();

        Self {
            models_dir,
            device: None,
            generation_model: hw.recommended_model,
            embedding_model: "Qwen3-Embedding-0.6B".to_string(),
            classification_model: "Qwen3-0.6B".to_string(),
        }
    }
}

fn dirs_next() -> Option<std::path::PathBuf> {
    std::env::var("HOME")
        .ok()
        .map(std::path::PathBuf::from)
}

/// Result of an inference call, including trace ID for outcome tracking.
#[derive(Debug, Clone)]
pub struct InferenceResult {
    /// The generated text.
    pub text: String,
    /// Trace ID for reporting outcomes back to the tracker.
    pub trace_id: String,
    /// Which model was used.
    pub model_used: String,
    /// Wall-clock latency in ms.
    pub latency_ms: u64,
}

/// The main inference engine. Thread-safe, lazily loads models.
///
/// Now includes the unified registry, adaptive router, and outcome tracker
/// for schema-driven model selection with learned performance profiles.
pub struct InferenceEngine {
    pub config: InferenceConfig,
    /// Unified model registry (local + remote).
    pub unified_registry: UnifiedRegistry,
    /// Adaptive router with three-phase selection.
    pub adaptive_router: AdaptiveRouter,
    /// Outcome tracker for learning from results.
    pub outcome_tracker: Arc<RwLock<OutcomeTracker>>,
    /// HTTP client for remote API models.
    remote_backend: RemoteBackend,
    // Legacy fields kept for backward compatibility
    pub registry: models::ModelRegistry,
    pub router: ModelRouter,
    backend: Arc<RwLock<Option<CandleBackend>>>,
    embedding_backend: Arc<RwLock<Option<EmbeddingBackend>>>,
}

impl InferenceEngine {
    pub fn new(config: InferenceConfig) -> Self {
        let registry = models::ModelRegistry::new(config.models_dir.clone());
        let hw = HardwareInfo::detect();
        let router = ModelRouter::new(hw.clone());
        let unified_registry = UnifiedRegistry::new(config.models_dir.clone());
        let adaptive_router = AdaptiveRouter::with_default_config(hw);
        let outcome_tracker = Arc::new(RwLock::new(OutcomeTracker::new()));

        Self {
            config,
            unified_registry,
            adaptive_router,
            outcome_tracker,
            remote_backend: RemoteBackend::new(),
            registry,
            router,
            backend: Arc::new(RwLock::new(None)),
            embedding_backend: Arc::new(RwLock::new(None)),
        }
    }

    /// Get or initialize the generative backend, loading the specified model.
    async fn ensure_backend(&self, model_name: &str) -> Result<(), InferenceError> {
        let read = self.backend.read().await;
        if read.is_some() {
            return Ok(());
        }
        drop(read);

        let mut write = self.backend.write().await;
        if write.is_some() {
            return Ok(());
        }

        let model_path = self.registry.ensure_model(model_name).await?;
        let device = self.config.device.unwrap_or_else(Device::auto);
        let backend = CandleBackend::load(&model_path, device)?;
        *write = Some(backend);
        Ok(())
    }

    /// Get or initialize the embedding backend.
    async fn ensure_embedding_backend(&self) -> Result<(), InferenceError> {
        let read = self.embedding_backend.read().await;
        if read.is_some() {
            return Ok(());
        }
        drop(read);

        let mut write = self.embedding_backend.write().await;
        if write.is_some() {
            return Ok(());
        }

        let model_path = self.registry.ensure_model(&self.config.embedding_model).await?;
        let device = self.config.device.unwrap_or_else(Device::auto);
        let backend = EmbeddingBackend::load(&model_path, device)?;
        *write = Some(backend);
        Ok(())
    }

    /// Route a prompt using the adaptive router (new). Returns full decision context.
    pub async fn route_adaptive(&self, prompt: &str) -> AdaptiveRoutingDecision {
        let tracker = self.outcome_tracker.read().await;
        self.adaptive_router.route(prompt, &self.unified_registry, &tracker)
    }

    /// Route a prompt to the best model without executing (legacy compat).
    pub fn route(&self, prompt: &str) -> RoutingDecision {
        self.router.route_generate(prompt, &self.registry)
    }

    /// Generate text from a prompt with outcome tracking.
    /// Returns `InferenceResult` with trace_id for reporting outcomes.
    pub async fn generate_tracked(&self, req: GenerateRequest) -> Result<InferenceResult, InferenceError> {
        let start = Instant::now();

        // Route using adaptive router
        let tracker_read = self.outcome_tracker.read().await;
        let decision = match req.model.clone() {
            Some(m) => AdaptiveRoutingDecision {
                model_id: m.clone(),
                model_name: m.clone(),
                task: InferenceTask::Generate,
                complexity: TaskComplexity::assess(&req.prompt),
                reason: "explicit model".into(),
                strategy: RoutingStrategy::Explicit,
                predicted_quality: 0.5,
                fallbacks: vec![],
            },
            None => self.adaptive_router.route(&req.prompt, &self.unified_registry, &tracker_read),
        };
        drop(tracker_read);

        // Record start
        let trace_id = {
            let mut tracker = self.outcome_tracker.write().await;
            tracker.record_start(&decision.model_id, decision.task, &decision.reason)
        };

        debug!(
            model = %decision.model_name,
            strategy = ?decision.strategy,
            reason = %decision.reason,
            trace = %trace_id,
            "adaptive-routed generate request"
        );

        // Execute — dispatch to local or remote backend, with fallback on failure
        let mut models_to_try = vec![decision.model_id.clone()];
        models_to_try.extend(decision.fallbacks.iter().cloned());

        let mut last_error = None;
        let mut used_model_name = decision.model_name.clone();

        for candidate_id in &models_to_try {
            let schema = self.unified_registry.get(candidate_id)
                .or_else(|| self.unified_registry.find_by_name(candidate_id))
                .cloned();

            let candidate_name = schema.as_ref()
                .map(|s| s.name.clone())
                .unwrap_or_else(|| candidate_id.clone());

            let is_remote = schema.as_ref().map(|s| s.is_remote()).unwrap_or(false);

            let result = if is_remote {
                let schema = schema.unwrap();
                self.remote_backend.generate(
                    &schema,
                    &req.prompt,
                    req.context.as_deref(),
                    req.params.temperature,
                    req.params.max_tokens,
                ).await
            } else {
                match self.ensure_backend(&candidate_name).await {
                    Ok(()) => {
                        let mut write = self.backend.write().await;
                        let backend = write.as_mut().unwrap();
                        tasks::generate::generate(backend, req.clone()).await
                    }
                    Err(e) => Err(e),
                }
            };

            match result {
                Ok(text) => {
                    let latency_ms = start.elapsed().as_millis() as u64;
                    let estimated_tokens = text.split_whitespace().count();
                    let mut tracker = self.outcome_tracker.write().await;
                    tracker.record_complete(&trace_id, latency_ms, 0, estimated_tokens);
                    used_model_name = candidate_name;
                    return Ok(InferenceResult {
                        text,
                        trace_id,
                        model_used: used_model_name,
                        latency_ms,
                    });
                }
                Err(e) => {
                    debug!(model = %candidate_name, error = %e, "model failed, trying fallback");
                    // Record failure for this model so the router learns
                    {
                        let mut tracker = self.outcome_tracker.write().await;
                        let fail_trace = tracker.record_start(candidate_id, decision.task, "fallback");
                        tracker.record_failure(&fail_trace, &e.to_string());
                    }
                    // Reset backend so next model can load
                    {
                        let mut write = self.backend.write().await;
                        *write = None;
                    }
                    last_error = Some(e);
                }
            }
        }

        // All models failed
        let e = last_error.unwrap_or(InferenceError::InferenceFailed("no models available".into()));
        let mut tracker = self.outcome_tracker.write().await;
        tracker.record_failure(&trace_id, &e.to_string());
        Err(e)
    }

    /// Generate text from a prompt (legacy API, no outcome tracking).
    /// When `req.model` is None, uses intelligent routing based on prompt complexity.
    pub async fn generate(&self, req: GenerateRequest) -> Result<String, InferenceError> {
        let model = match req.model.clone() {
            Some(m) => m,
            None => {
                let decision = self.router.route_generate(&req.prompt, &self.registry);
                debug!(
                    model = %decision.model,
                    reason = %decision.reason,
                    "auto-routed generate request"
                );
                decision.model
            }
        };
        self.ensure_backend(&model).await?;

        let mut write = self.backend.write().await;
        let backend = write.as_mut().unwrap();
        tasks::generate::generate(backend, req).await
    }

    /// Generate embeddings for text using the dedicated embedding model.
    /// Uses Qwen3-Embedding with proper last-token hidden state extraction.
    pub async fn embed(&self, req: EmbedRequest) -> Result<Vec<Vec<f32>>, InferenceError> {
        self.ensure_embedding_backend().await?;

        let mut write = self.embedding_backend.write().await;
        let backend = write.as_mut().unwrap();

        let instruction = req.instruction.as_deref()
            .unwrap_or("Retrieve relevant memory facts");

        let mut results = Vec::with_capacity(req.texts.len());
        for text in &req.texts {
            let embedding = if req.is_query {
                backend.embed_query(text, instruction)?
            } else {
                backend.embed_one(text)?
            };
            results.push(embedding);
        }
        Ok(results)
    }

    /// Classify text against candidate labels.
    /// When `req.model` is None, routes to the smallest available model.
    pub async fn classify(&self, req: ClassifyRequest) -> Result<Vec<ClassifyResult>, InferenceError> {
        let model = match req.model.clone() {
            Some(m) => m,
            None => {
                let m = self.router.route_small(&self.registry);
                debug!(model = %m, "auto-routed classify request");
                m
            }
        };
        self.ensure_backend(&model).await?;

        let mut write = self.backend.write().await;
        let backend = write.as_mut().unwrap();
        tasks::classify::classify(backend, req).await
    }

    /// List all known models and their status (new registry).
    pub fn list_models_unified(&self) -> Vec<ModelInfo> {
        self.unified_registry.list().iter().map(|m| ModelInfo::from(*m)).collect()
    }

    /// List all known models and their download status (legacy).
    pub fn list_models(&self) -> Vec<models::ModelInfo> {
        self.registry.list_models()
    }

    /// Download a model if not already present.
    pub async fn pull_model(&self, name: &str) -> Result<std::path::PathBuf, InferenceError> {
        self.registry.ensure_model(name).await
    }

    /// Remove a downloaded model.
    pub fn remove_model(&self, name: &str) -> Result<(), InferenceError> {
        self.registry.remove_model(name)
    }

    /// Register a model in the unified registry.
    pub fn register_model(&mut self, schema: ModelSchema) {
        self.unified_registry.register(schema);
    }

    /// Get outcome tracker for external use (e.g., memgine integration).
    pub fn outcome_tracker(&self) -> Arc<RwLock<OutcomeTracker>> {
        self.outcome_tracker.clone()
    }

    /// Export model performance profiles for persistence.
    pub async fn export_profiles(&self) -> Vec<ModelProfile> {
        let tracker = self.outcome_tracker.read().await;
        tracker.export_profiles()
    }

    /// Import model performance profiles (from persistence).
    pub async fn import_profiles(&self, profiles: Vec<ModelProfile>) {
        let mut tracker = self.outcome_tracker.write().await;
        tracker.import_profiles(profiles);
    }
}