Refactor chat storage design to implement a single-write strategy

- Updated the write strategy in CHAT_STORAGE_DESIGN.md to transition from a two-write to a single-write approach, enhancing efficiency by writing all messages (user input and assistant responses) to the database only once upon stream exit.
- Revised related documentation to clarify the new single-write process and its implications for message handling and database interactions.
- Adjusted implementation details in the Stream function to reflect the new strategy, ensuring that steps are only saved on error or interruption, thereby reducing unnecessary database writes.
This commit is contained in:
Max 2025-12-09 09:14:40 +08:00
parent 7ada8f6983
commit d4189e4d2d

View file

@ -503,20 +503,16 @@ If interrupted during delegate, the `space_snapshot` allows restoring `ctx.Space
## Write Strategy ## Write Strategy
### Two-Write Strategy ### Single-Write Strategy
All data is buffered in memory during execution and written to database only **twice**: All data is buffered in memory during execution and written to database **only once** when `Stream()` exits:
1. **Write 1 (Entry)**: When `Stream()` starts - save user input message
2. **Write 2 (Exit)**: When `Stream()` exits - batch save messages (and steps only on error/interrupt)
**Note**: Request tracking (status, tokens, duration) is handled by [OpenAPI Request Middleware](../../openapi/request/REQUEST_DESIGN.md). **Note**: Request tracking (status, tokens, duration) is handled by [OpenAPI Request Middleware](../../openapi/request/REQUEST_DESIGN.md).
``` ```
Stream() Entry Stream() Entry
├── 【Write 1】Save user input ├── Buffer user input message (role=user)
│ - User message (role=user)
├── Execution (all in memory) ├── Execution (all in memory)
│ - ctx.Send() → messageBuffer │ - ctx.Send() → messageBuffer
@ -524,10 +520,10 @@ Stream() Entry
│ - ctx.Replace() → update messageBuffer │ - ctx.Replace() → update messageBuffer
│ - Each step → stepBuffer │ - Each step → stepBuffer
└── 【Write 2】Save final state (via defer) └── 【Single Write】Save final state (via defer)
├── Always: ├── Always:
│ - Batch write all assistant messages │ - Batch write all messages (user input + assistant responses)
│ - Update token usage in openapi_request (via request_id) │ - Update token usage in openapi_request (via request_id)
└── Only on error/interrupt: └── Only on error/interrupt:
@ -536,59 +532,59 @@ Stream() Entry
### Write Points ### Write Points
| Event | Message Table | Step Table | Token Usage | | Event | Message Table | Step Table | Token Usage |
| ---------------- | -------------------- | ----------------------------------- | ----------- | | ---------------- | -------------------------------------- | ----------------------------------- | ----------- |
| Stream entry | Write 1 (user input) | - | - | | Stream entry | Buffer user input | - | - |
| During execution | Buffer in memory | Buffer in memory | - | | During execution | Buffer in memory | Buffer in memory | - |
| **Completed** | **Batch write all** | **❌ Skip (no need to resume)** | ✅ Update | | **Completed** | **Batch write all (user + assistant)** | **❌ Skip (no need to resume)** | ✅ Update |
| On interrupt | Batch write buffered | ✅ Batch write (status=interrupted) | ✅ Update | | On interrupt | Batch write all buffered | ✅ Batch write (status=interrupted) | ✅ Update |
| On error | Batch write buffered | ✅ Batch write (status=failed) | ✅ Update | | On error | Batch write all buffered | ✅ Batch write (status=failed) | ✅ Update |
**Why skip Steps on success?** **Why skip Steps on success?**
- Steps are only needed for resume/retry operations - Steps are only needed for resume/retry operations
- If completed successfully, there's nothing to resume - If completed successfully, there's nothing to resume
- Reduces database writes and keeps Step table clean - Reduces database writes and keeps Resume table clean
### Why Two Writes? ### Why Single Write?
| Scenario | What Happens | Data Safe? | | Scenario | What Happens | Data Safe? |
| ------------------ | ----------------------------------- | ---------- | | ------------------ | --------------------------------- | ---------- |
| Normal completion | `defer` triggers → Write 2 executes | ✅ | | Normal completion | `defer` triggers → Write executes | ✅ |
| User clicks stop | `defer` triggers → Write 2 executes | ✅ | | User clicks stop | `defer` triggers → Write executes | ✅ |
| LLM timeout | `defer` triggers → Write 2 executes | ✅ | | LLM timeout | `defer` triggers → Write executes | ✅ |
| Tool failure | `defer` triggers → Write 2 executes | ✅ | | Tool failure | `defer` triggers → Write executes | ✅ |
| Network disconnect | `defer` triggers → Write 2 executes | ✅ | | Network disconnect | `defer` triggers → Write executes | ✅ |
| Process crash | Service is down, user must retry | N/A | | Process crash | Service is down, user must retry | N/A |
**Note**: Process crash is a catastrophic failure handled at infrastructure level, not application level. **Note**: Process crash is a catastrophic failure handled at infrastructure level, not application level.
### Write Count Comparison ### Write Count Comparison
For a typical request: user input → hook_create → llm → tool → llm → hook_next → 5 messages For a typical request: user input → hook_create → llm → tool → hook_next → 5 messages
| Strategy | Database Writes | Notes | | Strategy | Database Writes | Notes |
| ---------------------- | --------------- | ------------------ | | ------------------------- | --------------- | --------------------- |
| Write per operation | 1 + 5 + 5 = 11 | One write per step | | Write per operation | 1 + 5 + 5 = 11 | One write per step |
| **Two-write strategy** | **2** | Entry + Exit only | | **Single-write strategy** | **1** | Exit only (via defer) |
### Implementation ### Implementation
````go ````go
func (ast *Assistant) Stream(ctx, inputMessages, options) { func (ast *Assistant) Stream(ctx, inputMessages, options) {
// ========== Write 1: Entry ==========
userMsg := createUserMessage(ctx, inputMessages)
chatStore.SaveMessages(ctx.ChatID, []*Message{userMsg})
// ========== Memory Buffers ========== // ========== Memory Buffers ==========
messageBuffer := NewMessageBuffer() messageBuffer := NewMessageBuffer()
stepBuffer := NewStepBuffer() stepBuffer := NewStepBuffer()
// Buffer user input message (not written yet)
userMsg := createUserMessage(ctx, inputMessages)
messageBuffer.Add(userMsg)
// Track current step for error handling // Track current step for error handling
var currentStep *Step var currentStep *Step
defer func() { defer func() {
// ========== Write 2: Exit (always executes) ========== // ========== Single Write: Exit (always executes) ==========
// Determine final status for incomplete steps // Determine final status for incomplete steps
finalStatus := "completed" finalStatus := "completed"
if ctx.IsInterrupted() { if ctx.IsInterrupted() {
@ -603,9 +599,13 @@ func (ast *Assistant) Stream(ctx, inputMessages, options) {
currentStep.Status = finalStatus currentStep.Status = finalStatus
} }
// Batch write all buffered data // Batch write all buffered messages (user input + assistant responses)
chatStore.SaveMessages(ctx.ChatID, messageBuffer.GetAll()) chatStore.SaveMessages(ctx.ChatID, messageBuffer.GetAll())
chatStore.SaveSteps(stepBuffer.GetAll())
// Only save steps on error/interrupt (not on success)
if finalStatus != "completed" {
chatStore.SaveResume(stepBuffer.GetAll())
}
// Update token usage in OpenAPI request record // Update token usage in OpenAPI request record
if ctx.RequestID != "" && completionResponse != nil { if ctx.RequestID != "" && completionResponse != nil {