Function generate_text_stream

cake-core/src/cake/sharding/api/text.rs:182–307

(
    state: web::Data<Arc<RwLock<Master<M>>>>,
    request: ChatRequest,
)

}

async fn generate_text_stream<M: Model>(
    state: web::Data<Arc<RwLock<Master<M>>>>,
    request: ChatRequest,
) -> HttpResponse {
    let id = format!("chatcmpl-{}", uuid::Uuid::new_v4());
    let created = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs();
    let model = M::MODEL_NAME.to_string();

    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<Option<String>>();

    let state_clone = state.clone();
    tokio::spawn(async move {
        let mut master = state_clone.write().await;

        if let Err(e) = master.reset() {
            log::error!("reset error: {e}");
            let _ = tx.send(None);
            return;
        }

        let model = match master.model.as_mut() {
            Some(m) if m.output_modality() == OutputModality::Text => m,
            _ => {
                log::error!("no text model loaded");
                let _ = tx.send(None);
                return;
            }
        };
        for message in request.messages {
            if let Err(e) = model.add_message(message) {
                log::error!("add_message error: {e}");
                let _ = tx.send(None);
                return;
            }
        }

        if let Err(e) = master
            .generate_text(request.max_tokens, |data| {
                if data.is_empty() {
                    let _ = tx.send(None);
                } else {
                    let _ = tx.send(Some(data.to_string()));
                }
            })
            .await
        {
            log::error!("generate_text error: {e}");
            let _ = tx.send(None);
        }

        let _ = master.goodbye().await;
    });

    let stream = async_stream::stream! {
        // Send initial role chunk