From adb288c5bb7ac6753daf898cb99c1458f9663773 Mon Sep 17 00:00:00 2001
From: Saifeddine ALOUI <aloui.seifeddine@gmail.com>
Date: Fri, 10 Jan 2025 21:39:25 +0100
Subject: [PATCH] added timeout

---
 lightrag/api/lightrag_server.py | 23 +++++++++++++++++++++++
 lightrag/llm.py                 |  4 +++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py
index 4f8e38cd..1175afab 100644
--- a/lightrag/api/lightrag_server.py
+++ b/lightrag/api/lightrag_server.py
@@ -101,6 +101,12 @@ def parse_args():
         help="Embedding model name (default: bge-m3:latest)",
     )
 
+    parser.add_argument(
+        "--timeout",
+        default=300,
+        help="Timeout is seconds (useful when using slow AI)",
+    )
+
     # RAG configuration
     parser.add_argument(
         "--max-async", type=int, default=4, help="Maximum async operations (default: 4)"
@@ -139,6 +145,22 @@ def parse_args():
         default=None,
     )
 
+    # Optional https parameters
+    parser.add_argument(
+        "--ssl",
+        action="store_true",
+        help="Enable HTTPS (default: False)"
+    )
+    parser.add_argument(
+        "--ssl-certfile",
+        default=None,
+        help="Path to SSL certificate file (required if --ssl is enabled)"
+    )
+    parser.add_argument(
+        "--ssl-keyfile",
+        default=None, 
+        help="Path to SSL private key file (required if --ssl is enabled)"
+    )
     return parser.parse_args()
 
 
@@ -284,6 +306,7 @@ def create_app(args):
         llm_model_max_token_size=args.max_tokens,
         llm_model_kwargs={
             "host": args.llm_binding_host,
+            "timeout":args.timeout
             "options": {"num_ctx": args.max_tokens},
         },
         embedding_func=EmbeddingFunc(
diff --git a/lightrag/llm.py b/lightrag/llm.py
index 0c17019a..4e01dd51 100644
--- a/lightrag/llm.py
+++ b/lightrag/llm.py
@@ -336,6 +336,7 @@ async def hf_model_if_cache(
         (RateLimitError, APIConnectionError, APITimeoutError)
     ),
 )
+
 async def ollama_model_if_cache(
     model,
     prompt,
@@ -406,8 +407,9 @@ 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
 
-    async with aiohttp.ClientSession() as session:
+    async with aiohttp.ClientSession(timeout=timeout) as session:
         if stream:
 
             async def inner():