diff --git a/shinkai-bin/shinkai-node/src/llm_provider/providers/claude.rs b/shinkai-bin/shinkai-node/src/llm_provider/providers/claude.rs index 6f62b0674..dfbbfa0e7 100644 --- a/shinkai-bin/shinkai-node/src/llm_provider/providers/claude.rs +++ b/shinkai-bin/shinkai-node/src/llm_provider/providers/claude.rs @@ -225,14 +225,47 @@ async fn handle_streaming_response( session_id: String, tools: Option>, ) -> Result { - let res = client + // Use tokio::select! to allow cancellation during the initial request phase + let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(500)); + let response_fut = client .post(url) .header("anthropic-version", "2023-06-01") .header("x-api-key", api_key) .header("content-type", "application/json") .json(&payload) - .send() - .await?; + .send(); + let mut response_fut = Box::pin(response_fut); + + // Wait for response or cancellation + let res = loop { + tokio::select! { + _ = interval.tick() => { + if let Some(ref inbox_name) = inbox_name { + if llm_stopper.should_stop(&inbox_name.to_string()) { + eprintln!("LLM job stopped by user request before response arrived"); + llm_stopper.reset(&inbox_name.to_string()); + + // Send WS message indicating the job is done + let _ = send_ws_update( + &ws_manager_trait, + Some(inbox_name.clone()), + &session_id, + "".to_string(), + false, + true, + Some("Stopped by user request".to_string()), + ) + .await; + + return Ok(LLMInferenceResponse::new("".to_string(), None, json!({}), Vec::new(), Vec::new(), None)); + } + } + }, + response = &mut response_fut => { + break response?; + } + } + }; // Check if it's an error response if !res.status().is_success() { diff --git a/shinkai-bin/shinkai-node/src/llm_provider/providers/gemini.rs b/shinkai-bin/shinkai-node/src/llm_provider/providers/gemini.rs index 52c5f823e..29bd7760e 100644 --- a/shinkai-bin/shinkai-node/src/llm_provider/providers/gemini.rs +++ b/shinkai-bin/shinkai-node/src/llm_provider/providers/gemini.rs @@ -137,7 +137,7 @@ impl LLMService for Gemini { inbox_name: Option, ws_manager_trait: Option>>, config: Option, - _llm_stopper: Arc, + llm_stopper: Arc, db: Arc, tracing_message_id: Option, ) -> Result { @@ -236,12 +236,38 @@ impl LLMService for Gemini { } } - let res = client + // Use tokio::select! to allow cancellation during the initial request phase + let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(500)); + let response_fut = client .post(&url) .header("Content-Type", "application/json") .json(&payload) - .send() - .await?; + .send(); + let mut response_fut = Box::pin(response_fut); + + // Wait for response or cancellation + let res = loop { + tokio::select! { + _ = interval.tick() => { + if let Some(ref inbox_name) = inbox_name { + if llm_stopper.should_stop(&inbox_name.to_string()) { + shinkai_log( + ShinkaiLogOption::JobExecution, + ShinkaiLogLevel::Info, + "LLM job stopped by user request before response arrived", + ); + llm_stopper.reset(&inbox_name.to_string()); + + return Ok(LLMInferenceResponse::new("".to_string(), None, json!({}), Vec::new(), Vec::new(), None)); + } + } + }, + response = &mut response_fut => { + break response?; + } + } + }; + shinkai_log( ShinkaiLogOption::JobExecution, ShinkaiLogLevel::Debug, @@ -260,6 +286,39 @@ impl LLMService for Gemini { let mut thinking_started = false; while let Some(item) = stream.next().await { + // Check if we need to stop the LLM job + if let Some(ref inbox_name) = inbox_name { + if llm_stopper.should_stop(&inbox_name.to_string()) { + shinkai_log( + ShinkaiLogOption::JobExecution, + ShinkaiLogLevel::Info, + "LLM job stopped by user request during streaming", + ); + llm_stopper.reset(&inbox_name.to_string()); + + // Send WS message indicating the job is done + let _ = send_ws_update( + &ws_manager_trait, + Some(inbox_name.clone()), + &session_id, + regular_content.clone(), + false, + true, + Some("Stopped by user request".to_string()), + ) + .await; + + return Ok(LLMInferenceResponse::new( + regular_content, + if thinking_content.is_empty() { None } else { Some(thinking_content) }, + json!({}), + function_calls, + generated_files, + None, + )); + } + } + match item { Ok(chunk) => { process_chunk( diff --git a/shinkai-bin/shinkai-node/src/llm_provider/providers/grok.rs b/shinkai-bin/shinkai-node/src/llm_provider/providers/grok.rs index 073cc5840..6ede1845a 100644 --- a/shinkai-bin/shinkai-node/src/llm_provider/providers/grok.rs +++ b/shinkai-bin/shinkai-node/src/llm_provider/providers/grok.rs @@ -169,13 +169,38 @@ async fn handle_streaming_response( session_id: String, tools: Option>, // Add tools parameter ) -> Result { - let res = client + // Use tokio::select! to allow cancellation during the initial request phase + let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(500)); + let response_fut = client .post(url) .bearer_auth(api_key) .header("Content-Type", "application/json") .json(&payload) - .send() - .await?; + .send(); + let mut response_fut = Box::pin(response_fut); + + // Wait for response or cancellation + let res = loop { + tokio::select! { + _ = interval.tick() => { + if let Some(ref inbox_name) = inbox_name { + if llm_stopper.should_stop(&inbox_name.to_string()) { + shinkai_log( + ShinkaiLogOption::JobExecution, + ShinkaiLogLevel::Info, + "LLM job stopped by user request before response arrived", + ); + llm_stopper.reset(&inbox_name.to_string()); + + return Ok(LLMInferenceResponse::new("".to_string(), None, json!({}), Vec::new(), Vec::new(), None)); + } + } + }, + response = &mut response_fut => { + break response?; + } + } + }; // Check if it's an error response let status = res.status(); diff --git a/shinkai-bin/shinkai-node/src/llm_provider/providers/groq.rs b/shinkai-bin/shinkai-node/src/llm_provider/providers/groq.rs index af60e6181..a779437b7 100644 --- a/shinkai-bin/shinkai-node/src/llm_provider/providers/groq.rs +++ b/shinkai-bin/shinkai-node/src/llm_provider/providers/groq.rs @@ -215,13 +215,38 @@ async fn handle_streaming_response( session_id: String, tools: Option>, // Add tools parameter ) -> Result { - let res = client + // Use tokio::select! to allow cancellation during the initial request phase + let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(500)); + let response_fut = client .post(url) .bearer_auth(api_key) .header("Content-Type", "application/json") .json(&payload) - .send() - .await?; + .send(); + let mut response_fut = Box::pin(response_fut); + + // Wait for response or cancellation + let res = loop { + tokio::select! { + _ = interval.tick() => { + if let Some(ref inbox_name) = inbox_name { + if llm_stopper.should_stop(&inbox_name.to_string()) { + shinkai_log( + ShinkaiLogOption::JobExecution, + ShinkaiLogLevel::Info, + "LLM job stopped by user request before response arrived", + ); + llm_stopper.reset(&inbox_name.to_string()); + + return Ok(LLMInferenceResponse::new("".to_string(), None, json!({}), Vec::new(), Vec::new(), None)); + } + } + }, + response = &mut response_fut => { + break response?; + } + } + }; // Check if it's an error response if !res.status().is_success() { diff --git a/shinkai-bin/shinkai-node/src/llm_provider/providers/openai.rs b/shinkai-bin/shinkai-node/src/llm_provider/providers/openai.rs index feff46626..e5ce0c614 100644 --- a/shinkai-bin/shinkai-node/src/llm_provider/providers/openai.rs +++ b/shinkai-bin/shinkai-node/src/llm_provider/providers/openai.rs @@ -724,7 +724,9 @@ pub async fn handle_streaming_response( tools: Option>, headers: Option, ) -> Result { - let res = client + // Use tokio::select! to allow cancellation during the initial request phase + let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(500)); + let response_fut = client .post(url) .bearer_auth(api_key) .header("Content-Type", "application/json") @@ -776,8 +778,43 @@ pub async fn handle_streaming_response( .unwrap_or(""), ) .json(&payload) - .send() - .await?; + .send(); + let mut response_fut = Box::pin(response_fut); + + // Wait for response or cancellation + let res = loop { + tokio::select! { + _ = interval.tick() => { + if let Some(ref inbox_name) = inbox_name { + if llm_stopper.should_stop(&inbox_name.to_string()) { + shinkai_log( + ShinkaiLogOption::JobExecution, + ShinkaiLogLevel::Info, + "LLM job stopped by user request before response arrived", + ); + llm_stopper.reset(&inbox_name.to_string()); + + // Send WS message indicating the job is done + let _ = send_ws_update( + &ws_manager_trait, + Some(inbox_name.clone()), + &session_id, + "".to_string(), + false, + true, + Some("Stopped by user request".to_string()), + ) + .await; + + return Ok(LLMInferenceResponse::new("".to_string(), None, json!({}), Vec::new(), Vec::new(), None)); + } + } + }, + response = &mut response_fut => { + break response?; + } + } + }; // Check if it's an error response if !res.status().is_success() { diff --git a/shinkai-bin/shinkai-node/src/llm_provider/providers/openrouter.rs b/shinkai-bin/shinkai-node/src/llm_provider/providers/openrouter.rs index 947d7ac49..fa05dab84 100644 --- a/shinkai-bin/shinkai-node/src/llm_provider/providers/openrouter.rs +++ b/shinkai-bin/shinkai-node/src/llm_provider/providers/openrouter.rs @@ -170,14 +170,39 @@ async fn handle_streaming_response( session_id: String, tools: Option>, // Add tools parameter ) -> Result { - let res = client + // Use tokio::select! to allow cancellation during the initial request phase + let mut interval = tokio::time::interval(tokio::time::Duration::from_millis(500)); + let response_fut = client .post(url) .bearer_auth(api_key) .header("Content-Type", "application/json") .header("X-Title", "Shinkai") .json(&payload) - .send() - .await?; + .send(); + let mut response_fut = Box::pin(response_fut); + + // Wait for response or cancellation + let res = loop { + tokio::select! { + _ = interval.tick() => { + if let Some(ref inbox_name) = inbox_name { + if llm_stopper.should_stop(&inbox_name.to_string()) { + shinkai_log( + ShinkaiLogOption::JobExecution, + ShinkaiLogLevel::Info, + "LLM job stopped by user request before response arrived", + ); + llm_stopper.reset(&inbox_name.to_string()); + + return Ok(LLMInferenceResponse::new("".to_string(), None, json!({}), Vec::new(), Vec::new(), None)); + } + } + }, + response = &mut response_fut => { + break response?; + } + } + }; // Check if it's an error response if !res.status().is_success() {