feat: Add multi-prompt stage support per assistant #88
Open
AasheeshLikePanner wants to merge 6 commits into
Open
feat: Add multi-prompt stage support per assistant #88AasheeshLikePanner wants to merge 6 commits into
AasheeshLikePanner wants to merge 6 commits into
Conversation
- Add AssistantPromptStage entity for storing multiple prompts per assistant - Add current_stage_id to conversations for tracking current stage - Add service methods for CRUD operations on prompt stages - Add stage template selection in LLM executor - Add database migration for prompt_stages table This enables step-based behavior where bot can use different system prompts at different conversation stages (e.g., identity_check -> order_placement). Stage transitions can be triggered via tool calls.
…ai#72) - Add built-in transition_stage tool that LLM can call to switch stages - Add stage_transition_caller that emits directive packets - Add handleStageTransition in dispatch to process stage changes - Update conversation.current_stage_id when transition is requested Now the bot can automatically transition between prompts: 1. LLM calls transition_stage(stage_name='order_placement') 2. Tool emits directive packet with transition_stage arg 3. Dispatcher processes and updates conversation stage 4. Next LLM request uses the new stage's prompt
AasheeshLikePanner
changed the title
There was a problem hiding this comment.
Pull request overview
This PR introduces “prompt stages” to support multi-step assistants by allowing different system prompts to be used at different points in a conversation, with a built-in tool to transition between stages.
Changes:
- Adds DB schema for
assistant_prompt_stagesand trackscurrent_stage_idonassistant_conversations. - Adds assistant-service CRUD methods for prompt stages and a method to update the conversation’s current stage.
- Registers a built-in
transition_stagetool and adds stage-selection logic in the LLM executor + adapter dispatch path.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 18 comments.
Show a summary per file
| File | Description |
|---|---|
| api/assistant-api/migrations/000012_add_prompt_stages.up.sql | Creates prompt stage table and adds current_stage_id to conversations. |
| api/assistant-api/migrations/000012_add_prompt_stages.down.sql | Rolls back the prompt stage table and conversation column. |
| api/assistant-api/internal/services/assistant/assistant.impl.service.go | Implements prompt stage CRUD + conversation stage update in AssistantService. |
| api/assistant-api/internal/services/assistant.service.go | Extends AssistantService interface with prompt-stage methods. |
| api/assistant-api/internal/entity/conversations/conversation.assistant.go | Adds CurrentStageId field to the conversation entity. |
| api/assistant-api/internal/entity/assistants/prompt_stage.assistant.go | Adds AssistantPromptStage entity + template setter/getter. |
| api/assistant-api/internal/entity/assistants/prompt_stage_test.go | Adds unit tests around prompt stage template parsing behavior. |
| api/assistant-api/internal/agent/executor/tool/tool.go | Registers the built-in transition_stage tool. |
| api/assistant-api/internal/agent/executor/tool/internal/local/stage_transition_caller.go | Implements the stage transition tool caller. |
| api/assistant-api/internal/agent/executor/llm/internal/model/model.go | Adds stage-based prompt selection logic (via cached templates). |
| api/assistant-api/internal/adapters/internal/dispatch.go | Hooks stage transitions into directive handling and updates conversation stage. |
Comments suppressed due to low confidence (1)
api/assistant-api/internal/adapters/internal/dispatch.go:839
- Even when the directive contains
transition_stage, this code still notifiesConversationDirective_END_CONVERSATIONto the client/stream. That will terminate the conversation while also switching the stage, which is likely unintended. Stage transitions should not be routed through the END_CONVERSATION directive path (handle them separately and skip the end-conversation notification).
case protos.ConversationDirective_END_CONVERSATION:
// Check for stage transition request
if stageName, ok := vl.Arguments["transition_stage"].(string); ok && stageName != "" {
talking.handleStageTransition(ctx, vl.ContextID, stageName)
}
if err := talking.Notify(ctx, &protos.ConversationDirective{
Id: vl.ContextID,
Type: vl.Directive,
Args: anyArgs,
Time: timestamppb.Now(),
}); err != nil {
- Fix migration column names (created_date/updated_date, order) - Add schema qualification (public.*) for consistency - Use clause.OrderByColumn instead of backticks for PostgreSQL - Set audit fields (CreatedBy, UpdatedBy, Status) on stage creation/update - Use soft delete instead of hard delete for DeletePromptStage - Return proper error when no stages found in GetDefaultPromptStage - Add auth scoping to UpdateConversationStage for tenant safety - Fix SetPrompt to return error for invalid JSON - Update in-memory conversation after stage transition - Use ConversationEventPacket instead of END_CONVERSATION to avoid ending conversation
- Use communication.GetMetadata() to get stage_template in executor - Set stage_template in requestor metadata after stage transition - Simplify executor template lookup to use metadata only
- Use 'prompt' instead of 'text_chat_complete' - Use 'promptVariables' instead of 'variables' - Use 'defaultValue' instead of 'default' - Update tests to check for error return from SetPrompt
- Use stage.Template.GetTextChatCompleteTemplate() instead of stage.GetTemplate() - Add gorm_types import to dispatch.go - Template now correctly passed as *TextChatCompletePromptTemplate to executor
Member
|
@AasheeshLikePanner Please update the branch with new changes. |
Author
|
@iamprashant I’ve pulled the latest changes from main. This branch is up to date and ready to merge. |
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
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.
Description
Enables a single assistant to use different system prompts at different conversation stages. This addresses the use case where a bot needs to function in steps - e.g., first verify user identity with one prompt, then handle order placement with a different prompt.
Type of Change
Related Issues
Checklist
General
Testing
Documentation
Security
Additional Notes