diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index d4dea1161c..5841a59b8b 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -280,8 +280,10 @@ jobs: tags: ${{ format('{0}/{1}:{2}', steps.login-ecr.outputs.registry, steps.ref.outputs.repo-name, steps.ref.outputs.image-tag) }} push: true provenance: false - cache-from: type=gha,scope=${{ github.workflow }} - cache-to: type=gha,mode=max,scope=${{ github.workflow }},ignore-error=true + # GHA cache (type=gha) disabled: restoring the mode=max cache fills the + # runner disk and the build fails with "no space left on device" during + # COPY. Restore via registry cache (type=registry to a dedicated mutable + # ECR repo) if build times become a problem. deploy: name: Deploy to Kubernetes diff --git a/src/__test__/components/data-exploration/cell-sets-tool/AnnotateClustersTool.test.jsx b/src/__test__/components/data-exploration/cell-sets-tool/AnnotateClustersTool.test.jsx new file mode 100644 index 0000000000..31d0b55abc --- /dev/null +++ b/src/__test__/components/data-exploration/cell-sets-tool/AnnotateClustersTool.test.jsx @@ -0,0 +1,71 @@ +import React from 'react'; +import { render, screen, waitFor } from '@testing-library/react'; +import { Provider } from 'react-redux'; +import userEvent from '@testing-library/user-event'; + +import { makeStore } from 'redux/store'; +import AnnotateClustersTool from 'components/data-exploration/cell-sets-tool/AnnotateClustersTool'; +import { runCassiaAnnotation } from 'redux/actions/cellSets'; +import createTestComponentFactory from '__test__/test-utils/testComponentFactory'; +import fake from '__test__/test-utils/constants'; + +jest.mock('redux/actions/cellSets', () => ({ + __esModule: true, + runCellSetsAnnotation: jest.fn(() => () => {}), + runCassiaAnnotation: jest.fn(() => () => {}), +})); + +const experimentId = fake.EXPERIMENT_ID; +const AnnotateClustersToolFactory = createTestComponentFactory( + AnnotateClustersTool, + { experimentId, onRunAnnotation: jest.fn() }, +); + +const renderTool = () => render( + + {AnnotateClustersToolFactory()} + , +); + +describe('AnnotateClustersTool', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('defaults to CASSIA and shows the additional-context input', () => { + renderTool(); + + expect(screen.getByRole('radio', { name: 'CASSIA' })).toBeChecked(); + expect(screen.getByRole('radio', { name: 'ScType' })).not.toBeChecked(); + // CASSIA-only free-text context box is present + expect(screen.getByText(/Additional context/i)).toBeInTheDocument(); + }); + + it('sends the entered context to runCassiaAnnotation on confirm', async () => { + renderTool(); + + userEvent.type(screen.getByPlaceholderText(/Large Intestine/i), 'Large Intestine'); + userEvent.type(screen.getByPlaceholderText(/e\.g\. Human/i), 'Human'); + userEvent.type( + screen.getByPlaceholderText(/colorectal tumor/i), + '3 tumor, 2 normal', + ); + + userEvent.click(screen.getByRole('button', { name: /Compute/i })); + + // Confirmation modal opens; it names Amazon Bedrock (appears in both the + // alert title and body, so assert the unique modal title). + await waitFor(() => { + expect(screen.getByText(/Confirm CASSIA annotation/i)).toBeInTheDocument(); + }); + expect(screen.getAllByText(/Amazon Bedrock/i).length).toBeGreaterThan(0); + + userEvent.click(screen.getByRole('button', { name: /Continue/i })); + + await waitFor(() => { + expect(runCassiaAnnotation).toHaveBeenCalledWith( + experimentId, 'Human', 'Large Intestine', '3 tumor, 2 normal', + ); + }); + }); +}); diff --git a/src/__test__/components/data-exploration/cell-sets-tool/CellSetsTool.test.jsx b/src/__test__/components/data-exploration/cell-sets-tool/CellSetsTool.test.jsx index 105bedb37c..c84d930dc6 100644 --- a/src/__test__/components/data-exploration/cell-sets-tool/CellSetsTool.test.jsx +++ b/src/__test__/components/data-exploration/cell-sets-tool/CellSetsTool.test.jsx @@ -723,14 +723,13 @@ describe('CellSetsTool', () => { userEvent.click(deleteAnnotatedCellClassButton); }); - // get all the cell set groups - const numCellGroupsAfterDelete = screen.getAllByRole('img', { name: 'down' }).length; - await waitFor(() => { - // This test used to assert that "Annotated cell set" text is not found in "screen" - // in order to verify that deletion was successful. - // 4 Cell groups: louvain, custom cell sets, sample, Track_1 (metadata) - expect(numCellGroupsAfterDelete).toEqual(4); + // Deletion is optimistic, so the annotated cell class is gone immediately. + // Count the expand carets: of the remaining groups (louvain, custom cell + // sets, sample, Track_1) only the three non-empty ones show a "down" caret + // - the empty "Custom cell sets" group shows none. + const numCellGroupsAfterDelete = screen.getAllByRole('img', { name: 'down' }).length; + expect(numCellGroupsAfterDelete).toEqual(3); expect(screen.queryByText('Annotated cell class')).toBeNull(); @@ -765,6 +764,10 @@ describe('AnnotateClustersTool', () => { // Switch to tab userEvent.click(annotateClustersTabTitle); + + // CASSIA is the default method; these tests exercise the ScType select flow + // (dropdowns + direct dispatch), so switch to ScType first. + userEvent.click(screen.getByRole('radio', { name: 'ScType' })); }); it('Renders correctly', async () => { diff --git a/src/__test__/redux/actions/cellSets/runCassiaAnnotation.test.js b/src/__test__/redux/actions/cellSets/runCassiaAnnotation.test.js new file mode 100644 index 0000000000..6cfd3c65b0 --- /dev/null +++ b/src/__test__/redux/actions/cellSets/runCassiaAnnotation.test.js @@ -0,0 +1,57 @@ +import configureMockStore from 'redux-mock-store'; +import thunk from 'redux-thunk'; + +import runCassiaAnnotation from 'redux/actions/cellSets/runCassiaAnnotation'; +import fake from '__test__/test-utils/constants'; +import fetchWork from 'utils/work/fetchWork'; + +jest.mock('utils/work/fetchWork', () => jest.fn(() => ({}))); +jest.mock('utils/getTimeoutForWorkerTask', () => jest.fn(() => 100)); + +const mockStore = configureMockStore([thunk]); +const experimentId = fake.EXPERIMENT_ID; + +const newStore = () => mockStore({ + cellSets: { error: false, updatingClustering: false, loading: false }, +}); + +describe('runCassiaAnnotation', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('sends name, species, tissue and additionalInfo to the worker', async () => { + await newStore().dispatch( + runCassiaAnnotation(experimentId, 'Human', 'Large Intestine', '3 tumor, 2 normal'), + ); + + expect(fetchWork).toHaveBeenCalledTimes(1); + const [calledExperimentId, body] = fetchWork.mock.calls[0]; + expect(calledExperimentId).toEqual(experimentId); + expect(body).toEqual({ + name: 'CASSIAAnnotate', + species: 'Human', + tissue: 'Large Intestine', + additionalInfo: '3 tumor, 2 normal', + }); + }); + + it('defaults additionalInfo to an empty string when omitted', async () => { + await newStore().dispatch( + runCassiaAnnotation(experimentId, 'Human', 'Large Intestine'), + ); + + const [, body] = fetchWork.mock.calls[0]; + expect(body.additionalInfo).toBe(''); + }); + + it('does not dispatch a request when cell sets are already updating', async () => { + const store = mockStore({ + cellSets: { error: false, updatingClustering: true, loading: true }, + }); + + await store.dispatch(runCassiaAnnotation(experimentId, 'Human', 'Large Intestine')); + + expect(fetchWork).not.toHaveBeenCalled(); + }); +}); diff --git a/src/__test__/redux/reducers/componentConfig/cellClassDelete.test.js b/src/__test__/redux/reducers/componentConfig/cellClassDelete.test.js new file mode 100644 index 0000000000..382c5acddb --- /dev/null +++ b/src/__test__/redux/reducers/componentConfig/cellClassDelete.test.js @@ -0,0 +1,42 @@ +import cellClassDelete from 'redux/reducers/componentConfig/cellClassDelete'; +import { CELL_CLASS_DELETE } from 'redux/actionTypes/cellSets'; + +const deleteAction = (key) => ({ type: CELL_CLASS_DELETE, payload: { key } }); + +describe('cellClassDelete (componentConfig)', () => { + it('removes the deleted cell class from heatmap groupedTracks', () => { + const state = { + interactiveHeatmap: { + config: { groupedTracks: ['louvain', 'CASSIA-Gut-Human-1', 'sample'] }, + }, + }; + + const newState = cellClassDelete(state, deleteAction('CASSIA-Gut-Human-1')); + + expect(newState.interactiveHeatmap.config.groupedTracks).toEqual(['louvain', 'sample']); + }); + + it('leaves groupedTracks untouched when it does not contain the deleted key', () => { + const state = { + interactiveHeatmap: { config: { groupedTracks: ['louvain', 'sample'] } }, + }; + + const newState = cellClassDelete(state, deleteAction('CASSIA-Gut-Human-1')); + + expect(newState.interactiveHeatmap.config.groupedTracks).toEqual(['louvain', 'sample']); + }); + + // Regression: plots whose config has no groupedTracks used to throw + // "Cannot read properties of undefined (reading 'includes')". + it('does not throw for plot configs without groupedTracks', () => { + const state = { + embeddingCategoricalMain: { config: { selectedCellSet: 'louvain' } }, + violinMain: { config: {} }, + noConfigPlot: {}, + }; + + expect( + () => cellClassDelete(state, deleteAction('CASSIA-Gut-Human-1')), + ).not.toThrow(); + }); +}); diff --git a/src/components/Loader.jsx b/src/components/Loader.jsx index d616af4334..6c8633e42e 100644 --- a/src/components/Loader.jsx +++ b/src/components/Loader.jsx @@ -1,6 +1,6 @@ import PropTypes from 'prop-types'; import { BounceLoader } from 'react-spinners'; -import { Typography, Spin } from 'antd'; +import { Typography, Spin, Progress } from 'antd'; import useSWR from 'swr'; import React from 'react'; import { useSelector } from 'react-redux'; @@ -39,15 +39,21 @@ const slowLoad = () => ( ); -const fastLoad = (message) => ( +const fastLoad = (message, percent) => (
-
- -
+ {typeof percent === 'number' ? ( +
+ +
+ ) : ( +
+ +
+ )}

{message || "We're getting your data ..."} @@ -74,10 +80,10 @@ const Loader = ({ experimentId }) => { } if (workerInfo && workerInfo.userMessage) { - const { userMessage } = workerInfo; + const { userMessage, progressPercent } = workerInfo; return (

- {fastLoad(userMessage)} + {fastLoad(userMessage, progressPercent)}
); } diff --git a/src/components/data-exploration/cell-sets-tool/AnnotateClustersTool.jsx b/src/components/data-exploration/cell-sets-tool/AnnotateClustersTool.jsx index 440e444915..15a9b97b9b 100644 --- a/src/components/data-exploration/cell-sets-tool/AnnotateClustersTool.jsx +++ b/src/components/data-exploration/cell-sets-tool/AnnotateClustersTool.jsx @@ -3,12 +3,29 @@ import PropTypes from 'prop-types'; import _ from 'lodash'; import { + Alert, Button, - Radio, Select, Space, Tooltip, + Input, Modal, Radio, Select, Space, Table, Tooltip, Typography, } from 'antd'; -import { runCellSetsAnnotation } from 'redux/actions/cellSets'; +import { runCellSetsAnnotation, runCassiaAnnotation } from 'redux/actions/cellSets'; import { useDispatch } from 'react-redux'; +const { Text, Paragraph } = Typography; + +// Illustrative sample of what is actually sent to CASSIA's LLM provider: for +// each cluster, only the marker gene *names*, ranked by avg_log2FC from highest +// to lowest (top 50). The summary statistics (avg_log2FC, pct.1, pct.2) are used +// only locally to rank and filter the genes — they are NOT included in the prompt. +const exampleMarkerColumns = [ + { title: 'cluster', dataIndex: 'cluster', key: 'cluster' }, + { title: 'ranked marker genes (highest → lowest)', dataIndex: 'genes', key: 'genes' }, +]; + +const exampleMarkerData = [ + { key: 1, cluster: 1, genes: 'CD3D, IL7R, CD3E, TRAC, CD2, LTB, … (top 50)' }, + { key: 2, cluster: 2, genes: 'MS4A1, CD79A, CD79B, CD19, HLA-DRA, TCL1A, … (top 50)' }, +]; + const tissueOptions = [ 'Immune system', 'Pancreas', @@ -59,54 +76,221 @@ const scTypeTooltipText = ( ); +const cassiaTooltipText = ( + <> + Automatic annotation is performed using CASSIA, a multi-agent large language + model system for interpretable cell type annotation developed by Elliot Xie + et al. It computes marker genes per cluster and uses an LLM to predict cell + types. Enter the species and tissue as free text (e.g. "Human" + and "Large Intestine"). + More details can be found in + {' '} + the CASSIA github repo + . + +); + +const ANNOTATION_METHODS = { + SCTYPE: 'sctype', + CASSIA: 'cassia', +}; + const AnnotateClustersTool = ({ experimentId, onRunAnnotation }) => { const dispatch = useDispatch(); + const [method, setMethod] = useState(ANNOTATION_METHODS.CASSIA); const [tissue, setTissue] = useState(null); const [species, setSpecies] = useState(null); + const [additionalInfo, setAdditionalInfo] = useState(''); + const [cassiaModalVisible, setCassiaModalVisible] = useState(false); + + const isCassia = method === ANNOTATION_METHODS.CASSIA; + + // tissue/species have different valid values per method (enum vs free text) + const onMethodChange = (e) => { + setMethod(e.target.value); + setTissue(null); + setSpecies(null); + setAdditionalInfo(''); + }; + + const onCompute = () => { + // CASSIA sends data to an external LLM provider, so confirm first + if (isCassia) { + setCassiaModalVisible(true); + return; + } + dispatch(runCellSetsAnnotation(experimentId, species, tissue)); + onRunAnnotation(); + }; + + const onConfirmCassia = () => { + setCassiaModalVisible(false); + dispatch(runCassiaAnnotation(experimentId, species, tissue, additionalInfo)); + onRunAnnotation(); + }; return ( - + + + CASSIA + - ScType + ScType Tissue type: - setTissue(e.target.value || null)} + style={{ width: '100%' }} + size='small' + /> + ) : ( + ({ label: option, value: option }))} - value={species} - placeholder='Select a species' - onChange={setSpecies} - style={{ width: '100%' }} - size='small' - /> + {isCassia ? ( + setSpecies(e.target.value || null)} + style={{ width: '100%' }} + size='small' + /> + ) : ( +