mirror of
https://github.com/open-webui/open-webui.git
synced 2025-12-13 12:55:19 +00:00
506 lines
17 KiB
Python
506 lines
17 KiB
Python
import time
|
||
import logging
|
||
import sys
|
||
|
||
from aiocache import cached
|
||
from typing import Any, Optional
|
||
import random
|
||
import json
|
||
import inspect
|
||
import uuid
|
||
import asyncio
|
||
|
||
from fastapi import Request, status
|
||
from starlette.responses import Response, StreamingResponse, JSONResponse
|
||
|
||
|
||
from open_webui.models.users import UserModel
|
||
|
||
from open_webui.socket.main import (
|
||
sio,
|
||
get_event_call,
|
||
get_event_emitter,
|
||
)
|
||
from open_webui.functions import generate_function_chat_completion
|
||
|
||
from open_webui.routers.openai import (
|
||
generate_chat_completion as generate_openai_chat_completion,
|
||
)
|
||
|
||
from open_webui.routers.ollama import (
|
||
generate_chat_completion as generate_ollama_chat_completion,
|
||
)
|
||
|
||
from open_webui.routers.pipelines import (
|
||
process_pipeline_inlet_filter,
|
||
process_pipeline_outlet_filter,
|
||
)
|
||
|
||
from open_webui.models.functions import Functions
|
||
from open_webui.models.models import Models
|
||
|
||
|
||
from open_webui.utils.plugin import (
|
||
load_function_module_by_id,
|
||
get_function_module_from_cache,
|
||
)
|
||
from open_webui.utils.models import get_all_models, check_model_access
|
||
from open_webui.utils.payload import convert_payload_openai_to_ollama
|
||
from open_webui.utils.response import (
|
||
convert_response_ollama_to_openai,
|
||
convert_streaming_response_ollama_to_openai,
|
||
)
|
||
from open_webui.utils.filter import (
|
||
get_sorted_filter_ids,
|
||
process_filter_functions,
|
||
)
|
||
|
||
from open_webui.env import SRC_LOG_LEVELS, GLOBAL_LOG_LEVEL, BYPASS_MODEL_ACCESS_CONTROL
|
||
|
||
|
||
logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
|
||
log = logging.getLogger(__name__)
|
||
log.setLevel(SRC_LOG_LEVELS["MAIN"])
|
||
|
||
|
||
async def generate_direct_chat_completion(
|
||
request: Request,
|
||
form_data: dict,
|
||
user: Any,
|
||
models: dict,
|
||
):
|
||
log.info("generate_direct_chat_completion")
|
||
|
||
metadata = form_data.pop("metadata", {})
|
||
|
||
user_id = metadata.get("user_id")
|
||
session_id = metadata.get("session_id")
|
||
request_id = str(uuid.uuid4()) # Generate a unique request ID
|
||
|
||
event_caller = get_event_call(metadata)
|
||
|
||
channel = f"{user_id}:{session_id}:{request_id}"
|
||
logging.info(f"WebSocket channel: {channel}")
|
||
|
||
if form_data.get("stream"):
|
||
q = asyncio.Queue()
|
||
|
||
async def message_listener(sid, data):
|
||
"""
|
||
Handle received socket messages and push them into the queue.
|
||
"""
|
||
await q.put(data)
|
||
|
||
# Register the listener
|
||
sio.on(channel, message_listener)
|
||
|
||
# Start processing chat completion in background
|
||
res = await event_caller(
|
||
{
|
||
"type": "request:chat:completion",
|
||
"data": {
|
||
"form_data": form_data,
|
||
"model": models[form_data["model"]],
|
||
"channel": channel,
|
||
"session_id": session_id,
|
||
},
|
||
}
|
||
)
|
||
|
||
log.info(f"res: {res}")
|
||
|
||
if res.get("status", False):
|
||
# Define a generator to stream responses
|
||
async def event_generator():
|
||
nonlocal q
|
||
try:
|
||
while True:
|
||
data = await q.get() # Wait for new messages
|
||
if isinstance(data, dict):
|
||
if "done" in data and data["done"]:
|
||
break # Stop streaming when 'done' is received
|
||
|
||
yield f"data: {json.dumps(data)}\n\n"
|
||
elif isinstance(data, str):
|
||
if "data:" in data:
|
||
yield f"{data}\n\n"
|
||
else:
|
||
yield f"data: {data}\n\n"
|
||
except Exception as e:
|
||
log.debug(f"Error in event generator: {e}")
|
||
pass
|
||
|
||
# Define a background task to run the event generator
|
||
async def background():
|
||
try:
|
||
del sio.handlers["/"][channel]
|
||
except Exception as e:
|
||
pass
|
||
|
||
# Return the streaming response
|
||
return StreamingResponse(
|
||
event_generator(), media_type="text/event-stream", background=background
|
||
)
|
||
else:
|
||
raise Exception(str(res))
|
||
else:
|
||
res = await event_caller(
|
||
{
|
||
"type": "request:chat:completion",
|
||
"data": {
|
||
"form_data": form_data,
|
||
"model": models[form_data["model"]],
|
||
"channel": channel,
|
||
"session_id": session_id,
|
||
},
|
||
}
|
||
)
|
||
|
||
if "error" in res and res["error"]:
|
||
raise Exception(res["error"])
|
||
|
||
return res
|
||
|
||
|
||
async def generate_chat_completion(
|
||
request: Request,
|
||
form_data: dict,
|
||
user: Any,
|
||
bypass_filter: bool = False,
|
||
chatting_completion: bool = False
|
||
):
|
||
"""
|
||
聊天完成生成函数 - 根据模型类型分发到不同的底层 API 处理器
|
||
|
||
这是聊天完成的核心路由函数，负责：
|
||
1. 验证模型存在性和用户权限
|
||
2. 处理 Direct 模式（直连外部 API）
|
||
3. 处理 Arena 模式（随机选择模型进行对比）
|
||
4. 根据模型类型分发到对应处理器：
|
||
- Pipe: Pipeline 插件函数
|
||
- Ollama: Ollama 本地模型
|
||
- OpenAI: OpenAI 兼容 API (含 Claude, Gemini 等)
|
||
|
||
Args:
|
||
request: FastAPI Request 对象
|
||
form_data: OpenAI 格式的聊天请求数据
|
||
user: 用户对象
|
||
bypass_filter: 是否绕过权限和 Pipeline Filter 检查
|
||
|
||
Returns:
|
||
- 流式: StreamingResponse (SSE 格式)
|
||
- 非流式: dict (OpenAI 兼容格式)
|
||
|
||
Raises:
|
||
Exception: 模型不存在或无权限访问
|
||
"""
|
||
log.debug(f"generate_chat_completion: {form_data}")
|
||
|
||
# === 1. 权限检查配置 ===
|
||
if BYPASS_MODEL_ACCESS_CONTROL:
|
||
bypass_filter = True # 全局配置：绕过所有权限检查
|
||
|
||
# === 2. 合并元数据 ===
|
||
# 从 request.state.metadata 获取上游传递的元数据（chat_id, user_id 等）
|
||
if hasattr(request.state, "metadata"):
|
||
if "metadata" not in form_data:
|
||
form_data["metadata"] = request.state.metadata
|
||
else:
|
||
# 合并，request.state.metadata 优先级更高
|
||
form_data["metadata"] = {
|
||
**form_data["metadata"],
|
||
**request.state.metadata,
|
||
}
|
||
|
||
# === 3. 确定模型列表来源 ===
|
||
if getattr(request.state, "direct", False) and hasattr(request.state, "model"):
|
||
# Direct 模式：用户直接提供外部 API 配置（如 OpenAI API Key）
|
||
models = {
|
||
request.state.model["id"]: request.state.model,
|
||
}
|
||
log.debug(f"direct connection to model: {models}")
|
||
else:
|
||
# 标准模式：使用平台内置模型列表
|
||
models = request.app.state.MODELS
|
||
|
||
# === 4. 验证模型存在性 ===
|
||
model_id = form_data["model"]
|
||
if model_id not in models:
|
||
raise Exception("Model not found")
|
||
|
||
model = models[model_id]
|
||
|
||
# === 5. Direct 模式分支：直连外部 API ===
|
||
if getattr(request.state, "direct", False):
|
||
return await generate_direct_chat_completion(
|
||
request, form_data, user=user, models=models
|
||
)
|
||
else:
|
||
# === 6. 标准模式：检查用户权限 ===
|
||
if not bypass_filter and user.role == "user":
|
||
try:
|
||
check_model_access(user, model) # 验证 RBAC 权限
|
||
except Exception as e:
|
||
raise e
|
||
|
||
# === 7. Arena 模式：随机选择模型进行盲测对比 ===
|
||
if False:
|
||
if model.get("owned_by") == "arena":
|
||
# 获取候选模型列表
|
||
model_ids = model.get("info", {}).get("meta", {}).get("model_ids")
|
||
filter_mode = model.get("info", {}).get("meta", {}).get("filter_mode")
|
||
|
||
# 如果是排除模式，反选模型列表
|
||
if model_ids and filter_mode == "exclude":
|
||
model_ids = [
|
||
model["id"]
|
||
for model in list(request.app.state.MODELS.values())
|
||
if model.get("owned_by") != "arena" and model["id"] not in model_ids
|
||
]
|
||
|
||
# 随机选择一个模型
|
||
selected_model_id = None
|
||
if isinstance(model_ids, list) and model_ids:
|
||
selected_model_id = random.choice(model_ids)
|
||
else:
|
||
# 未指定则从所有非 Arena 模型中随机选择
|
||
model_ids = [
|
||
model["id"]
|
||
for model in list(request.app.state.MODELS.values())
|
||
if model.get("owned_by") != "arena"
|
||
]
|
||
selected_model_id = random.choice(model_ids)
|
||
|
||
# 替换模型 ID
|
||
form_data["model"] = selected_model_id
|
||
|
||
# 流式响应：在首个 chunk 中注入 selected_model_id
|
||
if form_data.get("stream") == True:
|
||
|
||
async def stream_wrapper(stream):
|
||
"""在流式响应前添加选中的模型 ID"""
|
||
yield f"data: {json.dumps({'selected_model_id': selected_model_id})}\n\n"
|
||
async for chunk in stream:
|
||
yield chunk
|
||
|
||
# 递归调用自身，绕过 Arena 逻辑
|
||
response = await generate_chat_completion(
|
||
request, form_data, user, bypass_filter=True
|
||
)
|
||
return StreamingResponse(
|
||
stream_wrapper(response.body_iterator),
|
||
media_type="text/event-stream",
|
||
background=response.background,
|
||
)
|
||
else:
|
||
# 非流式响应：直接在结果中添加 selected_model_id
|
||
return {
|
||
**(
|
||
await generate_chat_completion(
|
||
request, form_data, user, bypass_filter=True
|
||
)
|
||
),
|
||
"selected_model_id": selected_model_id,
|
||
}
|
||
|
||
# === 8. Pipeline 模式：调用自定义 Python 函数 ===
|
||
if False:
|
||
if model.get("pipe"):
|
||
return await generate_function_chat_completion(
|
||
request, form_data, user=user, models=models
|
||
)
|
||
|
||
# === 9. Ollama 模式：调用本地 Ollama 服务 ===
|
||
if False:
|
||
if model.get("owned_by") == "ollama":
|
||
# 转换 OpenAI 格式 → Ollama 格式
|
||
form_data = convert_payload_openai_to_ollama(form_data)
|
||
response = await generate_ollama_chat_completion(
|
||
request=request,
|
||
form_data=form_data,
|
||
user=user,
|
||
bypass_filter=bypass_filter,
|
||
)
|
||
|
||
# 流式响应：转换 Ollama SSE → OpenAI SSE
|
||
if form_data.get("stream"):
|
||
response.headers["content-type"] = "text/event-stream"
|
||
return StreamingResponse(
|
||
convert_streaming_response_ollama_to_openai(response),
|
||
headers=dict(response.headers),
|
||
background=response.background,
|
||
)
|
||
else:
|
||
# 非流式响应：转换 Ollama JSON → OpenAI JSON
|
||
return convert_response_ollama_to_openai(response)
|
||
|
||
# === 10. OpenAI 兼容模式：调用 OpenAI API 或兼容服务 ===
|
||
# >>>
|
||
return await generate_openai_chat_completion(
|
||
request=request,
|
||
form_data=form_data,
|
||
user=user,
|
||
bypass_filter=bypass_filter,
|
||
chatting_completion = chatting_completion
|
||
)
|
||
|
||
|
||
chat_completion = generate_chat_completion
|
||
|
||
|
||
async def chat_completed(request: Request, form_data: dict, user: Any):
|
||
if not request.app.state.MODELS:
|
||
await get_all_models(request, user=user)
|
||
|
||
if getattr(request.state, "direct", False) and hasattr(request.state, "model"):
|
||
models = {
|
||
request.state.model["id"]: request.state.model,
|
||
}
|
||
else:
|
||
models = request.app.state.MODELS
|
||
|
||
data = form_data
|
||
model_id = data["model"]
|
||
if model_id not in models:
|
||
raise Exception("Model not found")
|
||
|
||
model = models[model_id]
|
||
|
||
try:
|
||
data = await process_pipeline_outlet_filter(request, data, user, models)
|
||
except Exception as e:
|
||
return Exception(f"Error: {e}")
|
||
|
||
metadata = {
|
||
"chat_id": data["chat_id"],
|
||
"message_id": data["id"],
|
||
"filter_ids": data.get("filter_ids", []),
|
||
"session_id": data["session_id"],
|
||
"user_id": user.id,
|
||
}
|
||
|
||
extra_params = {
|
||
"__event_emitter__": get_event_emitter(metadata),
|
||
"__event_call__": get_event_call(metadata),
|
||
"__user__": user.model_dump() if isinstance(user, UserModel) else {},
|
||
"__metadata__": metadata,
|
||
"__request__": request,
|
||
"__model__": model,
|
||
}
|
||
|
||
try:
|
||
filter_functions = [
|
||
Functions.get_function_by_id(filter_id)
|
||
for filter_id in get_sorted_filter_ids(
|
||
request, model, metadata.get("filter_ids", [])
|
||
)
|
||
]
|
||
|
||
result, _ = await process_filter_functions(
|
||
request=request,
|
||
filter_functions=filter_functions,
|
||
filter_type="outlet",
|
||
form_data=data,
|
||
extra_params=extra_params,
|
||
)
|
||
return result
|
||
except Exception as e:
|
||
return Exception(f"Error: {e}")
|
||
|
||
|
||
async def chat_action(request: Request, action_id: str, form_data: dict, user: Any):
|
||
if "." in action_id:
|
||
action_id, sub_action_id = action_id.split(".")
|
||
else:
|
||
sub_action_id = None
|
||
|
||
action = Functions.get_function_by_id(action_id)
|
||
if not action:
|
||
raise Exception(f"Action not found: {action_id}")
|
||
|
||
if not request.app.state.MODELS:
|
||
await get_all_models(request, user=user)
|
||
|
||
if getattr(request.state, "direct", False) and hasattr(request.state, "model"):
|
||
models = {
|
||
request.state.model["id"]: request.state.model,
|
||
}
|
||
else:
|
||
models = request.app.state.MODELS
|
||
|
||
data = form_data
|
||
model_id = data["model"]
|
||
|
||
if model_id not in models:
|
||
raise Exception("Model not found")
|
||
model = models[model_id]
|
||
|
||
__event_emitter__ = get_event_emitter(
|
||
{
|
||
"chat_id": data["chat_id"],
|
||
"message_id": data["id"],
|
||
"session_id": data["session_id"],
|
||
"user_id": user.id,
|
||
}
|
||
)
|
||
__event_call__ = get_event_call(
|
||
{
|
||
"chat_id": data["chat_id"],
|
||
"message_id": data["id"],
|
||
"session_id": data["session_id"],
|
||
"user_id": user.id,
|
||
}
|
||
)
|
||
|
||
function_module, _, _ = get_function_module_from_cache(request, action_id)
|
||
|
||
if hasattr(function_module, "valves") and hasattr(function_module, "Valves"):
|
||
valves = Functions.get_function_valves_by_id(action_id)
|
||
function_module.valves = function_module.Valves(**(valves if valves else {}))
|
||
|
||
if hasattr(function_module, "action"):
|
||
try:
|
||
action = function_module.action
|
||
|
||
# Get the signature of the function
|
||
sig = inspect.signature(action)
|
||
params = {"body": data}
|
||
|
||
# Extra parameters to be passed to the function
|
||
extra_params = {
|
||
"__model__": model,
|
||
"__id__": sub_action_id if sub_action_id is not None else action_id,
|
||
"__event_emitter__": __event_emitter__,
|
||
"__event_call__": __event_call__,
|
||
"__request__": request,
|
||
}
|
||
|
||
# Add extra params in contained in function signature
|
||
for key, value in extra_params.items():
|
||
if key in sig.parameters:
|
||
params[key] = value
|
||
|
||
if "__user__" in sig.parameters:
|
||
__user__ = user.model_dump() if isinstance(user, UserModel) else {}
|
||
|
||
try:
|
||
if hasattr(function_module, "UserValves"):
|
||
__user__["valves"] = function_module.UserValves(
|
||
**Functions.get_user_valves_by_id_and_user_id(
|
||
action_id, user.id
|
||
)
|
||
)
|
||
except Exception as e:
|
||
log.exception(f"Failed to get user values: {e}")
|
||
|
||
params = {**params, "__user__": __user__}
|
||
|
||
if inspect.iscoroutinefunction(action):
|
||
data = await action(**params)
|
||
else:
|
||
data = action(**params)
|
||
|
||
except Exception as e:
|
||
return Exception(f"Error: {e}")
|
||
|
||
return data
|