flow_adapter_ai/

registry.rs

use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;

use crate::error::CloudAiError;
use crate::request::{
    CloudAiRequest, CloudAiResponse, EmbeddingResponse, ToolDispatcher, ToolSpec,
};
use crate::stream::{LlmStreamEvent, LlmStreamSink};

/// Where a provider runs. Drives the Settings + node-inspector UI grouping
/// (one "LLM" area split into Local vs Cloud) without the frontend having to
/// hardcode provider names. `Cloud` providers make outbound network calls
/// gated by `allow_cloud_ai`; `Local` providers hit a user-configured
/// on-device endpoint gated by `allow_local_ai`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ProviderCategory {
    Local,
    Cloud,
}

/// Cloud-AI provider abstraction. `Cloud` implementations make outbound
/// network calls - the egress is intentional and gated by the
/// `allow_cloud_ai` setting at the executor layer. `Local` implementations
/// target a user-configured localhost endpoint and are gated by
/// `allow_local_ai`.
#[async_trait]
pub trait CloudAiProvider: Send + Sync {
    fn name(&self) -> &str;

    /// Convention-based environment variable name for this provider's API key.
    fn env_var(&self) -> &str;

    /// Default models suggested in the UI dropdown. Users can also type any
    /// model name the provider accepts.
    fn default_models(&self) -> &[&str];

    /// Per-model capability tags advertised to the UI, using the same
    /// vocabulary as the Hub catalog (`code | tool_use | reasoning | vision |
    /// embedding`). The node inspector gates capability-specific fields (image
    /// input, thinking, tool binding) on these the same way it does for local
    /// models resolved from the Hub. Cloud providers override this; the default
    /// advertises nothing, so the inspector falls back to showing every field.
    fn model_capabilities(&self) -> Vec<ModelCapabilityInfo> {
        Vec::new()
    }

    /// Local vs cloud. Defaults to `Cloud` so existing providers need no
    /// change; the local OpenAI-compatible provider overrides it.
    fn category(&self) -> ProviderCategory {
        ProviderCategory::Cloud
    }

    async fn invoke(&self, req: &CloudAiRequest) -> Result<CloudAiResponse, CloudAiError>;

    /// Streaming variant of `invoke`. Providers that support real
    /// streaming override this to push tokens through `sink` as they
    /// arrive; the default impl wraps `invoke` and emits a single
    /// `done = true` event so non-streaming providers are still
    /// observable on the chip without per-call branching at the call
    /// site.
    ///
    /// Implementations MUST emit exactly one event with `done = true`
    /// (either with the final partial delta or alone), even on failure.
    async fn invoke_stream(
        &self,
        req: &CloudAiRequest,
        sink: &dyn LlmStreamSink,
    ) -> Result<CloudAiResponse, CloudAiError> {
        let call_id = req.call_id.clone().unwrap_or_default();
        match self.invoke(req).await {
            Ok(resp) => {
                sink.emit(LlmStreamEvent::final_delta(&call_id, &resp)).await;
                Ok(resp)
            }
            Err(e) => {
                sink.emit(LlmStreamEvent::error(&call_id, self.name(), e.to_string()))
                    .await;
                Err(e)
            }
        }
    }

    /// Tool-use / function-calling loop. Providers that support tools override
    /// this: they expose `tools` to the model, run each requested call through
    /// `dispatcher`, feed the results back, and repeat until the model stops
    /// asking (or `max_iters` is hit). The default ignores tools and does a
    /// single `invoke`, so providers without tool support degrade gracefully.
    async fn invoke_tools(
        &self,
        req: &CloudAiRequest,
        _tools: &[ToolSpec],
        _dispatcher: &dyn ToolDispatcher,
        _max_iters: usize,
    ) -> Result<CloudAiResponse, CloudAiError> {
        self.invoke(req).await
    }

    /// Embeddings call for `task: "embedding"`. Providers with an embeddings
    /// endpoint override this; the default reports the capability unsupported
    /// (e.g. Anthropic exposes no embeddings API).
    async fn embed(&self, _req: &CloudAiRequest) -> Result<EmbeddingResponse, CloudAiError> {
        Err(CloudAiError::Unsupported {
            provider: self.name().to_string(),
            capability: "embeddings".to_string(),
        })
    }
}

#[derive(Default)]
pub struct CloudAiRegistry {
    providers: HashMap<String, Arc<dyn CloudAiProvider>>,
}

impl CloudAiRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn register(&mut self, provider: Arc<dyn CloudAiProvider>) {
        self.providers.insert(provider.name().to_string(), provider);
    }

    pub fn get(&self, name: &str) -> Option<Arc<dyn CloudAiProvider>> {
        self.providers.get(name).cloned()
    }

    pub fn list(&self) -> Vec<ProviderInfo> {
        let mut out: Vec<_> = self
            .providers
            .values()
            .map(|p| ProviderInfo {
                name: p.name().to_string(),
                env_var: p.env_var().to_string(),
                default_models: p.default_models().iter().map(|s| s.to_string()).collect(),
                model_capabilities: p.model_capabilities(),
                category: p.category(),
            })
            .collect();
        out.sort_by(|a, b| a.name.cmp(&b.name));
        out
    }
}

#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ProviderInfo {
    pub name: String,
    pub env_var: String,
    pub default_models: Vec<String>,
    /// Capability tags per model id (see [`CloudAiProvider::model_capabilities`]).
    #[serde(default)]
    pub model_capabilities: Vec<ModelCapabilityInfo>,
    pub category: ProviderCategory,
}

/// Capability tags for a single model, surfaced to the node inspector so cloud
/// models get the same capability-gated fields as local (Hub-catalog) models.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ModelCapabilityInfo {
    pub model: String,
    /// Tags from the Hub vocabulary: `code | tool_use | reasoning | vision |
    /// embedding`.
    pub capabilities: Vec<String>,
}

impl ModelCapabilityInfo {
    /// Build an entry from a model id and its capability tags.
    pub fn new(model: &str, capabilities: &[&str]) -> Self {
        Self {
            model: model.to_string(),
            capabilities: capabilities.iter().map(|s| s.to_string()).collect(),
        }
    }
}