This commit is contained in:
Timothy Jaeryang Baek
2026-05-31 18:42:28 -07:00
parent ee47c9c833
commit 7139797be0
+23 -6
View File
@@ -785,6 +785,7 @@ async def embed(
log.info(f'generate_ollama_batch_embeddings {form_data}')
await check_model_access(user, await Models.get_model_by_id(form_data.model), BYPASS_MODEL_ACCESS_CONTROL)
await validate_ollama_backend_idx(request, form_data.model, url_idx, user)
if url_idx is None:
model = form_data.model
@@ -837,6 +838,7 @@ async def embeddings(
log.info(f'generate_ollama_embeddings {form_data}')
await check_model_access(user, await Models.get_model_by_id(form_data.model), BYPASS_MODEL_ACCESS_CONTROL)
await validate_ollama_backend_idx(request, form_data.model, url_idx, user)
if url_idx is None:
model = form_data.model
@@ -896,6 +898,7 @@ async def generate_completion(
raise HTTPException(status_code=503, detail=ERROR_MESSAGES.OLLAMA_API_DISABLED)
await check_model_access(user, await Models.get_model_by_id(form_data.model), BYPASS_MODEL_ACCESS_CONTROL)
await validate_ollama_backend_idx(request, form_data.model, url_idx, user)
if url_idx is None:
await get_all_models(request, user=user)
@@ -954,7 +957,21 @@ class GenerateChatCompletionForm(BaseModel):
model_config = ConfigDict(extra='allow')
async def get_ollama_url(request: Request, model: str, url_idx: int | None = None):
async def validate_ollama_backend_idx(request: Request, model: str, url_idx: int | None, user) -> None:
# A caller-supplied url_idx must point to a backend the model is actually
# served from; the None path is already constrained to that allow-list.
if url_idx is None or user is None or getattr(user, 'role', None) == 'admin' or BYPASS_MODEL_ACCESS_CONTROL:
return
models = request.app.state.OLLAMA_MODELS
if not models or model not in models:
await get_all_models(request, user=user)
models = request.app.state.OLLAMA_MODELS
if url_idx not in (models.get(model) or {}).get('urls', []):
raise HTTPException(status_code=403, detail=ERROR_MESSAGES.ACCESS_PROHIBITED)
async def get_ollama_url(request: Request, model: str, url_idx: int | None = None, user=None):
await validate_ollama_backend_idx(request, model, url_idx, user)
if url_idx is None:
models = request.app.state.OLLAMA_MODELS
if model not in models:
@@ -1026,7 +1043,7 @@ async def generate_chat_completion(
else:
await check_model_access(user, None, bypass_filter)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx, user)
api_config = _resolve_api_config(request, url_idx, url)
prefix_id = api_config.get('prefix_id')
@@ -1112,7 +1129,7 @@ async def generate_openai_completion(
else:
await check_model_access(user, None)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx, user)
api_config = _resolve_api_config(request, url_idx, url)
prefix_id = api_config.get('prefix_id')
@@ -1169,7 +1186,7 @@ async def generate_openai_chat_completion(
else:
await check_model_access(user, None)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx, user)
api_config = _resolve_api_config(request, url_idx, url)
prefix_id = api_config.get('prefix_id')
@@ -1218,7 +1235,7 @@ async def generate_anthropic_messages(
else:
await check_model_access(user, None)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx, user)
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(
str(url_idx),
request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support
@@ -1276,7 +1293,7 @@ async def generate_responses(
else:
await check_model_access(user, None)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx)
url, url_idx = await get_ollama_url(request, payload['model'], url_idx, user)
api_config = request.app.state.config.OLLAMA_API_CONFIGS.get(
str(url_idx),
request.app.state.config.OLLAMA_API_CONFIGS.get(url, {}), # Legacy support