Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix read timeout for ollama (#18417)

Browse source 
Supercedes: #18310

Release Notes:

- Fixed `low_speed_timeout_in_seconds` for Ollama
This commit is contained in:
Conrad Irwin 2024-09-27 00:36:17 -06:00 • committed by GitHub
parent 1be3c44550
commit 02d0561586
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 7 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

10
crates/ollama/src/ollama.rs
View file

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