From a619b010640a356d95e241ff07e17717ee4c2fe1 Mon Sep 17 00:00:00 2001
From: Saifeddine ALOUI <aloui.seifeddine@gmail.com>
Date: Fri, 10 Jan 2025 22:17:13 +0100
Subject: [PATCH] Next test of timeout

---
 lightrag/api/lightrag_server.py | 11 ++++++++---
 lightrag/llm.py                 |  3 +--
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py
index d4cddd6c..40b63463 100644
--- a/lightrag/api/lightrag_server.py
+++ b/lightrag/api/lightrag_server.py
@@ -101,12 +101,17 @@ def parse_args():
         help="Embedding model name (default: bge-m3:latest)",
     )
 
+    def timeout_type(value):
+        if value is None or value == "None":
+            return None
+        return int(value)
+
     parser.add_argument(
         "--timeout",
-        default=300,
-        help="Timeout is seconds (useful when using slow AI)",
+        default=None,
+        type=timeout_type,
+        help="Timeout in seconds (useful when using slow AI). Use None for infinite timeout",
     )
-
     # RAG configuration
     parser.add_argument(
         "--max-async", type=int, default=4, help="Maximum async operations (default: 4)"
diff --git a/lightrag/llm.py b/lightrag/llm.py
index 4e01dd51..7a51d025 100644
--- a/lightrag/llm.py
+++ b/lightrag/llm.py
@@ -407,11 +407,10 @@ async def lollms_model_if_cache(
     full_prompt += prompt
 
     request_data["prompt"] = full_prompt
-    timeout = aiohttp.ClientTimeout(total=kwargs.get("timeout", 300))  # 300 seconds = 5 minutes
+    timeout = aiohttp.ClientTimeout(total=kwargs.get("timeout", None))
 
     async with aiohttp.ClientSession(timeout=timeout) as session:
         if stream:
-
             async def inner():
                 async with session.post(
                     f"{base_url}/lollms_generate", json=request_data