systemprompt-ai 0.65.0

//! `AiService` construction: provider clients from policy plus registry.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use crate::error::Result;
use std::collections::HashMap;
use std::sync::Arc;

use crate::models::RequestStatus;
use crate::models::ai::AiRequest;
use crate::repository::AiRepositories;
use crate::services::config::{AiConfigError, ConfigValidator};
use crate::services::providers::{ProviderClient, ProviderClientParams, ProviderFactory};
use crate::services::tooled::{ResponseSynthesizer, TooledExecutor};
use crate::services::tools::ToolDiscovery;
use crate::services::upstream::UpstreamTarget;

use super::super::request_storage::{RequestStorage, StoreParams};

use systemprompt_config::SecretsBootstrap;
use systemprompt_manifest::services::{
    AiConfig, AiProviderConfig, ProviderEntry, ProviderRegistry,
};
use systemprompt_traits::{DynAiSessionProvider, ToolProvider};
use tokio_util::task::TaskTracker;

pub struct AiService {
    pub(super) providers: HashMap<String, Arc<dyn ProviderClient>>,
    pub(super) tool_provider: Arc<dyn ToolProvider>,
    pub(super) tool_discovery: Arc<ToolDiscovery>,
    pub(super) tooled_executor: TooledExecutor,
    pub(super) synthesizer: ResponseSynthesizer,
    pub(super) storage: RequestStorage,
    pub(super) audit_tasks: TaskTracker,
    default_provider: String,
    default_model: String,
    default_max_output_tokens: u32,
}

impl std::fmt::Debug for AiService {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AiService")
            .field("default_provider", &self.default_provider)
            .finish_non_exhaustive()
    }
}

#[derive(Clone)]
pub struct AiServiceProviders {
    pub tools: Arc<dyn ToolProvider>,
    pub sessions: DynAiSessionProvider,
}

impl std::fmt::Debug for AiServiceProviders {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AiServiceProviders").finish_non_exhaustive()
    }
}

impl AiService {
    pub fn new(
        registry: &ProviderRegistry,
        ai_config: &AiConfig,
        providers: AiServiceProviders,
        repositories: &AiRepositories,
    ) -> Result<Self> {
        let AiServiceProviders {
            tools: tool_provider,
            sessions: session_provider,
        } = providers;
        let mut missing_env_vars = Vec::new();
        let providers = Self::build_providers(registry, ai_config, &mut missing_env_vars)?;
        ConfigValidator::validate(ai_config, &providers, &missing_env_vars)?;

        let default_provider = ai_config.default_provider.clone();
        let provider = providers.get(&default_provider).ok_or_else(|| {
            crate::error::AiError::from(AiConfigError::DefaultProviderNoConnectivity {
                provider: default_provider.clone(),
                connected: providers.keys().cloned().collect(),
                withheld: None,
            })
        })?;

        let default_model = ai_config
            .providers
            .get(&default_provider)
            .filter(|pc| !pc.default_model.is_empty())
            .map_or_else(
                || provider.default_model().to_owned(),
                |pc| pc.default_model.clone(),
            );

        let tool_discovery = Arc::new(ToolDiscovery::new(Arc::clone(&tool_provider)));
        let tooled_executor = TooledExecutor::new(Arc::clone(&tool_provider));

        let storage = RequestStorage::new(
            repositories.requests.clone(),
            repositories.payloads.clone(),
            session_provider,
        );

        Ok(Self {
            providers,
            tool_provider,
            tool_discovery,
            tooled_executor,
            synthesizer: ResponseSynthesizer::new(),
            storage,
            audit_tasks: TaskTracker::new(),
            default_provider,
            default_model,
            default_max_output_tokens: ai_config.default_max_output_tokens.unwrap_or(8192),
        })
    }

    #[must_use]
    pub fn with_context_materializer(
        mut self,
        materializer: systemprompt_traits::DynContextMaterializer,
    ) -> Self {
        self.storage = self.storage.with_context_materializer(materializer);
        self
    }

    // Why: streaming audit rows are written by tasks that outlive the stream
    // that produced them; the owner drains this tracker before the runtime
    // is torn down so an abandoned stream's usage is never lost at shutdown.
    #[must_use]
    pub const fn audit_tasks(&self) -> &TaskTracker {
        &self.audit_tasks
    }

    fn build_providers(
        registry: &ProviderRegistry,
        ai_config: &AiConfig,
        missing_env_vars: &mut Vec<String>,
    ) -> Result<HashMap<String, Arc<dyn ProviderClient>>> {
        SecretsBootstrap::get()?;
        let mut providers: HashMap<String, Arc<dyn ProviderClient>> = HashMap::new();

        for (name, policy) in &ai_config.providers {
            if !policy.enabled {
                continue;
            }
            let Some(entry) = registry.find_provider(name) else {
                tracing::warn!(
                    provider = %name,
                    "AI policy enables provider but the profile registry has no connectivity \
                     entry — skipping"
                );
                continue;
            };

            // Why: a provider without a usable credential cannot serve a
            // request; registering it would fail every call at the wire
            // instead of at boot, so it is withheld and reported.
            let target = match UpstreamTarget::from_secrets(entry) {
                Ok(target) => target,
                Err(e) => {
                    tracing::warn!(
                        provider = %name,
                        error = %e,
                        "upstream credential unusable — provider withheld from the registry"
                    );
                    missing_env_vars.push(format!("Provider '{name}': {e}"));
                    continue;
                },
            };

            let provider = Self::build_one(entry, policy, target)?;
            providers.insert(name.clone(), provider);
        }

        Ok(providers)
    }

    fn build_one(
        entry: &ProviderEntry,
        policy: &AiProviderConfig,
        target: UpstreamTarget,
    ) -> Result<Arc<dyn ProviderClient>> {
        let params = ProviderClientParams {
            name: entry.name.as_str(),
            target,
            google_search_enabled: policy.google_search_enabled,
            resilience: &policy.resilience,
            models: &entry.models,
            default_model: (!policy.default_model.is_empty())
                .then_some(policy.default_model.as_str()),
        };
        ProviderFactory::create(&params)
    }

    pub fn default_provider(&self) -> &str {
        &self.default_provider
    }

    pub fn default_model(&self) -> &str {
        &self.default_model
    }

    pub const fn default_max_output_tokens(&self) -> u32 {
        self.default_max_output_tokens
    }

    pub(super) fn get_provider(&self, name: &str) -> Result<Arc<dyn ProviderClient>> {
        self.providers
            .get(name)
            .cloned()
            .ok_or_else(|| crate::error::AiError::ProviderNotFound {
                provider: name.to_owned(),
            })
    }

    pub(super) async fn audit(&self, params: &StoreParams<'_>) {
        if let Err(e) = self.storage.store(params).await {
            tracing::error!(
                error = %e,
                provider = %params.request.provider(),
                model = %params.request.model(),
                status = ?params.status,
                "audit write failed"
            );
        }
    }

    pub(super) async fn store_error(
        &self,
        request: &AiRequest,
        request_id: systemprompt_identifiers::AiRequestId,
        latency_ms: u64,
        error_message: String,
    ) {
        use crate::models::ai::AiResponse;

        let error_response = AiResponse::new(
            request_id,
            String::new(),
            request.provider().to_owned(),
            request.model().to_owned(),
        )
        .with_latency(latency_ms);
        self.audit(&StoreParams {
            request,
            response: &error_response,
            context: &request.context,
            status: RequestStatus::Failed,
            error_message: Some(&error_message),
            cost_microdollars: 0,
        })
        .await;
    }
}