This commit is contained in:
Timothy Jaeryang Baek
2026-03-17 17:58:01 -05:00
parent fcf7208352
commit de3317e26b
220 changed files with 17200 additions and 22836 deletions
+141 -170
View File
@@ -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)