Fix read timeout for ollama (#18417)

Conrad Irwin created 1 year ago

Supercedes: #18310

Release Notes:

- Fixed `low_speed_timeout_in_seconds` for Ollama

Change summary

crates/ollama/src/ollama.rs | 10 +++++++---
1 file changed, 7 insertions(+), 3 deletions(-)

Detailed changes

crates/ollama/src/ollama.rs

@@ -1,6 +1,6 @@
 use anyhow::{anyhow, Context, Result};
 use futures::{io::BufReader, stream::BoxStream, AsyncBufReadExt, AsyncReadExt, StreamExt};
-use http_client::{http, AsyncBody, HttpClient, Method, Request as HttpRequest};
+use http_client::{http, AsyncBody, HttpClient, HttpRequestExt, Method, Request as HttpRequest};
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use serde_json::{value::RawValue, Value};
@@ -262,14 +262,18 @@ pub async fn stream_chat_completion(
     client: &dyn HttpClient,
     api_url: &str,
     request: ChatRequest,
-    _: Option<Duration>,
+    low_speed_timeout: Option<Duration>,
 ) -> Result<BoxStream<'static, Result<ChatResponseDelta>>> {
     let uri = format!("{api_url}/api/chat");
-    let request_builder = http::Request::builder()
+    let mut request_builder = http::Request::builder()
         .method(Method::POST)
         .uri(uri)
         .header("Content-Type", "application/json");
 
+    if let Some(low_speed_timeout) = low_speed_timeout {
+        request_builder = request_builder.read_timeout(low_speed_timeout);
+    }
+
     let request = request_builder.body(AsyncBody::from(serde_json::to_string(&request)?))?;
     let mut response = client.send(request).await?;
     if response.status().is_success() {