Added abort request for generate_stream

ollama-rs/src/error.rs

@@ -15,6 +15,8 @@ pub enum OllamaError {
     ReqwestError(#[from] reqwest::Error),
     #[error("Internal Ollama error")]
     InternalError(InternalOllamaError),
+    #[error("Ollama aborted the request")]
+    Abort,
     #[error("Error in Ollama")]
     Other(String),
 }

ollama-rs/src/generation/completion/mod.rs

@@ -47,18 +47,25 @@ impl Ollama {
             ));
         }
 
-        let stream = Box::new(res.bytes_stream().map(|res| match res {
-            Ok(bytes) => {
-                let res = serde_json::Deserializer::from_slice(&bytes).into_iter();
-                let res = res
-                    .filter_map(Result::ok) // Filter out the errors
-                    .collect::<Vec<GenerationResponse>>();
-                Ok(res)
+        let stream = Box::new(res.bytes_stream().map(move |res| {
+            if let Some(abort_signal) = request.abort_signal.as_ref() {
+                if abort_signal.aborted() {
+                    return Err(OllamaError::Abort);
+                }
+            }
+            match res {
+                Ok(bytes) => {
+                    let res = serde_json::Deserializer::from_slice(&bytes).into_iter();
+                    let res = res
+                        .filter_map(Result::ok) // Filter out the errors
+                        .collect::<Vec<GenerationResponse>>();
+                    Ok(res)
+                }
+                Err(e) => Err(OllamaError::Other(format!(
+                    "Failed to read response: {}",
+                    e
+                ))),
             }
-            Err(e) => Err(OllamaError::Other(format!(
-                "Failed to read response: {}",
-                e
-            ))),
         }));
 
         Ok(std::pin::Pin::from(stream))

ollama-rs/src/generation/completion/request.rs

@@ -1,3 +1,5 @@
+use std::sync::{atomic::AtomicBool, Arc};
+
 use serde::Serialize;
 
 use crate::generation::{
@@ -8,6 +10,28 @@ use crate::generation::{
 
 use super::GenerationContext;
 
+#[derive(Debug, Clone)]
+pub struct AbortSignal {
+    pub(crate) abort_signal: Arc<AtomicBool>,
+}
+
+impl AbortSignal {
+    pub fn new() -> Self {
+        Self {
+            abort_signal: Arc::new(AtomicBool::new(false)),
+        }
+    }
+
+    pub fn abort(&self) {
+        self.abort_signal
+            .store(true, std::sync::atomic::Ordering::Relaxed);
+    }
+
+    pub fn aborted(&self) -> bool {
+        self.abort_signal.load(std::sync::atomic::Ordering::Relaxed)
+    }
+}
+
 /// A generation request to Ollama.
 #[derive(Debug, Clone, Serialize)]
 pub struct GenerationRequest {
@@ -24,6 +48,8 @@ pub struct GenerationRequest {
     pub format: Option<FormatType>,
     pub keep_alive: Option<KeepAlive>,
     pub(crate) stream: bool,
+    #[serde(skip)]
+    pub abort_signal: Option<AbortSignal>,
 }
 
 impl GenerationRequest {
@@ -41,6 +67,7 @@ impl GenerationRequest {
             keep_alive: None,
             // Stream value will be overwritten by Ollama::generate_stream() and Ollama::generate() methods
             stream: false,
+            abort_signal: None,
         }
     }
 
@@ -103,4 +130,10 @@ impl GenerationRequest {
         self.keep_alive = Some(keep_alive);
         self
     }
+
+    /// Sets the abort signal for the request
+    pub fn abort_signal(mut self, abort_signal: AbortSignal) -> Self {
+        self.abort_signal = Some(abort_signal);
+        self
+    }
 }