From 2985d88f976ab63b6ce31d1c9929506e37c288ae Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Sat, 15 Feb 2025 11:39:10 +0800
Subject: [PATCH] refactor: improve CORS and streaming response headers

- Add configurable CORS origins
- Remove duplicate CORS headers
- Add X-Accel-Buffering header
- Update env example file
- Clean up header configurations
---
 .env.example                    | 13 +++++++------
 lightrag/api/lightrag_server.py | 16 +++++++++++-----
 lightrag/api/ollama_api.py      |  8 ++------
 3 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/.env.example b/.env.example
index 022bd63d..2701335a 100644
--- a/.env.example
+++ b/.env.example
@@ -1,12 +1,13 @@
 ### Server Configuration
-#HOST=0.0.0.0
-#PORT=9621
-#NAMESPACE_PREFIX=lightrag  # separating data from difference Lightrag instances
+# HOST=0.0.0.0
+# PORT=9621
+# NAMESPACE_PREFIX=lightrag  # separating data from difference Lightrag instances
+# CORS_ORIGINS=http://localhost:3000,http://localhost:8080
 
 ### Optional SSL Configuration
-#SSL=true
-#SSL_CERTFILE=/path/to/cert.pem
-#SSL_KEYFILE=/path/to/key.pem
+# SSL=true
+# SSL_CERTFILE=/path/to/cert.pem
+# SSL_KEYFILE=/path/to/key.pem
 
 ### Security (empty for no api-key is needed)
 # LIGHTRAG_API_KEY=your-secure-api-key-here
diff --git a/lightrag/api/lightrag_server.py b/lightrag/api/lightrag_server.py
index ce182bc1..19552faf 100644
--- a/lightrag/api/lightrag_server.py
+++ b/lightrag/api/lightrag_server.py
@@ -847,10 +847,19 @@ def create_app(args):
         lifespan=lifespan,
     )
 
+    def get_cors_origins():
+        """Get allowed origins from environment variable
+        Returns a list of allowed origins, defaults to ["*"] if not set
+        """
+        origins_str = os.getenv("CORS_ORIGINS", "*")
+        if origins_str == "*":
+            return ["*"]
+        return [origin.strip() for origin in origins_str.split(",")]
+
     # Add CORS middleware
     app.add_middleware(
         CORSMiddleware,
-        allow_origins=["*"],
+        allow_origins=get_cors_origins(),
         allow_credentials=True,
         allow_methods=["*"],
         allow_headers=["*"],
@@ -1377,10 +1386,7 @@ def create_app(args):
                     "Cache-Control": "no-cache",
                     "Connection": "keep-alive",
                     "Content-Type": "application/x-ndjson",
-                    "Access-Control-Allow-Origin": "*",
-                    "Access-Control-Allow-Methods": "POST, OPTIONS",
-                    "Access-Control-Allow-Headers": "Content-Type",
-                    "X-Accel-Buffering": "no",  # Disable Nginx buffering
+                    "X-Accel-Buffering": "no",  # 确保在Nginx代理时正确处理流式响应
                 },
             )
         except Exception as e:
diff --git a/lightrag/api/ollama_api.py b/lightrag/api/ollama_api.py
index 01a883ca..94703dee 100644
--- a/lightrag/api/ollama_api.py
+++ b/lightrag/api/ollama_api.py
@@ -316,9 +316,7 @@ class OllamaAPI:
                             "Cache-Control": "no-cache",
                             "Connection": "keep-alive",
                             "Content-Type": "application/x-ndjson",
-                            "Access-Control-Allow-Origin": "*",
-                            "Access-Control-Allow-Methods": "POST, OPTIONS",
-                            "Access-Control-Allow-Headers": "Content-Type",
+                            "X-Accel-Buffering": "no",  # 确保在Nginx代理时正确处理流式响应
                         },
                     )
                 else:
@@ -534,9 +532,7 @@ class OllamaAPI:
                             "Cache-Control": "no-cache",
                             "Connection": "keep-alive",
                             "Content-Type": "application/x-ndjson",
-                            "Access-Control-Allow-Origin": "*",
-                            "Access-Control-Allow-Methods": "POST, OPTIONS",
-                            "Access-Control-Allow-Headers": "Content-Type",
+                            "X-Accel-Buffering": "no",  # 确保在Nginx代理时正确处理流式响应
                         },
                     )
                 else: