refac
This commit is contained in:
+141
-170
@@ -41,22 +41,22 @@ async def fetch_ollama_models(request: Request, user: UserModel = None):
|
||||
raw_ollama_models = await ollama.get_all_models(request, user=user)
|
||||
return [
|
||||
{
|
||||
"id": model["model"],
|
||||
"name": model["name"],
|
||||
"object": "model",
|
||||
"created": int(time.time()),
|
||||
"owned_by": "ollama",
|
||||
"ollama": model,
|
||||
"connection_type": model.get("connection_type", "local"),
|
||||
"tags": model.get("tags", []),
|
||||
'id': model['model'],
|
||||
'name': model['name'],
|
||||
'object': 'model',
|
||||
'created': int(time.time()),
|
||||
'owned_by': 'ollama',
|
||||
'ollama': model,
|
||||
'connection_type': model.get('connection_type', 'local'),
|
||||
'tags': model.get('tags', []),
|
||||
}
|
||||
for model in raw_ollama_models["models"]
|
||||
for model in raw_ollama_models['models']
|
||||
]
|
||||
|
||||
|
||||
async def fetch_openai_models(request: Request, user: UserModel = None):
|
||||
openai_response = await openai.get_all_models(request, user=user)
|
||||
return openai_response["data"]
|
||||
return openai_response['data']
|
||||
|
||||
|
||||
async def get_all_base_models(request: Request, user: UserModel = None):
|
||||
@@ -72,9 +72,7 @@ async def get_all_base_models(request: Request, user: UserModel = None):
|
||||
)
|
||||
function_task = get_function_models(request)
|
||||
|
||||
openai_models, ollama_models, function_models = await asyncio.gather(
|
||||
openai_task, ollama_task, function_task
|
||||
)
|
||||
openai_models, ollama_models, function_models = await asyncio.gather(openai_task, ollama_task, function_task)
|
||||
|
||||
return function_models + openai_models + ollama_models
|
||||
|
||||
@@ -103,15 +101,15 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
if len(request.app.state.config.EVALUATION_ARENA_MODELS) > 0:
|
||||
arena_models = [
|
||||
{
|
||||
"id": model["id"],
|
||||
"name": model["name"],
|
||||
"info": {
|
||||
"meta": model["meta"],
|
||||
'id': model['id'],
|
||||
'name': model['name'],
|
||||
'info': {
|
||||
'meta': model['meta'],
|
||||
},
|
||||
"object": "model",
|
||||
"created": int(time.time()),
|
||||
"owned_by": "arena",
|
||||
"arena": True,
|
||||
'object': 'model',
|
||||
'created': int(time.time()),
|
||||
'owned_by': 'arena',
|
||||
'arena': True,
|
||||
}
|
||||
for model in request.app.state.config.EVALUATION_ARENA_MODELS
|
||||
]
|
||||
@@ -119,45 +117,35 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
# Add default arena model
|
||||
arena_models = [
|
||||
{
|
||||
"id": DEFAULT_ARENA_MODEL["id"],
|
||||
"name": DEFAULT_ARENA_MODEL["name"],
|
||||
"info": {
|
||||
"meta": DEFAULT_ARENA_MODEL["meta"],
|
||||
'id': DEFAULT_ARENA_MODEL['id'],
|
||||
'name': DEFAULT_ARENA_MODEL['name'],
|
||||
'info': {
|
||||
'meta': DEFAULT_ARENA_MODEL['meta'],
|
||||
},
|
||||
"object": "model",
|
||||
"created": int(time.time()),
|
||||
"owned_by": "arena",
|
||||
"arena": True,
|
||||
'object': 'model',
|
||||
'created': int(time.time()),
|
||||
'owned_by': 'arena',
|
||||
'arena': True,
|
||||
}
|
||||
]
|
||||
models = models + arena_models
|
||||
|
||||
global_action_ids = [
|
||||
function.id for function in Functions.get_global_action_functions()
|
||||
]
|
||||
enabled_action_ids = [
|
||||
function.id
|
||||
for function in Functions.get_functions_by_type("action", active_only=True)
|
||||
]
|
||||
global_action_ids = [function.id for function in Functions.get_global_action_functions()]
|
||||
enabled_action_ids = [function.id for function in Functions.get_functions_by_type('action', active_only=True)]
|
||||
|
||||
global_filter_ids = [
|
||||
function.id for function in Functions.get_global_filter_functions()
|
||||
]
|
||||
enabled_filter_ids = [
|
||||
function.id
|
||||
for function in Functions.get_functions_by_type("filter", active_only=True)
|
||||
]
|
||||
global_filter_ids = [function.id for function in Functions.get_global_filter_functions()]
|
||||
enabled_filter_ids = [function.id for function in Functions.get_functions_by_type('filter', active_only=True)]
|
||||
|
||||
custom_models = Models.get_all_models()
|
||||
|
||||
# Single O(1) lookup: Ollama base names first, then exact IDs (exact wins).
|
||||
base_model_lookup = {}
|
||||
for model in models:
|
||||
if model.get("owned_by") == "ollama":
|
||||
base_model_lookup.setdefault(model["id"].split(":")[0], model)
|
||||
base_model_lookup[model["id"]] = model
|
||||
if model.get('owned_by') == 'ollama':
|
||||
base_model_lookup.setdefault(model['id'].split(':')[0], model)
|
||||
base_model_lookup[model['id']] = model
|
||||
|
||||
existing_ids = {m["id"] for m in models}
|
||||
existing_ids = {m['id'] for m in models}
|
||||
|
||||
for custom_model in custom_models:
|
||||
if custom_model.base_model_id is None:
|
||||
@@ -166,26 +154,22 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
|
||||
if model:
|
||||
if custom_model.is_active:
|
||||
model["name"] = custom_model.name
|
||||
model["info"] = custom_model.model_dump()
|
||||
model['name'] = custom_model.name
|
||||
model['info'] = custom_model.model_dump()
|
||||
|
||||
action_ids = []
|
||||
filter_ids = []
|
||||
|
||||
if "info" in model:
|
||||
if "meta" in model["info"]:
|
||||
action_ids.extend(
|
||||
model["info"]["meta"].get("actionIds", [])
|
||||
)
|
||||
filter_ids.extend(
|
||||
model["info"]["meta"].get("filterIds", [])
|
||||
)
|
||||
if 'info' in model:
|
||||
if 'meta' in model['info']:
|
||||
action_ids.extend(model['info']['meta'].get('actionIds', []))
|
||||
filter_ids.extend(model['info']['meta'].get('filterIds', []))
|
||||
|
||||
if "params" in model["info"]:
|
||||
del model["info"]["params"]
|
||||
if 'params' in model['info']:
|
||||
del model['info']['params']
|
||||
|
||||
model["action_ids"] = action_ids
|
||||
model["filter_ids"] = filter_ids
|
||||
model['action_ids'] = action_ids
|
||||
model['filter_ids'] = filter_ids
|
||||
else:
|
||||
models.remove(model)
|
||||
|
||||
@@ -193,38 +177,36 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
if custom_model.id in existing_ids:
|
||||
continue
|
||||
|
||||
owned_by = "openai"
|
||||
owned_by = 'openai'
|
||||
connection_type = None
|
||||
pipe = None
|
||||
|
||||
base_model = base_model_lookup.get(custom_model.base_model_id)
|
||||
if base_model is None:
|
||||
base_model = base_model_lookup.get(
|
||||
custom_model.base_model_id.split(":")[0]
|
||||
)
|
||||
base_model = base_model_lookup.get(custom_model.base_model_id.split(':')[0])
|
||||
if base_model:
|
||||
owned_by = base_model.get("owned_by", "unknown")
|
||||
if "pipe" in base_model:
|
||||
pipe = base_model["pipe"]
|
||||
connection_type = base_model.get("connection_type", None)
|
||||
owned_by = base_model.get('owned_by', 'unknown')
|
||||
if 'pipe' in base_model:
|
||||
pipe = base_model['pipe']
|
||||
connection_type = base_model.get('connection_type', None)
|
||||
|
||||
model = {
|
||||
"id": f"{custom_model.id}",
|
||||
"name": custom_model.name,
|
||||
"object": "model",
|
||||
"created": custom_model.created_at,
|
||||
"owned_by": owned_by,
|
||||
"connection_type": connection_type,
|
||||
"preset": True,
|
||||
**({"pipe": pipe} if pipe is not None else {}),
|
||||
'id': f'{custom_model.id}',
|
||||
'name': custom_model.name,
|
||||
'object': 'model',
|
||||
'created': custom_model.created_at,
|
||||
'owned_by': owned_by,
|
||||
'connection_type': connection_type,
|
||||
'preset': True,
|
||||
**({'pipe': pipe} if pipe is not None else {}),
|
||||
}
|
||||
|
||||
info = custom_model.model_dump()
|
||||
if "params" in info:
|
||||
if 'params' in info:
|
||||
# Remove params to avoid exposing sensitive info
|
||||
del info["params"]
|
||||
del info['params']
|
||||
|
||||
model["info"] = info
|
||||
model['info'] = info
|
||||
|
||||
action_ids = []
|
||||
filter_ids = []
|
||||
@@ -232,32 +214,32 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
if custom_model.meta:
|
||||
meta = custom_model.meta.model_dump()
|
||||
|
||||
if "actionIds" in meta:
|
||||
action_ids.extend(meta["actionIds"])
|
||||
if 'actionIds' in meta:
|
||||
action_ids.extend(meta['actionIds'])
|
||||
|
||||
if "filterIds" in meta:
|
||||
filter_ids.extend(meta["filterIds"])
|
||||
if 'filterIds' in meta:
|
||||
filter_ids.extend(meta['filterIds'])
|
||||
|
||||
model["action_ids"] = action_ids
|
||||
model["filter_ids"] = filter_ids
|
||||
model['action_ids'] = action_ids
|
||||
model['filter_ids'] = filter_ids
|
||||
|
||||
models.append(model)
|
||||
|
||||
# Process action_ids to get the actions
|
||||
def get_action_items_from_module(function, module):
|
||||
actions = []
|
||||
if hasattr(module, "actions"):
|
||||
if hasattr(module, 'actions'):
|
||||
actions = module.actions
|
||||
return [
|
||||
{
|
||||
"id": f"{function.id}.{action['id']}",
|
||||
"name": action.get("name", f"{function.name} ({action['id']})"),
|
||||
"description": function.meta.description,
|
||||
"icon": action.get(
|
||||
"icon_url",
|
||||
function.meta.manifest.get("icon_url", None)
|
||||
or getattr(module, "icon_url", None)
|
||||
or getattr(module, "icon", None),
|
||||
'id': f'{function.id}.{action["id"]}',
|
||||
'name': action.get('name', f'{function.name} ({action["id"]})'),
|
||||
'description': function.meta.description,
|
||||
'icon': action.get(
|
||||
'icon_url',
|
||||
function.meta.manifest.get('icon_url', None)
|
||||
or getattr(module, 'icon_url', None)
|
||||
or getattr(module, 'icon', None),
|
||||
),
|
||||
}
|
||||
for action in actions
|
||||
@@ -265,12 +247,12 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
else:
|
||||
return [
|
||||
{
|
||||
"id": function.id,
|
||||
"name": function.name,
|
||||
"description": function.meta.description,
|
||||
"icon": function.meta.manifest.get("icon_url", None)
|
||||
or getattr(module, "icon_url", None)
|
||||
or getattr(module, "icon", None),
|
||||
'id': function.id,
|
||||
'name': function.name,
|
||||
'description': function.meta.description,
|
||||
'icon': function.meta.manifest.get('icon_url', None)
|
||||
or getattr(module, 'icon_url', None)
|
||||
or getattr(module, 'icon', None),
|
||||
}
|
||||
]
|
||||
|
||||
@@ -278,27 +260,25 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
def get_filter_items_from_module(function, module):
|
||||
return [
|
||||
{
|
||||
"id": function.id,
|
||||
"name": function.name,
|
||||
"description": function.meta.description,
|
||||
"icon": function.meta.manifest.get("icon_url", None)
|
||||
or getattr(module, "icon_url", None)
|
||||
or getattr(module, "icon", None),
|
||||
"has_user_valves": hasattr(module, "UserValves"),
|
||||
'id': function.id,
|
||||
'name': function.name,
|
||||
'description': function.meta.description,
|
||||
'icon': function.meta.manifest.get('icon_url', None)
|
||||
or getattr(module, 'icon_url', None)
|
||||
or getattr(module, 'icon', None),
|
||||
'has_user_valves': hasattr(module, 'UserValves'),
|
||||
}
|
||||
]
|
||||
|
||||
# Batch-prefetch all needed function records to avoid N+1 queries
|
||||
all_function_ids = set()
|
||||
for model in models:
|
||||
all_function_ids.update(model.get("action_ids", []))
|
||||
all_function_ids.update(model.get("filter_ids", []))
|
||||
all_function_ids.update(model.get('action_ids', []))
|
||||
all_function_ids.update(model.get('filter_ids', []))
|
||||
all_function_ids.update(global_action_ids)
|
||||
all_function_ids.update(global_filter_ids)
|
||||
|
||||
functions_by_id = {
|
||||
f.id: f for f in Functions.get_functions_by_ids(list(all_function_ids))
|
||||
}
|
||||
functions_by_id = {f.id: f for f in Functions.get_functions_by_ids(list(all_function_ids))}
|
||||
|
||||
# Pre-warm the function module cache once per unique function ID.
|
||||
# This ensures each function's DB freshness check runs exactly once,
|
||||
@@ -307,28 +287,26 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
try:
|
||||
get_function_module_from_cache(request, function_id)
|
||||
except Exception as e:
|
||||
log.info(f"Failed to load function module for {function_id}: {e}")
|
||||
log.info(f'Failed to load function module for {function_id}: {e}')
|
||||
|
||||
# Apply global model defaults to all models
|
||||
# Per-model overrides take precedence over global defaults
|
||||
default_metadata = (
|
||||
getattr(request.app.state.config, "DEFAULT_MODEL_METADATA", None) or {}
|
||||
)
|
||||
default_metadata = getattr(request.app.state.config, 'DEFAULT_MODEL_METADATA', None) or {}
|
||||
|
||||
if default_metadata:
|
||||
for model in models:
|
||||
info = model.get("info")
|
||||
info = model.get('info')
|
||||
|
||||
if info is None:
|
||||
model["info"] = {"meta": copy.deepcopy(default_metadata)}
|
||||
model['info'] = {'meta': copy.deepcopy(default_metadata)}
|
||||
continue
|
||||
|
||||
meta = info.setdefault("meta", {})
|
||||
meta = info.setdefault('meta', {})
|
||||
for key, value in default_metadata.items():
|
||||
if key == "capabilities":
|
||||
if key == 'capabilities':
|
||||
# Merge capabilities: defaults as base, per-model overrides win
|
||||
existing = meta.get("capabilities") or {}
|
||||
meta["capabilities"] = {**value, **existing}
|
||||
existing = meta.get('capabilities') or {}
|
||||
meta['capabilities'] = {**value, **existing}
|
||||
elif meta.get(key) is None:
|
||||
meta[key] = copy.deepcopy(value)
|
||||
|
||||
@@ -339,10 +317,10 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
def get_action_priority(action_id):
|
||||
try:
|
||||
function_module = request.app.state.FUNCTIONS.get(action_id)
|
||||
if function_module and hasattr(function_module, "Valves"):
|
||||
if function_module and hasattr(function_module, 'Valves'):
|
||||
valves_db = all_function_valves.get(action_id)
|
||||
valves = function_module.Valves(**(valves_db if valves_db else {}))
|
||||
return getattr(valves, "priority", 0)
|
||||
return getattr(valves, 'priority', 0)
|
||||
except Exception:
|
||||
pass
|
||||
return 0
|
||||
@@ -350,51 +328,47 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
for model in models:
|
||||
action_ids = [
|
||||
action_id
|
||||
for action_id in list(set(model.pop("action_ids", []) + global_action_ids))
|
||||
for action_id in list(set(model.pop('action_ids', []) + global_action_ids))
|
||||
if action_id in enabled_action_ids
|
||||
]
|
||||
action_ids.sort(key=lambda aid: (get_action_priority(aid), aid))
|
||||
|
||||
filter_ids = [
|
||||
filter_id
|
||||
for filter_id in list(set(model.pop("filter_ids", []) + global_filter_ids))
|
||||
for filter_id in list(set(model.pop('filter_ids', []) + global_filter_ids))
|
||||
if filter_id in enabled_filter_ids
|
||||
]
|
||||
|
||||
model["actions"] = []
|
||||
model['actions'] = []
|
||||
for action_id in action_ids:
|
||||
action_function = functions_by_id.get(action_id)
|
||||
if action_function is None:
|
||||
log.info(f"Action not found: {action_id}")
|
||||
log.info(f'Action not found: {action_id}')
|
||||
continue
|
||||
|
||||
function_module = request.app.state.FUNCTIONS.get(action_id)
|
||||
if function_module is None:
|
||||
log.info(f"Failed to load action module: {action_id}")
|
||||
log.info(f'Failed to load action module: {action_id}')
|
||||
continue
|
||||
model["actions"].extend(
|
||||
get_action_items_from_module(action_function, function_module)
|
||||
)
|
||||
model['actions'].extend(get_action_items_from_module(action_function, function_module))
|
||||
|
||||
model["filters"] = []
|
||||
model['filters'] = []
|
||||
for filter_id in filter_ids:
|
||||
filter_function = functions_by_id.get(filter_id)
|
||||
if filter_function is None:
|
||||
log.info(f"Filter not found: {filter_id}")
|
||||
log.info(f'Filter not found: {filter_id}')
|
||||
continue
|
||||
|
||||
function_module = request.app.state.FUNCTIONS.get(filter_id)
|
||||
if function_module is None:
|
||||
log.info(f"Failed to load filter module: {filter_id}")
|
||||
log.info(f'Failed to load filter module: {filter_id}')
|
||||
continue
|
||||
if getattr(function_module, "toggle", None):
|
||||
model["filters"].extend(
|
||||
get_filter_items_from_module(filter_function, function_module)
|
||||
)
|
||||
if getattr(function_module, 'toggle', None):
|
||||
model['filters'].extend(get_filter_items_from_module(filter_function, function_module))
|
||||
|
||||
log.debug(f"get_all_models() returned {len(models)} models")
|
||||
log.debug(f'get_all_models() returned {len(models)} models')
|
||||
|
||||
models_dict = {model["id"]: model for model in models}
|
||||
models_dict = {model['id']: model for model in models}
|
||||
if isinstance(request.app.state.MODELS, RedisDict):
|
||||
request.app.state.MODELS.set(models_dict)
|
||||
else:
|
||||
@@ -404,81 +378,78 @@ async def get_all_models(request, refresh: bool = False, user: UserModel = None)
|
||||
|
||||
|
||||
def check_model_access(user, model, db=None):
|
||||
if model.get("arena"):
|
||||
meta = model.get("info", {}).get("meta", {})
|
||||
access_grants = meta.get("access_grants", [])
|
||||
if model.get('arena'):
|
||||
meta = model.get('info', {}).get('meta', {})
|
||||
access_grants = meta.get('access_grants', [])
|
||||
if not has_access(
|
||||
user.id,
|
||||
permission="read",
|
||||
permission='read',
|
||||
access_grants=access_grants,
|
||||
db=db,
|
||||
):
|
||||
raise Exception("Model not found")
|
||||
raise Exception('Model not found')
|
||||
else:
|
||||
model_info = Models.get_model_by_id(model.get("id"), db=db)
|
||||
model_info = Models.get_model_by_id(model.get('id'), db=db)
|
||||
if not model_info:
|
||||
raise Exception("Model not found")
|
||||
raise Exception('Model not found')
|
||||
elif not (
|
||||
user.id == model_info.user_id
|
||||
or AccessGrants.has_access(
|
||||
user_id=user.id,
|
||||
resource_type="model",
|
||||
resource_type='model',
|
||||
resource_id=model_info.id,
|
||||
permission="read",
|
||||
permission='read',
|
||||
db=db,
|
||||
)
|
||||
):
|
||||
raise Exception("Model not found")
|
||||
raise Exception('Model not found')
|
||||
|
||||
|
||||
def get_filtered_models(models, user, db=None):
|
||||
# Filter out models that the user does not have access to
|
||||
if (
|
||||
user.role == "user"
|
||||
or (user.role == "admin" and not BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
user.role == 'user' or (user.role == 'admin' and not BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
) and not BYPASS_MODEL_ACCESS_CONTROL:
|
||||
model_infos = {}
|
||||
for model in models:
|
||||
if model.get("arena"):
|
||||
if model.get('arena'):
|
||||
continue
|
||||
info = model.get("info")
|
||||
info = model.get('info')
|
||||
if info:
|
||||
model_infos[model["id"]] = info
|
||||
model_infos[model['id']] = info
|
||||
|
||||
user_group_ids = {
|
||||
group.id for group in Groups.get_groups_by_member_id(user.id, db=db)
|
||||
}
|
||||
user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user.id, db=db)}
|
||||
|
||||
# Batch-fetch accessible resource IDs in a single query instead of N has_access calls
|
||||
accessible_model_ids = AccessGrants.get_accessible_resource_ids(
|
||||
user_id=user.id,
|
||||
resource_type="model",
|
||||
resource_type='model',
|
||||
resource_ids=list(model_infos.keys()),
|
||||
permission="read",
|
||||
permission='read',
|
||||
user_group_ids=user_group_ids,
|
||||
db=db,
|
||||
)
|
||||
|
||||
filtered_models = []
|
||||
for model in models:
|
||||
if model.get("arena"):
|
||||
meta = model.get("info", {}).get("meta", {})
|
||||
access_grants = meta.get("access_grants", [])
|
||||
if model.get('arena'):
|
||||
meta = model.get('info', {}).get('meta', {})
|
||||
access_grants = meta.get('access_grants', [])
|
||||
if has_access(
|
||||
user.id,
|
||||
permission="read",
|
||||
permission='read',
|
||||
access_grants=access_grants,
|
||||
user_group_ids=user_group_ids,
|
||||
):
|
||||
filtered_models.append(model)
|
||||
continue
|
||||
|
||||
model_info = model_infos.get(model["id"])
|
||||
model_info = model_infos.get(model['id'])
|
||||
if model_info:
|
||||
if (
|
||||
(user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
or user.id == model_info.get("user_id")
|
||||
or model["id"] in accessible_model_ids
|
||||
(user.role == 'admin' and BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
or user.id == model_info.get('user_id')
|
||||
or model['id'] in accessible_model_ids
|
||||
):
|
||||
filtered_models.append(model)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user