Переглянути джерело

Merge pull request #13031 from gaby/anyio-workers

feat: Add support for configuring FastAPI/AnyIO Thread Pool
Tim Jaeryang Baek 3 місяців тому
батько
коміт
d2f3c82fd6

+ 3 - 0
backend/open_webui/config.py

@@ -1203,6 +1203,9 @@ ENABLE_USER_WEBHOOKS = PersistentConfig(
     os.environ.get("ENABLE_USER_WEBHOOKS", "True").lower() == "true",
 )
 
+# FastAPI / AnyIO settings
+THREAD_POOL_SIZE = int(os.getenv("THREAD_POOL_SIZE", "0"))
+
 
 def validate_cors_origins(origins):
     for origin in origins:

+ 8 - 0
backend/open_webui/main.py

@@ -17,6 +17,7 @@ from sqlalchemy import text
 from typing import Optional
 from aiocache import cached
 import aiohttp
+import anyio.to_thread
 import requests
 
 
@@ -106,6 +107,8 @@ from open_webui.config import (
     OPENAI_API_CONFIGS,
     # Direct Connections
     ENABLE_DIRECT_CONNECTIONS,
+    # Thread pool size for FastAPI/AnyIO
+    THREAD_POOL_SIZE,
     # Tool Server Configs
     TOOL_SERVER_CONNECTIONS,
     # Code Execution
@@ -434,6 +437,11 @@ async def lifespan(app: FastAPI):
     if LICENSE_KEY:
         get_license_data(app, LICENSE_KEY)
 
+    pool_size = THREAD_POOL_SIZE
+    if pool_size and pool_size > 0:
+        limiter = anyio.to_thread.current_default_thread_limiter()
+        limiter.total_tokens = pool_size
+
     asyncio.create_task(periodic_usage_pool_cleanup())
     yield
 

+ 2 - 1
backend/open_webui/routers/knowledge.py

@@ -742,5 +742,6 @@ def add_files_to_knowledge_batch(
         )
 
     return KnowledgeFilesResponse(
-        **knowledge.model_dump(), files=Files.get_file_metadatas_by_ids(existing_file_ids)
+        **knowledge.model_dump(),
+        files=Files.get_file_metadatas_by_ids(existing_file_ids),
     )

+ 8 - 4
backend/open_webui/routers/users.py

@@ -298,8 +298,10 @@ async def update_user_by_id(
             )
     except Exception as e:
         log.error(f"Error checking primary admin status: {e}")
-        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Could not verify primary admin status.")
-
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Could not verify primary admin status.",
+        )
 
     user = Users.get_user_by_id(user_id)
 
@@ -341,7 +343,6 @@ async def update_user_by_id(
     )
 
 
-
 ############################
 # DeleteUserById
 ############################
@@ -359,7 +360,10 @@ async def delete_user_by_id(user_id: str, user=Depends(get_admin_user)):
             )
     except Exception as e:
         log.error(f"Error checking primary admin status: {e}")
-        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Could not verify primary admin status.")
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail="Could not verify primary admin status.",
+        )
 
     if user.id != user_id:
         result = Auths.delete_auth_by_id(user_id)