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
+68 -106
View File
@@ -28,13 +28,13 @@ log = logging.getLogger(__name__)
# ModelParams is a model for the data stored in the params field of the Model table
class ModelParams(BaseModel):
model_config = ConfigDict(extra="allow")
model_config = ConfigDict(extra='allow')
pass
# ModelMeta is a model for the data stored in the meta field of the Model table
class ModelMeta(BaseModel):
profile_image_url: Optional[str] = "/static/favicon.png"
profile_image_url: Optional[str] = '/static/favicon.png'
description: Optional[str] = None
"""
@@ -43,13 +43,13 @@ class ModelMeta(BaseModel):
capabilities: Optional[dict] = None
model_config = ConfigDict(extra="allow")
model_config = ConfigDict(extra='allow')
pass
class Model(Base):
__tablename__ = "model"
__tablename__ = 'model'
id = Column(Text, primary_key=True, unique=True)
"""
@@ -139,10 +139,8 @@ class ModelForm(BaseModel):
class ModelsTable:
def _get_access_grants(
self, model_id: str, db: Optional[Session] = None
) -> list[AccessGrantModel]:
return AccessGrants.get_grants_by_resource("model", model_id, db=db)
def _get_access_grants(self, model_id: str, db: Optional[Session] = None) -> list[AccessGrantModel]:
return AccessGrants.get_grants_by_resource('model', model_id, db=db)
def _to_model_model(
self,
@@ -150,13 +148,9 @@ class ModelsTable:
access_grants: Optional[list[AccessGrantModel]] = None,
db: Optional[Session] = None,
) -> ModelModel:
model_data = ModelModel.model_validate(model).model_dump(
exclude={"access_grants"}
)
model_data["access_grants"] = (
access_grants
if access_grants is not None
else self._get_access_grants(model_data["id"], db=db)
model_data = ModelModel.model_validate(model).model_dump(exclude={'access_grants'})
model_data['access_grants'] = (
access_grants if access_grants is not None else self._get_access_grants(model_data['id'], db=db)
)
return ModelModel.model_validate(model_data)
@@ -167,37 +161,32 @@ class ModelsTable:
with get_db_context(db) as db:
result = Model(
**{
**form_data.model_dump(exclude={"access_grants"}),
"user_id": user_id,
"created_at": int(time.time()),
"updated_at": int(time.time()),
**form_data.model_dump(exclude={'access_grants'}),
'user_id': user_id,
'created_at': int(time.time()),
'updated_at': int(time.time()),
}
)
db.add(result)
db.commit()
db.refresh(result)
AccessGrants.set_access_grants(
"model", result.id, form_data.access_grants, db=db
)
AccessGrants.set_access_grants('model', result.id, form_data.access_grants, db=db)
if result:
return self._to_model_model(result, db=db)
else:
return None
except Exception as e:
log.exception(f"Failed to insert a new model: {e}")
log.exception(f'Failed to insert a new model: {e}')
return None
def get_all_models(self, db: Optional[Session] = None) -> list[ModelModel]:
with get_db_context(db) as db:
all_models = db.query(Model).all()
model_ids = [model.id for model in all_models]
grants_map = AccessGrants.get_grants_by_resources("model", model_ids, db=db)
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
return [
self._to_model_model(
model, access_grants=grants_map.get(model.id, []), db=db
)
for model in all_models
self._to_model_model(model, access_grants=grants_map.get(model.id, []), db=db) for model in all_models
]
def get_models(self, db: Optional[Session] = None) -> list[ModelUserResponse]:
@@ -209,7 +198,7 @@ class ModelsTable:
users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else []
users_dict = {user.id: user for user in users}
grants_map = AccessGrants.get_grants_by_resources("model", model_ids, db=db)
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
models = []
for model in all_models:
@@ -222,7 +211,7 @@ class ModelsTable:
access_grants=grants_map.get(model.id, []),
db=db,
).model_dump(),
"user": user.model_dump() if user else None,
'user': user.model_dump() if user else None,
}
)
)
@@ -232,28 +221,23 @@ class ModelsTable:
with get_db_context(db) as db:
all_models = db.query(Model).filter(Model.base_model_id == None).all()
model_ids = [model.id for model in all_models]
grants_map = AccessGrants.get_grants_by_resources("model", model_ids, db=db)
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
return [
self._to_model_model(
model, access_grants=grants_map.get(model.id, []), db=db
)
for model in all_models
self._to_model_model(model, access_grants=grants_map.get(model.id, []), db=db) for model in all_models
]
def get_models_by_user_id(
self, user_id: str, permission: str = "write", db: Optional[Session] = None
self, user_id: str, permission: str = 'write', db: Optional[Session] = None
) -> list[ModelUserResponse]:
models = self.get_models(db=db)
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)}
return [
model
for model in models
if model.user_id == user_id
or AccessGrants.has_access(
user_id=user_id,
resource_type="model",
resource_type='model',
resource_id=model.id,
permission=permission,
user_group_ids=user_group_ids,
@@ -261,13 +245,13 @@ class ModelsTable:
)
]
def _has_permission(self, db, query, filter: dict, permission: str = "read"):
def _has_permission(self, db, query, filter: dict, permission: str = 'read'):
return AccessGrants.has_permission_filter(
db=db,
query=query,
DocumentModel=Model,
filter=filter,
resource_type="model",
resource_type='model',
permission=permission,
)
@@ -285,22 +269,22 @@ class ModelsTable:
query = query.filter(Model.base_model_id != None)
if filter:
query_key = filter.get("query")
query_key = filter.get('query')
if query_key:
query = query.filter(
or_(
Model.name.ilike(f"%{query_key}%"),
Model.base_model_id.ilike(f"%{query_key}%"),
User.name.ilike(f"%{query_key}%"),
User.email.ilike(f"%{query_key}%"),
User.username.ilike(f"%{query_key}%"),
Model.name.ilike(f'%{query_key}%'),
Model.base_model_id.ilike(f'%{query_key}%'),
User.name.ilike(f'%{query_key}%'),
User.email.ilike(f'%{query_key}%'),
User.username.ilike(f'%{query_key}%'),
)
)
view_option = filter.get("view_option")
if view_option == "created":
view_option = filter.get('view_option')
if view_option == 'created':
query = query.filter(Model.user_id == user_id)
elif view_option == "shared":
elif view_option == 'shared':
query = query.filter(Model.user_id != user_id)
# Apply access control filtering
@@ -308,10 +292,10 @@ class ModelsTable:
db,
query,
filter,
permission="read",
permission='read',
)
tag = filter.get("tag")
tag = filter.get('tag')
if tag:
# TODO: This is a simple implementation and should be improved for performance
like_pattern = f'%"{tag.lower()}"%' # `"tag"` inside JSON array
@@ -319,21 +303,21 @@ class ModelsTable:
query = query.filter(meta_text.like(like_pattern))
order_by = filter.get("order_by")
direction = filter.get("direction")
order_by = filter.get('order_by')
direction = filter.get('direction')
if order_by == "name":
if direction == "asc":
if order_by == 'name':
if direction == 'asc':
query = query.order_by(Model.name.asc())
else:
query = query.order_by(Model.name.desc())
elif order_by == "created_at":
if direction == "asc":
elif order_by == 'created_at':
if direction == 'asc':
query = query.order_by(Model.created_at.asc())
else:
query = query.order_by(Model.created_at.desc())
elif order_by == "updated_at":
if direction == "asc":
elif order_by == 'updated_at':
if direction == 'asc':
query = query.order_by(Model.updated_at.asc())
else:
query = query.order_by(Model.updated_at.desc())
@@ -352,7 +336,7 @@ class ModelsTable:
items = query.all()
model_ids = [model.id for model, _ in items]
grants_map = AccessGrants.get_grants_by_resources("model", model_ids, db=db)
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
models = []
for model, user in items:
@@ -363,19 +347,13 @@ class ModelsTable:
access_grants=grants_map.get(model.id, []),
db=db,
).model_dump(),
user=(
UserResponse(**UserModel.model_validate(user).model_dump())
if user
else None
),
user=(UserResponse(**UserModel.model_validate(user).model_dump()) if user else None),
)
)
return ModelListResponse(items=models, total=total)
def get_model_by_id(
self, id: str, db: Optional[Session] = None
) -> Optional[ModelModel]:
def get_model_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ModelModel]:
try:
with get_db_context(db) as db:
model = db.get(Model, id)
@@ -383,16 +361,12 @@ class ModelsTable:
except Exception:
return None
def get_models_by_ids(
self, ids: list[str], db: Optional[Session] = None
) -> list[ModelModel]:
def get_models_by_ids(self, ids: list[str], db: Optional[Session] = None) -> list[ModelModel]:
try:
with get_db_context(db) as db:
models = db.query(Model).filter(Model.id.in_(ids)).all()
model_ids = [model.id for model in models]
grants_map = AccessGrants.get_grants_by_resources(
"model", model_ids, db=db
)
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
return [
self._to_model_model(
model,
@@ -404,9 +378,7 @@ class ModelsTable:
except Exception:
return []
def toggle_model_by_id(
self, id: str, db: Optional[Session] = None
) -> Optional[ModelModel]:
def toggle_model_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ModelModel]:
with get_db_context(db) as db:
try:
model = db.query(Model).filter_by(id=id).first()
@@ -422,30 +394,26 @@ class ModelsTable:
except Exception:
return None
def update_model_by_id(
self, id: str, model: ModelForm, db: Optional[Session] = None
) -> Optional[ModelModel]:
def update_model_by_id(self, id: str, model: ModelForm, db: Optional[Session] = None) -> Optional[ModelModel]:
try:
with get_db_context(db) as db:
# update only the fields that are present in the model
data = model.model_dump(exclude={"id", "access_grants"})
data = model.model_dump(exclude={'id', 'access_grants'})
result = db.query(Model).filter_by(id=id).update(data)
db.commit()
if model.access_grants is not None:
AccessGrants.set_access_grants(
"model", id, model.access_grants, db=db
)
AccessGrants.set_access_grants('model', id, model.access_grants, db=db)
return self.get_model_by_id(id, db=db)
except Exception as e:
log.exception(f"Failed to update the model by id {id}: {e}")
log.exception(f'Failed to update the model by id {id}: {e}')
return None
def delete_model_by_id(self, id: str, db: Optional[Session] = None) -> bool:
try:
with get_db_context(db) as db:
AccessGrants.revoke_all_access("model", id, db=db)
AccessGrants.revoke_all_access('model', id, db=db)
db.query(Model).filter_by(id=id).delete()
db.commit()
@@ -458,7 +426,7 @@ class ModelsTable:
with get_db_context(db) as db:
model_ids = [row[0] for row in db.query(Model.id).all()]
for model_id in model_ids:
AccessGrants.revoke_all_access("model", model_id, db=db)
AccessGrants.revoke_all_access('model', model_id, db=db)
db.query(Model).delete()
db.commit()
@@ -466,9 +434,7 @@ class ModelsTable:
except Exception:
return False
def sync_models(
self, user_id: str, models: list[ModelModel], db: Optional[Session] = None
) -> list[ModelModel]:
def sync_models(self, user_id: str, models: list[ModelModel], db: Optional[Session] = None) -> list[ModelModel]:
try:
with get_db_context(db) as db:
# Get existing models
@@ -483,37 +449,33 @@ class ModelsTable:
if model.id in existing_ids:
db.query(Model).filter_by(id=model.id).update(
{
**model.model_dump(exclude={"access_grants"}),
"user_id": user_id,
"updated_at": int(time.time()),
**model.model_dump(exclude={'access_grants'}),
'user_id': user_id,
'updated_at': int(time.time()),
}
)
else:
new_model = Model(
**{
**model.model_dump(exclude={"access_grants"}),
"user_id": user_id,
"updated_at": int(time.time()),
**model.model_dump(exclude={'access_grants'}),
'user_id': user_id,
'updated_at': int(time.time()),
}
)
db.add(new_model)
AccessGrants.set_access_grants(
"model", model.id, model.access_grants, db=db
)
AccessGrants.set_access_grants('model', model.id, model.access_grants, db=db)
# Remove models that are no longer present
for model in existing_models:
if model.id not in new_model_ids:
AccessGrants.revoke_all_access("model", model.id, db=db)
AccessGrants.revoke_all_access('model', model.id, db=db)
db.delete(model)
db.commit()
all_models = db.query(Model).all()
model_ids = [model.id for model in all_models]
grants_map = AccessGrants.get_grants_by_resources(
"model", model_ids, db=db
)
grants_map = AccessGrants.get_grants_by_resources('model', model_ids, db=db)
return [
self._to_model_model(
model,
@@ -523,7 +485,7 @@ class ModelsTable:
for model in all_models
]
except Exception as e:
log.exception(f"Error syncing models for user {user_id}: {e}")
log.exception(f'Error syncing models for user {user_id}: {e}')
return []