Refactor AiModelCategory and add Soma Detection model#9477
Open
hotzenklotz wants to merge 22 commits into
Open
Refactor AiModelCategory and add Soma Detection model#9477hotzenklotz wants to merge 22 commits into
hotzenklotz wants to merge 22 commits into
Conversation
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds three AI model categories, marks nuclei inference as deprecated, introduces Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Review rate limit: 0/1 reviews remaining, refill in 60 minutes.Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx (1)
72-85:⚠️ Potential issue | 🟡 MinorConsider handling additional categories in
mapCategoryToJobType.The switch only handles
EM_NEURONS,EM_NUCLEI,EM_GENERIC, andEM_SOMATA. If a customAiModelfrom the API has a category likeEM_SYNAPSES,EM_NEURON_TYPES, orEM_CELL_ORGANELLES, the default case will throw an error when the user selects it.🛡️ Suggested defensive handling
const mapCategoryToJobType = ( category: APIAiModelCategory, ): APIJobCommand.INFER_NEURONS | APIJobCommand.INFER_INSTANCES => { switch (category) { case APIAiModelCategory.EM_NEURONS: return APIJobCommand.INFER_NEURONS; case APIAiModelCategory.EM_NUCLEI: case APIAiModelCategory.EM_GENERIC: case APIAiModelCategory.EM_SOMATA: + case APIAiModelCategory.EM_SYNAPSES: + case APIAiModelCategory.EM_NEURON_TYPES: + case APIAiModelCategory.EM_CELL_ORGANELLES: return APIJobCommand.INFER_INSTANCES; default: - throw new Error(`Unsupported category: ${category}`); + // Fallback for any future categories + console.warn(`Unknown category: ${category}, defaulting to INFER_INSTANCES`); + return APIJobCommand.INFER_INSTANCES; } };🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx` around lines 72 - 85, mapCategoryToJobType currently throws for any APIAiModelCategory values beyond EM_NEURONS/EM_NUCLEI/EM_GENERIC/EM_SOMATA; update it to defensively handle new/custom categories (e.g., EM_SYNAPSES, EM_NEURON_TYPES, EM_CELL_ORGANELLES) by returning a safe default like APIJobCommand.INFER_INSTANCES instead of throwing, and emit a console/process warning or logger entry noting the unexpected category; modify the switch in mapCategoryToJobType to include a catch-all default that returns APIJobCommand.INFER_INSTANCES and logs the unsupported category value so unexpected enum values won’t crash the UI.frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_image_segmentation_job_context.tsx (1)
234-252:⚠️ Potential issue | 🟡 MinorMissing
customConfigurationinuseCallbackdependency array.
customConfigurationis used in the callback (lines 202, 214) but is not listed in the dependency array. This could cause the callback to use stale values.🐛 Proposed fix
}, [ areParametersValid, selectedModel, selectedJobType, selectedBoundingBox, newDatasetName, selectedLayer, dataset, isViewMode, annotationId, seedGeneratorDistanceThreshold, isEvaluationActive, splitMergerEvaluationSettings, userBoundingBoxCount, taskBoundingBoxes, skeletonAnnotation, datasetConfiguration, dispatch, + customConfiguration, ]);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_image_segmentation_job_context.tsx` around lines 234 - 252, The callback passed to useCallback references customConfiguration but the dependency array (the array containing areParametersValid, selectedModel, etc.) omits customConfiguration, risking stale values; update the useCallback dependency array to include customConfiguration (or memoize customConfiguration with useMemo if you need a stable identity) so the callback (the function created by useCallback) is re-created when customConfiguration changes and lint warnings are resolved.
🧹 Nitpick comments (2)
app/models/aimodels/AiModel.scala (1)
92-109: Consolidate pretrained voxel-size resolution in one path.
inferenceBBoxToTargetMag()now usesfindPretrainedModelVoxelSize(), butfindModelVoxelSize()below still carries the old boolean neuron-vs-nuclei assumption. Routing that helper through the same ID-based lookup would keep the pretrained model table in one place and avoid the two code paths drifting again when another model is added.Also applies to: 143-149
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@app/models/aimodels/AiModel.scala` around lines 92 - 109, inferenceBBoxToTargetMag uses findPretrainedModelVoxelSize for pretrained IDs but findModelVoxelSize still encodes the old boolean neuron-vs-nuclei branching; update findModelVoxelSize to route pretrainedModelId through findPretrainedModelVoxelSize (or change its signature to accept Option[String] and call findPretrainedModelVoxelSize when present) and remove the hardcoded boolean neuron/vs/nuclei assumption so both code paths share the same ID-based lookup; update callers (including inferenceBBoxToTargetMag and the code around the 143-149 region) to pass the pretrainedModelIdOpt into the unified helper and fall back to existing custom-AI-model voxel-size logic when no pretrained ID is provided.frontend/javascripts/admin/job/job_list_view.tsx (1)
151-154: Usejob.args.modelIdwhen naming active instance jobs.Now that nuclei and soma inference run via
INFER_INSTANCES, only legacyinfer_nucleijobs keep a specific label here. New nuclei and soma jobs will collapse into the generic instance-segmentation text/filter value, which makes job history harder to understand. Consider deriving the display name/description forINFER_INSTANCESentries fromjob.args.modelId(for examplenuclei_inferral_modelvssoma_inferral_model).Also applies to: 260-264, 441-447
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@frontend/javascripts/admin/job/job_list_view.tsx` around lines 151 - 154, Replace the hardcoded "AI Instance Segmentation" label for APIJobCommand.INFER_INSTANCES with a derived display name using job.args.modelId (e.g., `${job.args.modelId}_inferral_model`) and fallback to the generic text when modelId is missing; update any codepaths that build the job label/display string and any filters that currently match on the static "AI Instance Segmentation" value to use the derived name (look for the mapping/labeling logic around APIJobCommand.INFER_INSTANCES in job_list_view.tsx and the other similar mappings referenced in the file).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@app/models/job/JobCommand.scala`:
- Around line 14-21: lowPriorityJobs currently uses values which includes
deprecated commands (globalize_floodfills, train_model, infer_with_model,
infer_nuclei), so change the computation to exclude the legacy set: define a
legacyJobs Set including those deprecated enum members and compute
lowPriorityJobs as values.diff(highPriorityJobs ++ legacyJobs) (reference the
symbols lowPriorityJobs, highPriorityJobs, values and the deprecated enum
entries).
In `@frontend/javascripts/viewer/view/ai_jobs/utils.ts`:
- Around line 63-65: The runtime throw happens because aiModelId can be
undefined when callers run the logic; update the guard conditions so the
query/validator only run when a model is selected: in the component where the
query enabled is computed (the expression currently using selectedBoundingBox &&
selectedJobType && selectedLayer), add aiModelId to that boolean chain so it
becomes selectedBoundingBox && selectedJobType && selectedLayer && aiModelId;
likewise in the validator check (the expression currently using value &&
selectedLayer && selectedJobType), include the selected model variable used
there (selectedModel or aiModelId) so it becomes value && selectedLayer &&
selectedJobType && selectedModel/aiModelId to prevent the throw in the function
that expects aiModelId.
---
Outside diff comments:
In
`@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_image_segmentation_job_context.tsx`:
- Around line 234-252: The callback passed to useCallback references
customConfiguration but the dependency array (the array containing
areParametersValid, selectedModel, etc.) omits customConfiguration, risking
stale values; update the useCallback dependency array to include
customConfiguration (or memoize customConfiguration with useMemo if you need a
stable identity) so the callback (the function created by useCallback) is
re-created when customConfiguration changes and lint warnings are resolved.
In `@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx`:
- Around line 72-85: mapCategoryToJobType currently throws for any
APIAiModelCategory values beyond EM_NEURONS/EM_NUCLEI/EM_GENERIC/EM_SOMATA;
update it to defensively handle new/custom categories (e.g., EM_SYNAPSES,
EM_NEURON_TYPES, EM_CELL_ORGANELLES) by returning a safe default like
APIJobCommand.INFER_INSTANCES instead of throwing, and emit a console/process
warning or logger entry noting the unexpected category; modify the switch in
mapCategoryToJobType to include a catch-all default that returns
APIJobCommand.INFER_INSTANCES and logs the unsupported category value so
unexpected enum values won’t crash the UI.
---
Nitpick comments:
In `@app/models/aimodels/AiModel.scala`:
- Around line 92-109: inferenceBBoxToTargetMag uses findPretrainedModelVoxelSize
for pretrained IDs but findModelVoxelSize still encodes the old boolean
neuron-vs-nuclei branching; update findModelVoxelSize to route pretrainedModelId
through findPretrainedModelVoxelSize (or change its signature to accept
Option[String] and call findPretrainedModelVoxelSize when present) and remove
the hardcoded boolean neuron/vs/nuclei assumption so both code paths share the
same ID-based lookup; update callers (including inferenceBBoxToTargetMag and the
code around the 143-149 region) to pass the pretrainedModelIdOpt into the
unified helper and fall back to existing custom-AI-model voxel-size logic when
no pretrained ID is provided.
In `@frontend/javascripts/admin/job/job_list_view.tsx`:
- Around line 151-154: Replace the hardcoded "AI Instance Segmentation" label
for APIJobCommand.INFER_INSTANCES with a derived display name using
job.args.modelId (e.g., `${job.args.modelId}_inferral_model`) and fallback to
the generic text when modelId is missing; update any codepaths that build the
job label/display string and any filters that currently match on the static "AI
Instance Segmentation" value to use the derived name (look for the
mapping/labeling logic around APIJobCommand.INFER_INSTANCES in job_list_view.tsx
and the other similar mappings referenced in the file).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: f56dd748-2971-4f27-8448-c7ba9418df9f
⛔ Files ignored due to path filters (1)
frontend/assets/images/soma-inferral-example.pngis excluded by!**/*.png
📒 Files selected for processing (16)
app/controllers/AiModelController.scalaapp/models/aimodels/AiModel.scalaapp/models/aimodels/AiModelCategory.scalaapp/models/job/JobCommand.scalaconf/webknossos.latest.routesfrontend/javascripts/admin/api/jobs.tsfrontend/javascripts/admin/job/job_list_view.tsxfrontend/javascripts/types/api_types.tsfrontend/javascripts/viewer/view/action_bar_view.tsxfrontend/javascripts/viewer/view/ai_jobs/constants.tsfrontend/javascripts/viewer/view/ai_jobs/credit_information.tsxfrontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_analysis_settings.tsxfrontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_image_segmentation_job_context.tsxfrontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsxfrontend/javascripts/viewer/view/ai_jobs/train_ai_model/ai_training_job_context.tsxfrontend/javascripts/viewer/view/ai_jobs/utils.ts
Comment on lines
14
to
21
| val compute_mesh_file, compute_segment_index_file, convert_to_wkw, export_tiff, find_largest_segment_id, | ||
| materialize_volume_annotation, render_animation, align_sections, infer_nuclei, infer_neurons, infer_instances, | ||
| infer_mitochondria, train_neuron_model, train_instance_model, | ||
| materialize_volume_annotation, render_animation, align_sections, infer_neurons, infer_instances, infer_mitochondria, | ||
| train_neuron_model, train_instance_model, | ||
| // No-longer supported jobs, kept here to be able to display old existing jobs: | ||
| globalize_floodfills, train_model, infer_with_model = Value | ||
| globalize_floodfills, train_model, infer_with_model, infer_nuclei = Value | ||
|
|
||
| val highPriorityJobs: Set[Value] = Set(convert_to_wkw, export_tiff) | ||
| val lowPriorityJobs: Set[Value] = values.diff(highPriorityJobs) |
Contributor
There was a problem hiding this comment.
Exclude deprecated commands from lowPriorityJobs.
infer_nuclei moved to the legacy bucket, but lowPriorityJobs is still derived from values, so all deprecated commands remain part of the normal runnable set. That undermines the “no-longer supported” split and can leak legacy commands into worker capability or scheduling logic.
Suggested fix
val compute_mesh_file, compute_segment_index_file, convert_to_wkw, export_tiff, find_largest_segment_id,
materialize_volume_annotation, render_animation, align_sections, infer_neurons, infer_instances, infer_mitochondria,
train_neuron_model, train_instance_model,
// No-longer supported jobs, kept here to be able to display old existing jobs:
globalize_floodfills, train_model, infer_with_model, infer_nuclei = Value
val highPriorityJobs: Set[Value] = Set(convert_to_wkw, export_tiff)
- val lowPriorityJobs: Set[Value] = values.diff(highPriorityJobs)
+ private val deprecatedJobs: Set[Value] =
+ Set(globalize_floodfills, train_model, infer_with_model, infer_nuclei)
+ val lowPriorityJobs: Set[Value] = values.diff(highPriorityJobs ++ deprecatedJobs)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@app/models/job/JobCommand.scala` around lines 14 - 21, lowPriorityJobs
currently uses values which includes deprecated commands (globalize_floodfills,
train_model, infer_with_model, infer_nuclei), so change the computation to
exclude the legacy set: define a legacyJobs Set including those deprecated enum
members and compute lowPriorityJobs as values.diff(highPriorityJobs ++
legacyJobs) (reference the symbols lowPriorityJobs, highPriorityJobs, values and
the deprecated enum entries).
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@app/controllers/AiModelController.scala`:
- Around line 291-292: resolveAiModel currently accepts any known pretrained id
causing cross-route mismatches; before computing effectivePretrainedModelId in
AiModelController, validate that the supplied or derived pretrained model id is
allowed for the current inference route. Add route-aware validation (e.g., a
small whitelist/map keyed by route type used by the instance/neuron endpoints)
and reject pretrained ids that don’t match that whitelist (return a validation
error or None) instead of silently accepting them; update the resolution logic
that sets effectivePretrainedModelId (and the other similar spots you mentioned
around lines 349-350 and 462-473) to call this route-aware check or pass the
route into resolveAiModel so resolveAiModel enforces allowed pretrained ids per
route.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: cdb374fb-8785-4fed-9c7c-677c616cfca1
📒 Files selected for processing (3)
app/controllers/AiModelController.scalaconf/webknossos.latest.routesfrontend/javascripts/types/api_types.ts
🚧 Files skipped from review as they are similar to previous changes (2)
- conf/webknossos.latest.routes
- frontend/javascripts/types/api_types.ts
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
app/models/aimodels/AiModel.scala (1)
205-225:⚠️ Potential issue | 🟠 Major
readAccessQignoresisSuperUserOnly.
OR isPretrainedModel = TRUEmakes every pretrained row readable to every authenticated user, solistAiModelswill expose models that the newisSuperUserOnlyflag is supposed to hide. Please enforce that flag server-side instead of relying on frontend filtering.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@app/models/aimodels/AiModel.scala` around lines 205 - 225, readAccessQ currently lets any authenticated user see pretrained models via "OR isPretrainedModel = TRUE", ignoring the new isSuperUserOnly flag; change the query in readAccessQ to additionally require that either isSuperUserOnly IS NOT TRUE or the requesting user is a superuser. Concretely, add a clause like "AND (isSuperUserOnly IS NOT TRUE OR EXISTS (SELECT 1 FROM webknossos.users_ u WHERE u._id = $requestingUserId AND u.isSuperUser = TRUE))" to the model visibility logic in readAccessQ so pretrained models marked isSuperUserOnly=true are only returned when the requesting user is a superuser.
♻️ Duplicate comments (1)
app/controllers/AiModelController.scala (1)
284-286:⚠️ Potential issue | 🟠 MajorKeep raw pretrained ids route- and access-scoped.
The
findPretrainedModelVoxelSize(id).isDefinedbranch still treats any known pretrained token as valid on every inference route. That means/infer_instancescan accept neuron models,/infer_neuronscan accept nuclei/soma models, and raw ids also bypass any DB-backed flags likeisSuperUserOnlybecause noAiModelis loaded. Resolve pretrained ids against route-specific, access-checked metadata instead of accepting the string on its own.Also applies to: 346-348, 463-473
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@app/controllers/AiModelController.scala` around lines 284 - 286, The code currently accepts raw pretrained IDs (via effectivePretrainedModelId) when aiModelOpt is empty, allowing route-mismatched or access-restricted pretrained tokens; instead, validate any pretrainedModelId against route-specific and access-checked metadata rather than trusting the string. Update the logic around resolveAiModel and effectivePretrainedModelId so that when pretrainedModelIdOpt is defined but aiModelOpt.isEmpty you resolve and authorize that pretrained ID using the same route-scoped checks you use for DB-backed AiModel (e.g., call the route-specific resolver/authorizer or a helper that maps a pretrained id -> AiModel metadata and verifies isSuperUserOnly and model-type compatibility via findPretrainedModelVoxelSize), and only set effectivePretrainedModelId if that validation succeeds; apply the same pattern used in resolveAiModel for the other affected blocks referenced (lines ~346-348 and ~463-473).
🧹 Nitpick comments (2)
frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx (2)
65-70: Silent early return when category is missing may confuse users.If a model somehow lacks a category, clicking it does nothing with no feedback. Consider logging a warning or showing an error toast to aid debugging.
♻️ Optional: Add user feedback for invalid model selection
const onSelectModel = (model: AiModel) => { - if (!model.category) return; + if (!model.category) { + console.warn("Selected model has no category:", model); + return; + } const jobType = mapCategoryToJobType(model.category); setSelectedModel(model); setSelectedJobType(jobType); };🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx` around lines 65 - 70, The onSelectModel handler currently returns silently if model.category is missing; update onSelectModel (involving AiModel, mapCategoryToJobType, setSelectedModel, setSelectedJobType) to surface feedback: when model.category is falsy, log a warning (e.g., console.warn with the model id/name) and trigger a user-facing notification/toast indicating invalid model selection, otherwise proceed to compute jobType and call setSelectedModel and setSelectedJobType as before.
51-60: Query key should includeisSuperUserfor cache correctness.The
queryFnfilters models based onisSuperUser, but this value isn't included in thequeryKey. If the user's super user status were to change (e.g., in a different tab), the cached data would be stale.♻️ Proposed fix
const { data: allModels = [], isLoading } = useQueryWithErrorHandling( { - queryKey: ["aiModels"], + queryKey: ["aiModels", isSuperUser], queryFn: async () => { const models = await getAiModels(); return models.filter((aiModel) => (isSuperUser ? true : !aiModel.isSuperUserOnly)); }, }, "Could not load model list.", );🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx` around lines 51 - 60, The query key for useQueryWithErrorHandling must include the isSuperUser value so cached results reflect the user's privilege; update the queryKey in the useQueryWithErrorHandling call (currently ["aiModels"]) to include isSuperUser (e.g., ["aiModels", isSuperUser]) so the queryFn (which calls getAiModels and filters by isSuperUser) will be re-run when that flag changes and avoid stale allModels cache.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@app/controllers/AiModelController.scala`:
- Around line 298-306: The JSON payload currently sets "model_organization_id"
by passing an Option
(aiModelOpt.filterNot(_.isPretrainedModel).map(_._organization)) directly into
commandArgs, which serializes None as null; change commandArgs construction to
build a base Json.obj with the required fields and then conditionally merge in
the model_organization_id only when the Option is defined (e.g., compute val
modelOrgOpt = aiModelOpt.filterNot(_.isPretrainedModel).map(_._organization) and
do baseObj ++ modelOrgOpt.map(o => Json.obj("model_organization_id" ->
o)).getOrElse(Json.obj())), and apply the same conditional merge in the second
location (the other commandArgs around lines 367-375) so pretrained runs omit
the key entirely.
In `@app/models/aimodels/AiModel.scala`:
- Around line 103-105: The code is using the mutable user-facing AiModel.name to
lookup pretrained metadata; instead use the immutable pretrained identifier
field (pretrainedModelId) instead. Update the match branches that call
findPretrainedModelVoxelSize (including the case that currently uses
aiModel.name and the other occurrences around the noted blocks 149-155 and
160-161) to pass aiModel.pretrainedModelId or the model’s immutable pretrained
identifier, and adjust the related error message text to mention that identifier
instead of aiModel.name. Ensure any use of pretrainedModelIdOpt remains
consistent with this immutable id.
---
Outside diff comments:
In `@app/models/aimodels/AiModel.scala`:
- Around line 205-225: readAccessQ currently lets any authenticated user see
pretrained models via "OR isPretrainedModel = TRUE", ignoring the new
isSuperUserOnly flag; change the query in readAccessQ to additionally require
that either isSuperUserOnly IS NOT TRUE or the requesting user is a superuser.
Concretely, add a clause like "AND (isSuperUserOnly IS NOT TRUE OR EXISTS
(SELECT 1 FROM webknossos.users_ u WHERE u._id = $requestingUserId AND
u.isSuperUser = TRUE))" to the model visibility logic in readAccessQ so
pretrained models marked isSuperUserOnly=true are only returned when the
requesting user is a superuser.
---
Duplicate comments:
In `@app/controllers/AiModelController.scala`:
- Around line 284-286: The code currently accepts raw pretrained IDs (via
effectivePretrainedModelId) when aiModelOpt is empty, allowing route-mismatched
or access-restricted pretrained tokens; instead, validate any pretrainedModelId
against route-specific and access-checked metadata rather than trusting the
string. Update the logic around resolveAiModel and effectivePretrainedModelId so
that when pretrainedModelIdOpt is defined but aiModelOpt.isEmpty you resolve and
authorize that pretrained ID using the same route-scoped checks you use for
DB-backed AiModel (e.g., call the route-specific resolver/authorizer or a helper
that maps a pretrained id -> AiModel metadata and verifies isSuperUserOnly and
model-type compatibility via findPretrainedModelVoxelSize), and only set
effectivePretrainedModelId if that validation succeeds; apply the same pattern
used in resolveAiModel for the other affected blocks referenced (lines ~346-348
and ~463-473).
---
Nitpick comments:
In `@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx`:
- Around line 65-70: The onSelectModel handler currently returns silently if
model.category is missing; update onSelectModel (involving AiModel,
mapCategoryToJobType, setSelectedModel, setSelectedJobType) to surface feedback:
when model.category is falsy, log a warning (e.g., console.warn with the model
id/name) and trigger a user-facing notification/toast indicating invalid model
selection, otherwise proceed to compute jobType and call setSelectedModel and
setSelectedJobType as before.
- Around line 51-60: The query key for useQueryWithErrorHandling must include
the isSuperUser value so cached results reflect the user's privilege; update the
queryKey in the useQueryWithErrorHandling call (currently ["aiModels"]) to
include isSuperUser (e.g., ["aiModels", isSuperUser]) so the queryFn (which
calls getAiModels and filters by isSuperUser) will be re-run when that flag
changes and avoid stale allModels cache.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 2d23b33e-cefd-4a76-9411-e8ad17c8b9a1
📒 Files selected for processing (12)
app/controllers/AiModelController.scalaapp/controllers/InitialDataController.scalaapp/models/aimodels/AiModel.scalaapp/models/aimodels/AiModelCategory.scalafrontend/javascripts/admin/api/jobs.tsfrontend/javascripts/types/api_types.tsfrontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_analysis_settings.tsxfrontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_image_segmentation_job_context.tsxfrontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsxschema/evolutions/161-pretrained-ai-models.sqlschema/evolutions/reversions/161-pretrained-ai-models.sqlschema/schema.sql
🚧 Files skipped from review as they are similar to previous changes (2)
- app/models/aimodels/AiModelCategory.scala
- frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_analysis_settings.tsx
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
app/controllers/AiModelController.scala (1)
293-305:⚠️ Potential issue | 🟠 MajorOmit
model_organization_idinstead of serializingnull.
aiModelOpt.filterNot(_.isPretrainedModel).map(_._organization)still goes throughJson.obj, so pretrained runs emit"model_organization_id": nullrather than omitting the key. If the worker distinguishes absent from null, this still takes the wrong code path. Build the base payload first and mergemodel_organization_idonly when it exists.In Play JSON, when building Json.obj with an Option field whose value is None, is the field omitted or serialized as null?Also applies to: 356-373
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@app/controllers/AiModelController.scala` around lines 293 - 305, The Json.obj construction in AiModelController (the commandArgs payload) is currently serializing model_organization_id as null when the model is pretrained; instead, build the base Json object first (e.g., basePayload or commandArgsBase) using the existing keys, then detect the organization with aiModelOpt.filterNot(_.isPretrainedModel).map(_._organization) and only merge/add the "model_organization_id" key into the Json object when that Option is defined; apply the same pattern to the other payload block around lines 356-373 so optional fields are omitted rather than serialized as null.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@app/controllers/AiModelController.scala`:
- Around line 283-285: The aiModel resolution logic (the Fox.runOptional block
using ObjectId.fromString and aiModelDAO.findOne) needs to accept pretrained
string ids like "soma_inferral_model" / "nuclei_inferral_model": implement a
resolver function (e.g., resolveAiModelByIdOrPretrained) that takes the incoming
aiModelId string, returns the seeded AiModel from aiModelDAO when the string
matches known pretrained keys, or falls back to ObjectId.fromString(...) ->
aiModelDAO.findOne for DB ids; replace the existing
Fox.runOptional(request.body.aiModelId)(id =>
ObjectId.fromString(id).flatMap(oid => aiModelDAO.findOne(oid))) usages with
Fox.runOptional(request.body.aiModelId)(resolveAiModelByIdOrPretrained) (also
update the other similar blocks referenced), and ensure you still persist the
AiInference._aiModel as the resolved DB AiModel id while keeping the
worker-facing model_id string unchanged.
In `@app/models/aimodels/AiModel.scala`:
- Line 228: The DAO read filter currently allows pretrained models through
because of "isPretrainedModel = TRUE", which bypasses the isSuperUserOnly
restriction; update the aiModelDAO.findOne/read predicate to include the
isSuperUserOnly visibility rule (e.g., only return rows where isSuperUserOnly =
FALSE for non-superusers) so DAO-level reads respect model visibility, or
alternatively enforce a superuser check immediately after aiModelDAO.findOne in
the inference path (the request handlers that resolve models for
submission/inference) and reject non-superusers when the resolved
AiModel.isSuperUserOnly is true; change either the DAO query logic or add an
explicit gate in the model-resolution code paths (where models are fetched for
inference) to fix the bypass.
In `@schema/evolutions/reversions/162-aimodel-nullable-user.sql`:
- Around line 5-10: The rollback SQL currently only deletes pretrained models
but leaves the seeded sample_ai_model row with _user = NULL, causing the ALTER
COLUMN _user SET NOT NULL to fail; modify the migration to either (a) update the
seeded sample row in webknossos.aiModels to set a valid _user before running
ALTER TABLE (targeting the row where name = 'sample_ai_model' or seeded by
InitialDataController), or (b) add a single UPDATE that backfills every NULL
_user in webknossos.aiModels to a valid user id (or a system-owner id) before
executing ALTER TABLE, ensuring you still remove pretrained rows
(isPretrainedModel = TRUE) as currently done.
---
Duplicate comments:
In `@app/controllers/AiModelController.scala`:
- Around line 293-305: The Json.obj construction in AiModelController (the
commandArgs payload) is currently serializing model_organization_id as null when
the model is pretrained; instead, build the base Json object first (e.g.,
basePayload or commandArgsBase) using the existing keys, then detect the
organization with aiModelOpt.filterNot(_.isPretrainedModel).map(_._organization)
and only merge/add the "model_organization_id" key into the Json object when
that Option is defined; apply the same pattern to the other payload block around
lines 356-373 so optional fields are omitted rather than serialized as null.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0aa72f47-596d-418a-aee3-7438a9af307b
📒 Files selected for processing (6)
app/controllers/AiModelController.scalaapp/controllers/InitialDataController.scalaapp/models/aimodels/AiModel.scalaschema/evolutions/162-aimodel-nullable-user.sqlschema/evolutions/reversions/162-aimodel-nullable-user.sqlschema/schema.sql
| FROM webknossos.users_ | ||
| WHERE _id = $requestingUserId | ||
| )) | ||
| OR isPretrainedModel = TRUE |
Contributor
There was a problem hiding this comment.
isSuperUserOnly is bypassed by the DAO read filter.
This makes every pretrained row readable via aiModelDAO.findOne, including the nuclei/soma models seeded as superuser-only. The inference endpoints use that lookup directly, so a non-superuser can still submit jobs against those models by id. Keep the DAO predicate aligned with the visibility rules, or enforce the superuser gate before using the resolved model for inference.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@app/models/aimodels/AiModel.scala` at line 228, The DAO read filter currently
allows pretrained models through because of "isPretrainedModel = TRUE", which
bypasses the isSuperUserOnly restriction; update the aiModelDAO.findOne/read
predicate to include the isSuperUserOnly visibility rule (e.g., only return rows
where isSuperUserOnly = FALSE for non-superusers) so DAO-level reads respect
model visibility, or alternatively enforce a superuser check immediately after
aiModelDAO.findOne in the inference path (the request handlers that resolve
models for submission/inference) and reject non-superusers when the resolved
AiModel.isSuperUserOnly is true; change either the DAO query logic or add an
explicit gate in the model-resolution code paths (where models are fetched for
inference) to fix the bypass.
Comment on lines
+5
to
+10
| -- Remove pretrained models that have no user (they will be re-inserted by InitialDataController after rollback) | ||
| DELETE FROM webknossos.aiModels WHERE isPretrainedModel = TRUE; | ||
|
|
||
| ALTER TABLE webknossos.aiModels DROP CONSTRAINT _user_objectId; | ||
| ALTER TABLE webknossos.aiModels ALTER COLUMN _user SET NOT NULL; | ||
| ALTER TABLE webknossos.aiModels ADD CONSTRAINT _user_objectId CHECK (_user ~ '^[0-9a-f]{24}$'); |
Contributor
There was a problem hiding this comment.
Rollback still fails on the seeded sample model.
This only removes pretrained rows before restoring _user to NOT NULL, but InitialDataController now seeds sample_ai_model with _user = NULL and isPretrainedModel = FALSE. Rolling back a fresh 162 install will still fail on that row. Either keep the sample model owned, or clear/backfill every null _user row before SET NOT NULL.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@schema/evolutions/reversions/162-aimodel-nullable-user.sql` around lines 5 -
10, The rollback SQL currently only deletes pretrained models but leaves the
seeded sample_ai_model row with _user = NULL, causing the ALTER COLUMN _user SET
NOT NULL to fail; modify the migration to either (a) update the seeded sample
row in webknossos.aiModels to set a valid _user before running ALTER TABLE
(targeting the row where name = 'sample_ai_model' or seeded by
InitialDataController), or (b) add a single UPDATE that backfills every NULL
_user in webknossos.aiModels to a valid user id (or a system-owner id) before
executing ALTER TABLE, ensuring you still remove pretrained rows
(isPretrainedModel = TRUE) as currently done.
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (3)
app/models/aimodels/AiModel.scala (1)
149-149: 💤 Low valueTypo in variable name:
PretrainedMitchondriaVoxelSize.The variable name has a typo - "Mitchondria" should be "Mitochondria".
✏️ Suggested fix
- private lazy val PretrainedMitchondriaVoxelSize = VoxelSize(Vec3Double(7.96, 7.96, 31.2), LengthUnit.nanometer) + private lazy val PretrainedMitochondriaVoxelSize = VoxelSize(Vec3Double(7.96, 7.96, 31.2), LengthUnit.nanometer)Also update the reference on line 156:
- case "576500000000000000000002" => Some(PretrainedMitchondriaVoxelSize) + case "576500000000000000000002" => Some(PretrainedMitochondriaVoxelSize)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@app/models/aimodels/AiModel.scala` at line 149, Rename the misspelled identifier PretrainedMitchondriaVoxelSize to PretrainedMitochondriaVoxelSize across the file (and any references, e.g., the usage referenced near the variable declaration) to correct "Mitchondria" → "Mitochondria"; update the constant declaration name and all places that reference PretrainedMitchondriaVoxelSize (such as the subsequent initialization/usage) so identifiers remain consistent and compile.frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx (2)
65-70: 💤 Low valueSilent early return when model has no category.
If
model.categoryis falsy, the function returns early without any user feedback. While this shouldn't happen for properly configured models, it could lead to confusing UX where clicking a model does nothing.Consider logging a warning or showing a toast notification:
💡 Suggested improvement
const onSelectModel = (model: AiModel) => { - if (!model.category) return; + if (!model.category) { + console.warn(`Model ${model.name} has no category defined`); + return; + } const jobType = mapCategoryToJobType(model.category); setSelectedModel(model); setSelectedJobType(jobType); };🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx` around lines 65 - 70, The onSelectModel handler currently returns silently if model.category is falsy, causing no feedback; update onSelectModel to handle that branch by emitting a clear warning or user-facing notification instead of returning silently — e.g., call a toast/notification helper or console.warn with contextual info mentioning the model (use model.id or model.name if available), and avoid setting selected state when category is missing; ensure mapCategoryToJobType, setSelectedModel and setSelectedJobType are only called when a valid category exists.
185-187: 💤 Low valueTag label doesn't reflect all category types.
The tag only distinguishes between
EM_NEURONS("NEURONS") and everything else ("INSTANCES"), butEM_MITOCHONDRIAmodels would more accurately be labeled "MITOCHONDRIA". Consider extending the logic to show more descriptive tags for each category.💡 Suggested improvement
- <Tag> - {item.category === APIAiModelCategory.EM_NEURONS ? "NEURONS" : "INSTANCES"} - </Tag> + <Tag> + {item.category === APIAiModelCategory.EM_NEURONS + ? "NEURONS" + : item.category === APIAiModelCategory.EM_MITOCHONDRIA + ? "MITOCHONDRIA" + : "INSTANCES"} + </Tag>🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx` around lines 185 - 187, The Tag rendering currently maps only APIAiModelCategory.EM_NEURONS to "NEURONS" and defaults all other APIAiModelCategory values to "INSTANCES"; update the render logic in ai_model_selector.tsx (where Tag is used and item.category is read) to map each enum value to a descriptive label (e.g., APIAiModelCategory.EM_NEURONS => "NEURONS", APIAiModelCategory.EM_MITOCHONDRIA => "MITOCHONDRIA", etc.)—replace the ternary with a small switch/lookup keyed by APIAiModelCategory so new categories show accurate tags and are easy to extend.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@app/models/aimodels/AiModel.scala`:
- Line 149: Rename the misspelled identifier PretrainedMitchondriaVoxelSize to
PretrainedMitochondriaVoxelSize across the file (and any references, e.g., the
usage referenced near the variable declaration) to correct "Mitchondria" →
"Mitochondria"; update the constant declaration name and all places that
reference PretrainedMitchondriaVoxelSize (such as the subsequent
initialization/usage) so identifiers remain consistent and compile.
In `@frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx`:
- Around line 65-70: The onSelectModel handler currently returns silently if
model.category is falsy, causing no feedback; update onSelectModel to handle
that branch by emitting a clear warning or user-facing notification instead of
returning silently — e.g., call a toast/notification helper or console.warn with
contextual info mentioning the model (use model.id or model.name if available),
and avoid setting selected state when category is missing; ensure
mapCategoryToJobType, setSelectedModel and setSelectedJobType are only called
when a valid category exists.
- Around line 185-187: The Tag rendering currently maps only
APIAiModelCategory.EM_NEURONS to "NEURONS" and defaults all other
APIAiModelCategory values to "INSTANCES"; update the render logic in
ai_model_selector.tsx (where Tag is used and item.category is read) to map each
enum value to a descriptive label (e.g., APIAiModelCategory.EM_NEURONS =>
"NEURONS", APIAiModelCategory.EM_MITOCHONDRIA => "MITOCHONDRIA", etc.)—replace
the ternary with a small switch/lookup keyed by APIAiModelCategory so new
categories show accurate tags and are easy to extend.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: dc281518-8efe-4fb3-b10a-6194df96ccd0
📒 Files selected for processing (6)
app/controllers/AiModelController.scalaapp/controllers/InitialDataController.scalaapp/models/aimodels/AiModel.scalafrontend/javascripts/admin/voxelytics/ai_model_list_view.tsxfrontend/javascripts/types/api_types.tsfrontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx
🚧 Files skipped from review as they are similar to previous changes (1)
- app/controllers/InitialDataController.scala
fm3
requested changes
May 4, 2026
Member
There was a problem hiding this comment.
Nice stuff! Testing the happy path went well. I had a look at the backend code and already like the unification. I think we can go a little further in the cleanup. I added some comments and questions below, please have a look :-)
I think it would be great if someone from the frotnend team could have a look at the frontend changes.
The PR description doesn’t include the usual checkboxes, which I think are useful. For example, the migration guide entries are still missing.
| Some(AiModelCategory.em_somata), | ||
| isSuperUserOnly = true, | ||
| isPretrainedModel = true | ||
| ) |
Member
There was a problem hiding this comment.
I think it would be helpful to add the parameter names for all, not just the boolean parameters here. Otherwise all the Nones and List.empty are confusing.
| case "576500000000000000000001" => Some(PretrainedNeuronModelVoxelSize) | ||
| case "576500000000000000000002" => Some(PretrainedMitchondriaVoxelSize) | ||
| case "576500000000000000000003" => Some(PretrainedNucleiModelVoxelSize) | ||
| case "576500000000000000000004" => Some(PretrainedSomataModelVoxelSize) |
Member
There was a problem hiding this comment.
These strings should not be duplicated here and in InitialDataController. Maybe we could define them here as public variables, inject the AiModelService there and use them from it?
| voxelSize <- dataStoreClient.getEffectiveAiModelVoxelSize(modelPath) | ||
| } yield voxelSize | ||
| case None => | ||
| Fox.failure("Expected an AI model") |
Member
There was a problem hiding this comment.
Is this function ever called with aiModelOpt=None? If not, it should just take aiModel: AiModel, then we don’t need the match and can just use if(aiModel.isPretrainedModel)
| FROM webknossos.users_ | ||
| WHERE _id = $requestingUserId | ||
| )) | ||
| OR isPretrainedModel = TRUE |
Member
There was a problem hiding this comment.
Suggested change
| OR isPretrainedModel = TRUE | |
| -- The pretrained models should be visible by all users of all orgas | |
| OR isPretrainedModel |
Member
There was a problem hiding this comment.
What’s the idea behind isSuperUserOnly? Should they be hidden from non-superusers? Then this query should be adapted to reflect that. Currently it looks like the backend does not check this bool at all, neither in sql nor scala.
| trainingAnnotations: AiModelTrainingAnnotationSpecification[]; | ||
| name: string; | ||
| aiModelCategory: APIAiModelCategory.EM_NUCLEI; | ||
| aiModelCategory: APIAiModelCategory.EM_GENERIC; |
Member
There was a problem hiding this comment.
So generic is always instance? Should the name maybe reflect that?
| do $$ begin if (select schemaVersion from webknossos.releaseInformation) <> 161 then raise exception 'Previous schema version mismatch'; end if; end; $$ language plpgsql; | ||
|
|
||
| -- Remove pretrained models inserted by application startup | ||
| DELETE FROM webknossos.aiModels WHERE isPretrainedModel = TRUE; |
Member
There was a problem hiding this comment.
Suggested change
| DELETE FROM webknossos.aiModels WHERE isPretrainedModel = TRUE; | |
| DELETE FROM webknossos.aiModels WHERE isPretrainedModel; |
| _ <- organizationService.assertNoOrganizationsPresent | ||
| _ <- insertRootFolder() | ||
| _ <- insertOrganization() | ||
| _ <- insertPretrainedAiModels() |
Member
There was a problem hiding this comment.
If we want this in production and not only in dev, this needs to be moved above assertInitialDataEnabled. Maybe while we’re here you could add a comment at assertInitialDataEnabled saying something like All insert calls below this assertion are only executed in the dev setup (where initialDataEnabled is true)!
| modified TIMESTAMPTZ NOT NULL DEFAULT NOW(), | ||
| isDeleted BOOLEAN NOT NULL DEFAULT FALSE, | ||
| isSuperUserOnly BOOLEAN NOT NULL DEFAULT FALSE, | ||
| isPretrainedModel BOOLEAN NOT NULL DEFAULT FALSE, |
Member
There was a problem hiding this comment.
Suggested change
| isPretrainedModel BOOLEAN NOT NULL DEFAULT FALSE, | |
| isPretrained BOOLEAN NOT NULL DEFAULT FALSE, |
We’re already in the model context here, so I’d just call it isPretrained. This propagates to scala and typescript as well.
| private lazy val PretrainedSomataModelVoxelSize = VoxelSize(Vec3Double(179.84, 179.84, 224.0), LengthUnit.nanometer) | ||
|
|
||
| def findModelVoxelSize(aiModelOpt: Option[AiModel], usePretrainedNeuronModel: Boolean, dataStore: DataStore)( | ||
| def findPretrainedModelVoxelSize(pretrainedModelId: ObjectId): Option[VoxelSize] = |
Member
There was a problem hiding this comment.
Suggested change
| def findPretrainedModelVoxelSize(pretrainedModelId: ObjectId): Option[VoxelSize] = | |
| private def findPretrainedModelVoxelSize(pretrainedModelId: ObjectId): Option[VoxelSize] = |
|
|
||
| case class RunInferenceParameters(datasetId: ObjectId, | ||
| aiModelId: Option[ObjectId], | ||
| aiModelId: Option[String], |
Member
There was a problem hiding this comment.
If I understand correctly you moved away from non-objectId-strings again, right? So this could be changed back to ObjectId?
Also, does it still need to be optional? In my understanding a model id is now always passed? This could then also simplify runInstanceModelInference and runNeuronModelInference quite a bit
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR unifies how AI models are managed across the platform by moving pretrained models into the same database-backed system as custom-trained models. Instead of maintaining hardcoded frontend definitions for built-in models, all models are now loaded from a single API source, simplifying model management and making the system easier to extend.
Key Improvements
Unified model architecture
Expanded model category support
Added missing backend/frontend categories and introduced new categories for:
Simplified inference flow
More consistent model metadata
Improved permissions & visibility
Frontend cleanup
User Impact
Worker PR
https://github.com/scalableminds/voxelytics/pull/4527
Steps to test:
Fixes