From 02d05615866a35da38e31630ac41c61aaf64a59d Mon Sep 17 00:00:00 2001 From: Conrad Irwin Date: Fri, 27 Sep 2024 00:36:17 -0600 Subject: [PATCH] Fix read timeout for ollama (#18417) Supercedes: #18310 Release Notes: - Fixed `low_speed_timeout_in_seconds` for Ollama --- crates/ollama/src/ollama.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/crates/ollama/src/ollama.rs b/crates/ollama/src/ollama.rs index 84404afce1..a38b9e7a56 100644 --- a/crates/ollama/src/ollama.rs +++ b/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, + low_speed_timeout: Option, ) -> Result>> { 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() {