Empowering AI Workflows: Granular Control Over Models and Personas

The Quest for Granular Control in AI Workflows

In the world of AI-powered applications, flexibility is king. We've been building out a workflow engine that orchestrates sequences of large language model (LLM) calls, allowing users to build complex AI agents. Initially, our system offered workflow-wide settings for things like the chosen LLM provider/model (e.g., OpenAI GPT-4, Anthropic Claude) and the 'persona' – a set of instructions defining the AI's tone, role, and constraints.

While powerful, this workflow-wide approach had its limitations. What if the first step of a workflow needed to act as a stoic data extractor, but the second step needed to be a friendly, creative content generator? Or perhaps one step performs best with GPT-4, while another is more cost-effective and accurate with Claude 3 Opus. Our users needed more. They needed per-step provider/model selection and per-step persona overrides.

This past session, that's exactly what we set out to build. It was a journey through data model changes, API updates, frontend UI restructuring, and, as always, a few unexpected UX curveballs.

The Backend Blueprint: Data & Logic

The first domino to fall was the data model. To allow a WorkflowStep to have its own persona, we needed to link them.

Evolving the Schema

Our prisma/schema.prisma was the starting point. We added an optional personaId to the WorkflowStep model, establishing a one-to-many relationship where a Persona can be associated with multiple WorkflowSteps.

// prisma/schema.prisma

model Persona {
  id           String @id @default(uuid())
  name         String
  description  String?
  // ... other persona fields
  workflowSteps WorkflowStep[] // Reverse relation
}

model WorkflowStep {
  id          String    @id @default(uuid())
  workflowId  String
  // ... other step fields
  personaId   String?   @map("persona_id") // New: Optional persona override
  persona     Persona?  @relation(fields: [personaId], references: [id])

  @@index([workflowId])
  @@map("workflow_steps")
}

After updating the schema, a quick npm run db:push && npm run db:generate brought our database and Prisma client up to speed, giving us a shiny new persona_id column and the step.personaId field in our code.

API & Engine Integration

Next, we extended our tRPC API. The steps.update input in src/server/trpc/routers/workflows.ts now accepts an optional personaId. This allows the frontend to send the chosen persona for a specific step.

Crucially, when duplicating a workflow, we also ensured the personaId was carried through using persona: { connect: { id } }, preserving the per-step settings for the new workflow.

The core logic for applying these overrides lives in src/server/services/workflow-engine.ts. Inside our executeStep() function, before an LLM call is made, we now check if the current step has a personaId defined. If it does, this step-specific persona is loaded from the database and overrides any workflow-level persona that might have been set. This ensures the most granular setting always takes precedence.

A small but important detail: the personas.list query, which populates our persona dropdowns, is no longer gated behind an enabled: settingsOpen flag. It's now always available, as per-step selection means personas might be needed at any time, not just during workflow setup.

The Frontend Facelift: UI/UX Challenges & Solutions

With the backend ready, the real fun (and occasional pain) began on the frontend, specifically in src/app/(dashboard)/dashboard/workflows/[id]/page.tsx.

Restructuring for Interactivity

Our step headers needed to accommodate new interactive elements: a ProviderPicker for model selection and a <select> for persona override. This led to our first major UX hurdle:

Lesson 1: Avoid Nested Interactive Elements Like the Plague

The Problem: My initial instinct was to nest the ProviderPicker directly inside the existing step header <button> element. This button was responsible for expanding/collapsing the step details.

<!-- What I tried (and failed at) -->
<button onClick={toggleExpand}>
  Step Title
  <ProviderPicker /> <!-- Nested interactive element -->
</button>

The Failure: This is invalid HTML. Browsers might try to "fix" it, but the result is unpredictable click propagation, accessibility nightmares, and general UI weirdness. Clicks on the ProviderPicker were often also triggering the parent button's expand/collapse functionality.

The Workaround & Solution: The fix was to separate the concerns. We restructured the step header into a <div> wrapper. Inside this wrapper, the toggle <button> sits on the left, and the ProviderPicker (along with other controls) sits on the right. Crucially, we added onClick={(e) => e.stopPropagation()} to the ProviderPicker's wrapper to prevent its clicks from bubbling up and triggering the step expand/collapse.

// Simplified example of the solution
<div className="flex items-center justify-between">
  <button onClick={toggleExpand} className="flex-grow text-left">
    Step Title
  </button>
  <div onClick={(e) => e.stopPropagation()}>
    <ProviderPicker />
  </div>
</div>

This ensures clear separation of interaction areas, making the UI predictable and accessible.

Implementing the New Controls

Once the structural issues were ironed out:

• Provider Picker: We added the ProviderPicker to step headers, but only when the workflow is pending or paused. This ensures users can tweak model choices for upcoming steps but not for steps already completed or actively running. We used a non-compact version with filterAvailable to give users a clear view of the model name.
• Persona Dropdown: Inside the expanded step body, we introduced a <select> dropdown for persona selection. This select component is populated with all available personas, including their descriptions (truncated to 50 characters) to help users make informed choices.

The Dance of State Management: Optimistic UI

This brings us to our second significant UX challenge:

Lesson 2: Embrace Optimistic UI for Responsiveness (and Type Safety!)

The Problem: After implementing the persona select dropdown, users reported that their selection wasn't visually persisting. They'd choose a persona, the dropdown would briefly update, then revert to the previous value. The server was receiving the update, but the UI wasn't reflecting it immediately. I was also initially using (step as any).personaId for the select value binding, which hinted at underlying type confusion.

The Failure: The flickering UI created a poor user experience, making the feature feel broken. The (step as any) cast was a red flag, masking a potential type mismatch or an outdated type definition.

The Workaround & Solution: We implemented an optimistic UI update.

1. Introduced a local personaOverrides: Record<string, string | null> state.
2. When a user selects a persona from the dropdown, we immediately update this local personaOverrides state, causing the UI to reflect the new choice instantly.
3. Simultaneously, we trigger the tRPC mutation to update the server.
4. Once the server mutation successfully completes and the data is refetched, we clear the personaOverrides state. The UI then updates using the server's confirmed state, which should now match the optimistic update.

This pattern provides a snappy, responsive user experience. The (step as any).personaId issue was resolved by ensuring the Prisma client was properly regenerated (npm run db:generate), which correctly typed step.personaId as string | null, aligning with the select's expected value.

Other Refinements

A few other bug fixes (captured in commit c96f057) polished the experience:

• Removed the compact prop from the ProviderPicker to ensure the model name was always visible, improving clarity.
• Added the persona description after its name in the options list, providing more context at a glance.

The Current State and What's Next

Both 14edacf (initial implementation) and c96f057 (UX fixes) are now pushed to main. The features are live, allowing users to define highly flexible and nuanced AI workflows.

Our immediate next steps involve thorough manual verification:

1. Testing workflows with failed Anthropic steps, switching providers, and re-running.
2. Confirming per-step persona selection persists and demonstrably influences LLM output.
3. Ensuring workflow-level personas still function correctly when no per-step override is present.
4. Verifying completed/running workflows display read-only provider text instead of interactive pickers.
5. Considering mobile UX, as the ProviderPicker is currently hidden on smaller screens (hidden sm:block) and may need a mobile-specific fallback.

This session was a great example of how a seemingly straightforward feature can lead to interesting architectural decisions and common frontend pitfalls. By tackling these challenges head-on and applying patterns like optimistic UI, we've significantly enhanced the power and flexibility of our AI workflow engine.

Happy coding!

{"thingsDone":["Added per-step provider/model selection to workflow steps","Added per-step persona override to workflow steps","Implemented optimistic UI for persona selection","Fixed nested interactive element bug","Restructured step header UI","Updated Prisma schema and database migrations","Integrated new fields into tRPC API and workflow engine logic"],"pains":["Nesting interactive elements (ProviderPicker inside a button) causing click propagation issues and invalid HTML","Persona selection not visually persisting due to lack of optimistic UI and initial type confusion"],"successes":["Successfully implemented granular per-step control for AI models and personas","Resolved critical UX bugs through UI restructuring and optimistic state management","Ensured proper data model and API integration","Improved overall user experience and workflow flexibility"],"techStack":["Next.js","React","tRPC","Prisma","PostgreSQL","TypeScript","TailwindCSS","LLMs"]}