systemprompt-ai 0.65.0

//! Streaming generation execution.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use crate::error::{AiError, ProviderCapability, Result};
use futures::Stream;
use std::collections::HashMap;
use std::pin::Pin;
use systemprompt_identifiers::AiRequestId;

use crate::models::ai::{AiRequest, GoogleSearchParams, SearchGroundedResponse, StreamChunk};
use crate::services::providers::{
    GenerationParams, ModelPricing, ProviderClient, SearchGenerationParams, ToolGenerationParams,
};

use super::service::AiService;
use super::stream_wrapper::{StreamStorageParams, StreamStorageWrapper};

impl AiService {
    pub async fn generate_stream(
        &self,
        request: &AiRequest,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk>> + Send>>> {
        let request_id = AiRequestId::generate();
        let start = std::time::Instant::now();
        let provider = self.get_provider(request.provider())?;

        if !provider.supports_streaming() {
            return Err(AiError::CapabilityUnsupported {
                provider: request.provider().to_owned(),
                capability: ProviderCapability::Streaming,
            });
        }

        let mut params = GenerationParams::new(
            &request.messages,
            request.model(),
            request.max_output_tokens(),
        );
        if let Some(sampling) = request.sampling.as_ref() {
            params = params.with_sampling(sampling);
        }

        let pricing = priced_model(provider.as_ref(), request.model())?;
        let inner_stream = provider.generate_stream(params).await?;

        let wrapped_stream = StreamStorageWrapper::new(StreamStorageParams {
            inner: inner_stream,
            storage: self.storage.clone(),
            audit_tasks: self.audit_tasks.clone(),
            request: request.clone(),
            request_id,
            start,
            provider: request.provider().to_owned(),
            model: request.model().to_owned(),
            pricing,
        });

        Ok(Box::pin(wrapped_stream))
    }

    pub async fn generate_with_tools_stream(
        &self,
        request: &AiRequest,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamChunk>> + Send>>> {
        let request_id = AiRequestId::generate();
        let start = std::time::Instant::now();
        let provider = self.get_provider(request.provider())?;

        if !provider.supports_streaming() {
            return Err(AiError::CapabilityUnsupported {
                provider: request.provider().to_owned(),
                capability: ProviderCapability::Streaming,
            });
        }

        let tools = request.tools.clone().unwrap_or_else(Vec::new);
        let mut base = GenerationParams::new(
            &request.messages,
            request.model(),
            request.max_output_tokens(),
        );
        if let Some(sampling) = request.sampling.as_ref() {
            base = base.with_sampling(sampling);
        }
        let params = ToolGenerationParams::new(base, tools);

        let pricing = priced_model(provider.as_ref(), request.model())?;
        let inner_stream = provider.generate_with_tools_stream(params).await?;

        let wrapped_stream = StreamStorageWrapper::new(StreamStorageParams {
            inner: inner_stream,
            storage: self.storage.clone(),
            audit_tasks: self.audit_tasks.clone(),
            request: request.clone(),
            request_id,
            start,
            provider: request.provider().to_owned(),
            model: request.model().to_owned(),
            pricing,
        });

        Ok(Box::pin(wrapped_stream))
    }

    pub async fn generate_with_google_search(
        &self,
        params: GoogleSearchParams<'_>,
    ) -> Result<SearchGroundedResponse> {
        let provider = self
            .providers
            .values()
            .find(|p| p.supports_google_search())
            .ok_or(AiError::NoProviderWithCapability {
                capability: ProviderCapability::GoogleSearch,
            })?;
        let model = params
            .model
            .or_else(|| {
                let cfg = self.default_model();
                (!cfg.is_empty()).then_some(cfg)
            })
            .unwrap_or_else(|| provider.default_model());
        let mut base = GenerationParams::new(&params.messages, model, params.max_output_tokens);
        if let Some(sampling) = params.sampling.as_ref() {
            base = base.with_sampling(sampling);
        }
        let mut search_params = SearchGenerationParams::new(base);
        if let Some(urls) = params.urls {
            search_params = search_params.with_urls(urls);
        }
        if let Some(schema) = params.response_schema {
            search_params = search_params.with_response_schema(schema);
        }
        provider.generate_with_google_search(search_params).await
    }

    pub async fn health_check(&self) -> Result<HashMap<String, bool>> {
        let mut health = HashMap::new();
        for (name, provider) in &self.providers {
            health.insert(format!("provider_{name}"), provider.is_available());
        }
        let tool_health = self.tool_provider.health_check().await?;
        for (service_id, is_healthy) in tool_health {
            health.insert(format!("tool_{service_id}"), is_healthy);
        }
        Ok(health)
    }
}

fn priced_model(provider: &dyn ProviderClient, model: &str) -> Result<ModelPricing> {
    provider
        .get_pricing(model)
        .ok_or_else(|| AiError::UnknownModel {
            provider: provider.name().to_owned(),
            model: model.to_owned(),
        })
}