Skip to content

Commit

Permalink
Merge pull request #164 from PySpur-Dev/fix/partialRun
Browse files Browse the repository at this point in the history
  • Loading branch information
srijanpatel authored Feb 18, 2025
2 parents 9d73887 + 748dbe2 commit 12487e3
Show file tree
Hide file tree
Showing 5 changed files with 101 additions and 85 deletions.
48 changes: 15 additions & 33 deletions frontend/src/components/Header.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -23,17 +23,16 @@ import { useHotkeys } from 'react-hotkeys-hook'
import { useDispatch, useSelector } from 'react-redux'
import { useSaveWorkflow } from '../hooks/useSaveWorkflow'
import { useWorkflowExecution } from '../hooks/useWorkflowExecution'
import { setProjectName } from '../store/flowSlice'
import { useWorkflowFileOperations } from '../hooks/useWorkflowFileOperations'
import { setProjectName, setRunModalOpen } from '../store/flowSlice'
import { RootState } from '../store/store'
import { AlertState } from '../types/alert'
import { getRunStatus, getWorkflow } from '../utils/api'
import ConfirmationModal from './modals/ConfirmationModal'
import DeployModal from './modals/DeployModal'
import HelpModal from './modals/HelpModal'
import RunModal from './modals/RunModal'
import SettingsCard from './modals/SettingsModal'
import { initializeFlow } from '../store/flowSlice'
import { RootState } from '../store/store'
import { useWorkflowFileOperations } from '../hooks/useWorkflowFileOperations'
import ConfirmationModal from './modals/ConfirmationModal'

interface HeaderProps {
activePage: 'dashboard' | 'workflow' | 'evals' | 'trace' | 'rag'
Expand All @@ -57,8 +56,9 @@ const Header: React.FC<HeaderProps> = ({ activePage, associatedWorkflowId, runId
isVisible: false,
})
const testInputs = useSelector((state: RootState) => state.flow.testInputs)
const [selectedRow, setSelectedRow] = useState<string | null>(null)
const selectedTestInputId = useSelector((state: RootState) => state.flow.selectedTestInputId)
const [isHelpModalOpen, setIsHelpModalOpen] = useState<boolean>(false)
const isRunModalOpen = useSelector((state: RootState) => state.flow.isRunModalOpen)

const router = useRouter()
const { id } = router.query
Expand All @@ -81,22 +81,11 @@ const Header: React.FC<HeaderProps> = ({ activePage, associatedWorkflowId, runId

const saveWorkflow = useSaveWorkflow()

const {
handleFileUpload,
isConfirmationOpen,
setIsConfirmationOpen,
handleConfirmOverwrite,
pendingWorkflowData
} = useWorkflowFileOperations({ showAlert })

useEffect(() => {
if (testInputs.length > 0 && !selectedRow) {
setSelectedRow(testInputs[0].id.toString())
}
}, [testInputs])
const { handleFileUpload, isConfirmationOpen, setIsConfirmationOpen, handleConfirmOverwrite, pendingWorkflowData } =
useWorkflowFileOperations({ showAlert })

const handleRunWorkflow = async (): Promise<void> => {
setIsDebugModalOpen(true)
dispatch(setRunModalOpen(true))
}

const handleProjectNameChange = (e: React.ChangeEvent<HTMLInputElement>): void => {
Expand Down Expand Up @@ -155,7 +144,7 @@ const Header: React.FC<HeaderProps> = ({ activePage, associatedWorkflowId, runId
return
}

const testCase = testInputs.find((row) => row.id.toString() === selectedRow) ?? testInputs[0]
const testCase = testInputs.find((row) => row.id.toString() === selectedTestInputId) ?? testInputs[0]

if (testCase) {
const { id, ...inputValues } = testCase
Expand Down Expand Up @@ -465,12 +454,7 @@ const Header: React.FC<HeaderProps> = ({ activePage, associatedWorkflowId, runId
</NavbarItem>
<NavbarItem className="hidden sm:flex">
<Tooltip content="Upload Workflow JSON">
<Button
isIconOnly
radius="full"
variant="light"
onPress={handleFileUpload}
>
<Button isIconOnly radius="full" variant="light" onPress={handleFileUpload}>
<Icon className="text-foreground/60" icon="solar:upload-linear" width={24} />
</Button>
</Tooltip>
Expand Down Expand Up @@ -552,20 +536,18 @@ const Header: React.FC<HeaderProps> = ({ activePage, associatedWorkflowId, runId
</NavbarContent>
</Navbar>
<RunModal
isOpen={isDebugModalOpen}
onOpenChange={setIsDebugModalOpen}
isOpen={isRunModalOpen}
onOpenChange={(isOpen) => dispatch(setRunModalOpen(isOpen))}
onRun={async (selectedInputs) => {
await executeWorkflow(selectedInputs)
setIsDebugModalOpen(false)
dispatch(setRunModalOpen(false))
}}
selectedRow={selectedRow}
onSelectedRowChange={setSelectedRow}
/>
<DeployModal
isOpen={isDeployModalOpen}
onOpenChange={setIsDeployModalOpen}
workflowId={workflowId}
testInput={testInputs.find((row) => row.id.toString() === selectedRow) ?? testInputs[0]}
testInput={testInputs.find((row) => row.id.toString() === selectedTestInputId) ?? testInputs[0]}
/>
<HelpModal isOpen={isHelpModalOpen} onClose={() => setIsHelpModalOpen(false)} />
</>
Expand Down
47 changes: 17 additions & 30 deletions frontend/src/components/modals/RunModal.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ import { Icon } from '@iconify/react'
import React, { useEffect, useState } from 'react'
import { useDispatch, useSelector } from 'react-redux'
import { useSaveWorkflow } from '../../hooks/useSaveWorkflow'
import { addTestInput, deleteTestInput, updateTestInput } from '../../store/flowSlice'
import { addTestInput, deleteTestInput, setSelectedTestInputId, updateTestInput } from '../../store/flowSlice'
import { getNodeMissingRequiredFields } from '../../store/nodeTypesSlice'
import { AppDispatch, RootState } from '../../store/store'
import FileUploadBox from '../FileUploadBox'
Expand All @@ -33,27 +33,19 @@ interface RunModalProps {
onOpenChange: (isOpen: boolean) => void
onRun: (initialInputs: Record<string, any>, files?: Record<string, string[]>) => void
onSave?: () => void
selectedRow?: string | null
onSelectedRowChange?: (rowId: string | null) => void
}

interface EditingCell {
rowId: number
field: string
}

const RunModal: React.FC<RunModalProps> = ({
isOpen,
onOpenChange,
onRun,
onSave,
selectedRow: externalSelectedRow,
onSelectedRowChange,
}) => {
const RunModal: React.FC<RunModalProps> = ({ isOpen, onOpenChange, onRun, onSave }) => {
const nodes = useSelector((state: RootState) => state.flow.nodes)
const nodeConfigs = useSelector((state: RootState) => state.flow.nodeConfigs)
const nodeTypesMetadata = useSelector((state: RootState) => state.nodeTypes).metadata
const workflowID = useSelector((state: RootState) => state.flow.workflowID)
const selectedTestInputId = useSelector((state: RootState) => state.flow.selectedTestInputId)
const inputNode = nodes.find((node) => node.type === 'InputNode')
const workflowInputVariables = inputNode ? nodeConfigs[inputNode.id]?.output_schema || {} : {}
const workflowInputVariableNames = Object.keys(workflowInputVariables)
Expand All @@ -69,7 +61,6 @@ const RunModal: React.FC<RunModalProps> = ({

const [testData, setTestData] = useState<TestInput[]>([])
const [editingCell, setEditingCell] = useState<EditingCell | null>(null)
const [selectedRow, setSelectedRow] = useState<string | null>(externalSelectedRow || null)
const [editorContents, setEditorContents] = useState<Record<string, string>>({})
const [uploadedFiles, setUploadedFiles] = useState<Record<string, File[]>>({})
const [filePaths, setFilePaths] = useState<Record<string, string[]>>({})
Expand All @@ -84,18 +75,11 @@ const RunModal: React.FC<RunModalProps> = ({
}, [testInputs])

useEffect(() => {
if (isOpen && testData.length > 0 && !selectedRow) {
if (isOpen && testData.length > 0 && !selectedTestInputId) {
const newSelectedRow = testData[0].id.toString()
setSelectedRow(newSelectedRow)
onSelectedRowChange?.(newSelectedRow)
dispatch(setSelectedTestInputId(newSelectedRow))
}
}, [isOpen, testData, selectedRow])

useEffect(() => {
if (externalSelectedRow !== selectedRow) {
setSelectedRow(externalSelectedRow)
}
}, [externalSelectedRow])
}, [isOpen, testData, selectedTestInputId, dispatch])

const getNextId = () => {
const maxId = testData.reduce((max, row) => Math.max(max, row.id), 0)
Expand All @@ -114,13 +98,15 @@ const RunModal: React.FC<RunModalProps> = ({
}
setTestData([...testData, newTestInput])
setEditorContents({}) // Clear editor contents
setSelectedRow(newId.toString()) // Convert to string for selection
dispatch(addTestInput(newTestInput))
saveWorkflow()
}

const handleDeleteRow = (id: number) => {
setTestData(testData.filter((row) => row.id !== id))
if (selectedTestInputId === id.toString()) {
dispatch(setSelectedTestInputId(null))
}
dispatch(deleteTestInput({ id }))
saveWorkflow()
}
Expand Down Expand Up @@ -278,11 +264,11 @@ const RunModal: React.FC<RunModalProps> = ({
}
setTestData([...testData, newTestInput])
dispatch(addTestInput(newTestInput))
setSelectedRow(newId.toString())
dispatch(setSelectedTestInputId(newId.toString()))
setEditorContents({})
testCaseToRun = newTestInput
} else {
testCaseToRun = testData.find((row) => row.id.toString() === selectedRow)
testCaseToRun = testData.find((row) => row.id.toString() === selectedTestInputId)
}

if (!testCaseToRun) return false
Expand All @@ -308,7 +294,7 @@ const RunModal: React.FC<RunModalProps> = ({
}
setTestData([...testData, newTestInput])
dispatch(addTestInput(newTestInput))
setSelectedRow(newId.toString())
dispatch(setSelectedTestInputId(newId.toString()))
setEditorContents({}) // Clear editor contents
saveWorkflow()
}
Expand Down Expand Up @@ -353,11 +339,10 @@ const RunModal: React.FC<RunModalProps> = ({
disabledKeys={
editingCell ? new Set([editingCell.rowId.toString()]) : new Set()
}
selectedKeys={selectedRow ? [selectedRow] : new Set()}
selectedKeys={selectedTestInputId ? [selectedTestInputId] : new Set()}
onSelectionChange={(selection) => {
const selectedKey = Array.from(selection)[0]?.toString() || null
setSelectedRow(selectedKey)
onSelectedRowChange?.(selectedKey)
dispatch(setSelectedTestInputId(selectedKey))
}}
classNames={{
base: 'min-w-[800px]',
Expand Down Expand Up @@ -525,7 +510,9 @@ const RunModal: React.FC<RunModalProps> = ({
onClose()
}
}}
isDisabled={!selectedRow && !Object.values(editorContents).some((v) => v?.trim())}
isDisabled={
!selectedTestInputId && !Object.values(editorContents).some((v) => v?.trim())
}
startContent={<Icon icon="material-symbols:play-arrow" />}
>
Run
Expand Down
39 changes: 35 additions & 4 deletions frontend/src/hooks/usePartialRun.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
import { setRunModalOpen, updateNodeDataOnly, updateNodesFromPartialRun } from '@/store/flowSlice'
import { AppDispatch, RootState } from '@/store/store'
import { useState } from 'react'
import { useSelector } from 'react-redux'
import { runPartialWorkflow } from '../utils/api'
import { AppDispatch } from '@/store/store'
import { updateNodeDataOnly, updateNodesFromPartialRun } from '@/store/flowSlice'

interface PartialRunResult {
// Add specific result type properties based on your API response
Expand All @@ -17,7 +18,7 @@ interface PartialRunError {
export interface PartialRunParams {
workflowId: string
nodeId: string
initialInputs: Record<string, any>
initialInputs?: Record<string, any>
partialOutputs: Record<string, any>
rerunPredecessors: boolean
}
Expand All @@ -26,6 +27,9 @@ const usePartialRun = (dispatch: AppDispatch) => {
const [loading, setLoading] = useState(false)
const [error, setError] = useState<PartialRunError | null>(null)
const [result, setResult] = useState<PartialRunResult | null>(null)
const selectedTestInputId = useSelector((state: RootState) => state.flow.selectedTestInputId)
const testInputs = useSelector((state: RootState) => state.flow.testInputs)
const nodes = useSelector((state: RootState) => state.flow.nodes)

const executePartialRun = async ({
workflowId,
Expand All @@ -38,7 +42,34 @@ const usePartialRun = (dispatch: AppDispatch) => {
setError(null)

try {
const data = await runPartialWorkflow(workflowId, nodeId, initialInputs, partialOutputs, rerunPredecessors)
// If no initialInputs provided, use the selected test input
let effectiveInitialInputs = initialInputs
if (!effectiveInitialInputs && testInputs.length > 0) {
const testCase = testInputs.find((row) => row.id.toString() === selectedTestInputId) ?? testInputs[0]
if (testCase) {
const { id, ...inputValues } = testCase
const inputNode = nodes.find((node) => node.type === 'InputNode')
if (inputNode?.id) {
effectiveInitialInputs = {
[inputNode.id]: inputValues,
}
}
}
}

// If no effectiveInitialInputs and no test inputs, open the RunModal
if (!effectiveInitialInputs && testInputs.length === 0) {
dispatch(setRunModalOpen(true))
return
}

const data = await runPartialWorkflow(
workflowId,
nodeId,
effectiveInitialInputs || {},
partialOutputs,
rerunPredecessors
)
setResult(data)

// Update nodes with their outputs using the action creator
Expand Down
36 changes: 23 additions & 13 deletions frontend/src/store/flowSlice.ts
Original file line number Diff line number Diff line change
Expand Up @@ -28,10 +28,12 @@ const initialState: FlowState = {
workflowInputVariables: {},
testInputs: [],
inputNodeValues: {},
selectedTestInputId: null,
history: {
past: [],
future: [],
},
isRunModalOpen: false,
}

const saveToHistory = (state: FlowState) => {
Expand All @@ -45,10 +47,8 @@ const saveToHistory = (state: FlowState) => {
const generateJsonSchema = (schema: Record<string, any>): string => {
const jsonSchema = {
type: 'object',
properties: Object.fromEntries(
Object.entries(schema).map(([key, type]) => [key, { type }])
),
required: Object.keys(schema)
properties: Object.fromEntries(Object.entries(schema).map(([key, type]) => [key, { type }])),
required: Object.keys(schema),
}
return JSON.stringify(jsonSchema, null, 2)
}
Expand Down Expand Up @@ -247,25 +247,25 @@ const flowSlice = createSlice({

updateNodeConfigOnly: (state, action: PayloadAction<{ id: string; data: any }>) => {
const { id, data } = action.payload
const currentConfig = state.nodeConfigs[id] || {};
const currentConfig = state.nodeConfigs[id] || {}
if (data.few_shot_examples) {
const oldExamples = currentConfig.few_shot_examples || [];
const newExamples = data.few_shot_examples;
const maxLength = Math.max(oldExamples.length, newExamples.length);
const mergedExamples = [];
const oldExamples = currentConfig.few_shot_examples || []
const newExamples = data.few_shot_examples
const maxLength = Math.max(oldExamples.length, newExamples.length)
const mergedExamples = []
for (let i = 0; i < maxLength; i++) {
mergedExamples[i] = { ...(oldExamples[i] || {}), ...(newExamples[i] || {}) };
mergedExamples[i] = { ...(oldExamples[i] || {}), ...(newExamples[i] || {}) }
}
state.nodeConfigs[id] = {
...currentConfig,
...data,
few_shot_examples: mergedExamples
};
few_shot_examples: mergedExamples,
}
} else {
state.nodeConfigs[id] = {
...currentConfig,
...data,
};
}
}

// If output_schema changed, rebuild connected RouterNode/CoalesceNode schemas
Expand Down Expand Up @@ -753,6 +753,14 @@ const flowSlice = createSlice({
return node // Return unchanged node if no output found
})
},

setSelectedTestInputId: (state, action: PayloadAction<string | null>) => {
state.selectedTestInputId = action.payload
},

setRunModalOpen: (state, action: PayloadAction<boolean>) => {
state.isRunModalOpen = action.payload
},
},
})

Expand Down Expand Up @@ -791,6 +799,8 @@ export const {
addNodeWithConfig,
updateNodeParentAndCoordinates,
updateNodesFromPartialRun,
setSelectedTestInputId,
setRunModalOpen,
} = flowSlice.actions

export default flowSlice.reducer
Expand Down
Loading

0 comments on commit 12487e3

Please sign in to comment.