From e130ad74d147a6871df9a9f043a16b89577b2f75 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mikael=20Tur=C3=B8y?= <mikael.turoy@holbergeeg.org>
Date: Wed, 12 Jun 2024 22:37:35 +0200
Subject: [PATCH] Added timeout setting for ollama streaming response

---
 TROUBLESHOOTING.md          | 4 ++++
 backend/apps/ollama/main.py | 3 ++-
 backend/config.py           | 1 +
 3 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/TROUBLESHOOTING.md b/TROUBLESHOOTING.md
index 8e8f89da0..d1a6dba28 100644
--- a/TROUBLESHOOTING.md
+++ b/TROUBLESHOOTING.md
@@ -18,6 +18,10 @@ If you're experiencing connection issues, it’s often due to the WebUI docker c
 docker run -d --network=host -v open-webui:/app/backend/data -e OLLAMA_BASE_URL=http://127.0.0.1:11434 --name open-webui --restart always ghcr.io/open-webui/open-webui:main
 ```
 
+### Error on Slow Reponses for Ollama
+
+Open WebUI has a default timeout of 15 minutes for Ollama to finish generating the response. If needed, this can be adjusted via the environment variable OLLAMA_GENERATE_TIMEOUT, which sets the timeout in seconds.
+
 ### General Connection Errors
 
 **Ensure Ollama Version is Up-to-Date**: Always start by checking that you have the latest version of Ollama. Visit [Ollama's official site](https://ollama.com/) for the latest updates.
diff --git a/backend/apps/ollama/main.py b/backend/apps/ollama/main.py
index 144755418..f9a906986 100644
--- a/backend/apps/ollama/main.py
+++ b/backend/apps/ollama/main.py
@@ -46,6 +46,7 @@ from config import (
     SRC_LOG_LEVELS,
     OLLAMA_BASE_URLS,
     ENABLE_OLLAMA_API,
+    OLLAMA_GENERATE_TIMEOUT,
     ENABLE_MODEL_FILTER,
     MODEL_FILTER_LIST,
     UPLOAD_DIR,
@@ -154,7 +155,7 @@ async def cleanup_response(
 async def post_streaming_url(url: str, payload: str):
     r = None
     try:
-        session = aiohttp.ClientSession(trust_env=True)
+        session = aiohttp.ClientSession(trust_env=True, timeout=aiohttp.ClientTimeout(total=OLLAMA_GENERATE_TIMEOUT))
         r = await session.post(url, data=payload)
         r.raise_for_status()
 
diff --git a/backend/config.py b/backend/config.py
index 30a23f29e..995a48a01 100644
--- a/backend/config.py
+++ b/backend/config.py
@@ -425,6 +425,7 @@ OLLAMA_API_BASE_URL = os.environ.get(
 )
 
 OLLAMA_BASE_URL = os.environ.get("OLLAMA_BASE_URL", "")
+OLLAMA_GENERATE_TIMEOUT = int(os.environ.get("OLLAMA_GENERATE_TIMEOUT", "900"))
 K8S_FLAG = os.environ.get("K8S_FLAG", "")
 USE_OLLAMA_DOCKER = os.environ.get("USE_OLLAMA_DOCKER", "false")