Enhance Assistant initialization and configuration management
- Improved the Assistant's initialization process by integrating dynamic configuration loading from `agent/kb.yml`, allowing for more flexible chat session settings. - Refactored the `initKBConfig` function to streamline the reading and parsing of Knowledge Base settings. - Optimized conversation initialization to prepare KB collections asynchronously, enhancing performance during chat interactions. - Updated tests to ensure the correct application of KB settings and robust integration with the Assistant's functionality.
This commit is contained in:
Max
parent
31b627aba0
commit
467d4e2398
2 changed files with 1515 additions and 0 deletions
641
agent/store/CHAT_STORAGE_DESIGN.md
Normal file
641
agent/store/CHAT_STORAGE_DESIGN.md
Normal file
|
|
@ -0,0 +1,641 @@
|
|||
# Chat Storage Design
|
||||
|
||||
This document describes the design for storing chat conversations, messages, and execution steps in the YAO Agent system.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Architecture](#architecture)
|
||||
- [Data Models](#data-models)
|
||||
- [Write Strategy](#write-strategy)
|
||||
- [API Interface](#api-interface)
|
||||
- [Usage Examples](#usage-examples)
|
||||
- [Related Documents](#related-documents)
|
||||
|
||||
## Overview
|
||||
|
||||
The chat storage system is designed to:
|
||||
|
||||
1. **Store user-visible messages** - All messages sent via `ctx.Send()`, including text, images, loading states, etc.
|
||||
2. **Support resume/retry** - Track execution steps to enable recovery from interruptions or failures
|
||||
3. **Efficient writes** - Batch message writes at request end
|
||||
|
||||
### Design Goals
|
||||
|
||||
| Goal | Solution |
|
||||
| ------------------------ | ------------------------------------------------ |
|
||||
| Complete chat history | Store final content of all `ctx.Send()` messages |
|
||||
| Resume from interruption | Track step status and input/output |
|
||||
| Retry failed operations | Store step input for re-execution |
|
||||
| Minimize database writes | Batch writes at request end |
|
||||
|
||||
### Non-Goals
|
||||
|
||||
- **Tracing/debugging** - Handled by separate [Trace module](../../trace/README.md)
|
||||
- **Streaming replay** - Not needed, history shows final content only
|
||||
- **Request tracking/billing** - Handled by [OpenAPI Request module](../../openapi/request/REQUEST_DESIGN.md)
|
||||
|
||||
### Relationship with OpenAPI Request
|
||||
|
||||
The Agent storage focuses on **chat content and execution state**, while request tracking (billing, rate limiting, auditing) is handled globally by the OpenAPI layer:
|
||||
|
||||
| Concern | Module | Table |
|
||||
| ------------------ | ----------------- | -------------------- |
|
||||
| Request tracking | `openapi/request` | `openapi_request` |
|
||||
| Billing (tokens) | `openapi/request` | `openapi_request` |
|
||||
| Rate limiting | `openapi/request` | - |
|
||||
| Chat conversations | `agent/store` | `agent_conversation` |
|
||||
| Chat messages | `agent/store` | `agent_message` |
|
||||
| Execution steps | `agent/store` | `agent_step` |
|
||||
|
||||
The `request_id` from OpenAPI middleware is passed to Agent and stored in messages/steps for correlation.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Chat Storage │
|
||||
├─────────────────────────────────────────────────────────────┤
|
||||
│ │
|
||||
│ ┌─────────────────┐ │
|
||||
│ │ Conversation │ Metadata: title, assistant, user │
|
||||
│ └────────┬────────┘ │
|
||||
│ │ │
|
||||
│ │ 1:N │
|
||||
│ ▼ │
|
||||
│ ┌─────────────────┐ │
|
||||
│ │ Message │ User-visible: type, props, role │
|
||||
│ └────────┬────────┘ │
|
||||
│ │ │
|
||||
│ │ N:N (via request_id) │
|
||||
│ ▼ │
|
||||
│ ┌─────────────────┐ │
|
||||
│ │ Step │ Execution: type, status, input/output │
|
||||
│ └─────────────────┘ │
|
||||
│ │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## Data Models
|
||||
|
||||
### 1. Conversation Table
|
||||
|
||||
Stores conversation metadata and session information.
|
||||
|
||||
**Table Name:** `agent_conversation`
|
||||
|
||||
| Column | Type | Nullable | Index | Description |
|
||||
| ----------------- | ----------- | -------- | ------ | ----------------------------------- |
|
||||
| `id` | ID | No | PK | Auto-increment primary key |
|
||||
| `conversation_id` | string(64) | No | Unique | Unique conversation identifier |
|
||||
| `title` | string(500) | Yes | - | Conversation title |
|
||||
| `assistant_id` | string(200) | No | Yes | Associated assistant ID |
|
||||
| `user_id` | string(200) | No | Yes | Owner user ID |
|
||||
| `team_id` | string(200) | Yes | Yes | Team ID for access control |
|
||||
| `mode` | string(50) | No | - | Conversation mode (default: "chat") |
|
||||
| `status` | enum | No | Yes | Status: `active`, `archived` |
|
||||
| `last_message_at` | timestamp | Yes | Yes | Timestamp of last message |
|
||||
| `metadata` | json | Yes | - | Additional metadata |
|
||||
| `created_at` | timestamp | No | Yes | Creation timestamp |
|
||||
| `updated_at` | timestamp | No | - | Last update timestamp |
|
||||
|
||||
**Indexes:**
|
||||
|
||||
| Name | Columns | Type |
|
||||
| -------------------- | ------------------- | ----- |
|
||||
| `idx_conv_user` | `user_id`, `status` | index |
|
||||
| `idx_conv_team` | `team_id`, `status` | index |
|
||||
| `idx_conv_assistant` | `assistant_id` | index |
|
||||
| `idx_conv_last_msg` | `last_message_at` | index |
|
||||
|
||||
### 2. Message Table
|
||||
|
||||
Stores user-visible messages (both user input and assistant responses).
|
||||
|
||||
**Table Name:** `agent_message`
|
||||
|
||||
| Column | Type | Nullable | Index | Description |
|
||||
| ----------------- | ----------- | -------- | ------ | ----------------------------------------- |
|
||||
| `id` | ID | No | PK | Auto-increment primary key |
|
||||
| `message_id` | string(64) | No | Unique | Unique message identifier |
|
||||
| `conversation_id` | string(64) | No | Yes | Parent conversation ID |
|
||||
| `request_id` | string(64) | Yes | Yes | Request ID for grouping |
|
||||
| `role` | enum | No | Yes | Role: `user`, `assistant` |
|
||||
| `type` | string(50) | No | - | Message type (text, image, loading, etc.) |
|
||||
| `props` | json | No | - | Message properties (content, url, etc.) |
|
||||
| `block_id` | string(64) | Yes | Yes | Block grouping ID |
|
||||
| `thread_id` | string(64) | Yes | Yes | Thread grouping ID |
|
||||
| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) |
|
||||
| `sequence` | integer | No | Yes | Message order within conversation |
|
||||
| `metadata` | json | Yes | - | Additional metadata |
|
||||
| `created_at` | timestamp | No | Yes | Creation timestamp |
|
||||
| `updated_at` | timestamp | No | - | Last update timestamp |
|
||||
|
||||
**Indexes:**
|
||||
|
||||
| Name | Columns | Type |
|
||||
| ------------------- | ----------------------------- | ----- |
|
||||
| `idx_msg_conv_seq` | `conversation_id`, `sequence` | index |
|
||||
| `idx_msg_request` | `request_id` | index |
|
||||
| `idx_msg_block` | `block_id` | index |
|
||||
| `idx_msg_assistant` | `assistant_id` | index |
|
||||
|
||||
**Message Types:**
|
||||
|
||||
| Type | Description | Props Example |
|
||||
| --------- | ----------------- | ------------------------------------------------- |
|
||||
| `text` | Text message | `{"content": "Hello world"}` |
|
||||
| `image` | Image message | `{"url": "...", "alt": "...", "caption": "..."}` |
|
||||
| `loading` | Loading indicator | `{"message": "Processing...", "done": false}` |
|
||||
| `error` | Error message | `{"message": "...", "code": "..."}` |
|
||||
| `action` | Action buttons | `{"buttons": [...]}` |
|
||||
| `file` | File attachment | `{"url": "...", "filename": "...", "size": 1024}` |
|
||||
|
||||
### 3. Step Table
|
||||
|
||||
Stores execution steps for resume/retry functionality.
|
||||
|
||||
**Table Name:** `agent_step`
|
||||
|
||||
| Column | Type | Nullable | Index | Description |
|
||||
| ----------------- | ----------- | -------- | ------ | -------------------------------- |
|
||||
| `id` | ID | No | PK | Auto-increment primary key |
|
||||
| `step_id` | string(64) | No | Unique | Unique step identifier |
|
||||
| `conversation_id` | string(64) | No | Yes | Parent conversation ID |
|
||||
| `request_id` | string(64) | No | Yes | Request ID |
|
||||
| `assistant_id` | string(200) | No | Yes | Assistant executing this step |
|
||||
| `stack_id` | string(64) | No | Yes | Stack node ID for this execution |
|
||||
| `stack_parent_id` | string(64) | Yes | Yes | Parent stack ID (for A2A calls) |
|
||||
| `stack_depth` | integer | No | - | Call depth (0=root, 1+=nested) |
|
||||
| `type` | enum | No | Yes | Step type |
|
||||
| `status` | enum | No | Yes | Step status |
|
||||
| `input` | json | Yes | - | Step input data |
|
||||
| `output` | json | Yes | - | Step output data |
|
||||
| `error` | text | Yes | - | Error message if failed |
|
||||
| `sequence` | integer | No | Yes | Step order within request |
|
||||
| `metadata` | json | Yes | - | Additional metadata |
|
||||
| `created_at` | timestamp | No | Yes | Creation timestamp |
|
||||
| `updated_at` | timestamp | No | - | Last update timestamp |
|
||||
|
||||
**Step Types:**
|
||||
|
||||
| Type | Description | Input | Output |
|
||||
| ------------- | --------------------- | ---------------------- | ------------------------------------- |
|
||||
| `input` | User input received | `{messages: [...]}` | - |
|
||||
| `hook_create` | Create hook execution | `{messages: [...]}` | `{messages: [...], ...}` |
|
||||
| `llm` | LLM completion call | `{messages: [...]}` | `{content: "...", tool_calls: [...]}` |
|
||||
| `tool` | Tool/MCP execution | `{server, tool, args}` | `{result: ...}` |
|
||||
| `hook_next` | Next hook execution | `{completion, tools}` | `{data: ...}` |
|
||||
| `delegate` | A2A delegation | `{agent_id, messages}` | `{response: ...}` |
|
||||
|
||||
**Step Status:**
|
||||
|
||||
| Status | Description | Can Resume |
|
||||
| ------------- | --------------------- | -------------- |
|
||||
| `pending` | Not started | Yes |
|
||||
| `running` | In progress | Yes (restart) |
|
||||
| `completed` | Finished successfully | No |
|
||||
| `failed` | Failed with error | Yes (retry) |
|
||||
| `interrupted` | User interrupted | Yes (continue) |
|
||||
|
||||
**Indexes:**
|
||||
|
||||
| Name | Columns | Type |
|
||||
| -------------------- | ------------------------ | ----- |
|
||||
| `idx_step_conv` | `conversation_id` | index |
|
||||
| `idx_step_request` | `request_id`, `sequence` | index |
|
||||
| `idx_step_status` | `status` | index |
|
||||
| `idx_step_stack` | `stack_id` | index |
|
||||
| `idx_step_parent` | `stack_parent_id` | index |
|
||||
| `idx_step_assistant` | `assistant_id` | index |
|
||||
|
||||
## Write Strategy
|
||||
|
||||
### Two-Write Strategy
|
||||
|
||||
All data is buffered in memory during execution and written to database only **twice**:
|
||||
|
||||
1. **Write 1 (Entry)**: When `Stream()` starts - save user input message
|
||||
2. **Write 2 (Exit)**: When `Stream()` exits - batch save all assistant messages and steps
|
||||
|
||||
**Note**: Request tracking (status, tokens, duration) is handled by [OpenAPI Request Middleware](../../openapi/request/REQUEST_DESIGN.md).
|
||||
|
||||
```
|
||||
Stream() Entry
|
||||
│
|
||||
├── 【Write 1】Save user input
|
||||
│ - User message (role=user)
|
||||
│
|
||||
├── Execution (all in memory)
|
||||
│ - ctx.Send() → messageBuffer
|
||||
│ - ctx.Append() → update messageBuffer
|
||||
│ - ctx.Replace() → update messageBuffer
|
||||
│ - Each step → stepBuffer
|
||||
│
|
||||
└── 【Write 2】Save final state (via defer)
|
||||
- Batch write all assistant messages
|
||||
- Batch write all steps (with final status)
|
||||
- Update token usage in openapi_request (via request_id)
|
||||
```
|
||||
|
||||
### Write Points
|
||||
|
||||
| Event | Message Table | Step Table |
|
||||
| ---------------- | -------------------- | ----------------------------------------- |
|
||||
| Stream entry | Write 1 (user input) | - |
|
||||
| During execution | Buffer in memory | Buffer in memory |
|
||||
| **Stream exit** | **Batch write all** | **Batch write all (status=completed)** |
|
||||
| On interrupt | Batch write buffered | Batch write buffered (status=interrupted) |
|
||||
| On error | Batch write buffered | Batch write buffered (status=failed) |
|
||||
|
||||
### Why Two Writes?
|
||||
|
||||
| Scenario | What Happens | Data Safe? |
|
||||
| ------------------ | ----------------------------------- | ---------- |
|
||||
| Normal completion | `defer` triggers → Write 2 executes | ✅ |
|
||||
| User clicks stop | `defer` triggers → Write 2 executes | ✅ |
|
||||
| LLM timeout | `defer` triggers → Write 2 executes | ✅ |
|
||||
| Tool failure | `defer` triggers → Write 2 executes | ✅ |
|
||||
| Network disconnect | `defer` triggers → Write 2 executes | ✅ |
|
||||
| Process crash | Service is down, user must retry | N/A |
|
||||
|
||||
**Note**: Process crash is a catastrophic failure handled at infrastructure level, not application level.
|
||||
|
||||
### Write Count Comparison
|
||||
|
||||
For a typical request: user input → hook_create → llm → tool → llm → hook_next → 5 messages
|
||||
|
||||
| Strategy | Database Writes | Notes |
|
||||
| ---------------------- | --------------- | ------------------ |
|
||||
| Write per operation | 1 + 5 + 5 = 11 | One write per step |
|
||||
| **Two-write strategy** | **2** | Entry + Exit only |
|
||||
|
||||
### Implementation
|
||||
|
||||
````go
|
||||
func (ast *Assistant) Stream(ctx, inputMessages, options) {
|
||||
// ========== Write 1: Entry ==========
|
||||
userMsg := createUserMessage(ctx, inputMessages)
|
||||
chatStore.SaveMessages(ctx.ChatID, []*Message{userMsg})
|
||||
|
||||
// ========== Memory Buffers ==========
|
||||
messageBuffer := NewMessageBuffer()
|
||||
stepBuffer := NewStepBuffer()
|
||||
|
||||
// Track current step for error handling
|
||||
var currentStep *Step
|
||||
|
||||
defer func() {
|
||||
// ========== Write 2: Exit (always executes) ==========
|
||||
// Determine final status for incomplete steps
|
||||
finalStatus := "completed"
|
||||
if ctx.IsInterrupted() {
|
||||
finalStatus = "interrupted"
|
||||
}
|
||||
if r := recover(); r != nil {
|
||||
finalStatus = "failed"
|
||||
}
|
||||
|
||||
// Update status of any incomplete step
|
||||
if currentStep != nil && currentStep.Status == "running" {
|
||||
currentStep.Status = finalStatus
|
||||
}
|
||||
|
||||
// Batch write all buffered data
|
||||
chatStore.SaveMessages(ctx.ChatID, messageBuffer.GetAll())
|
||||
chatStore.SaveSteps(stepBuffer.GetAll())
|
||||
|
||||
// Update token usage in OpenAPI request record
|
||||
if ctx.RequestID != "" && completionResponse != nil {
|
||||
request.UpdateTokenUsage(
|
||||
ctx.RequestID,
|
||||
completionResponse.Usage.PromptTokens,
|
||||
completionResponse.Usage.CompletionTokens,
|
||||
)
|
||||
}
|
||||
}()
|
||||
|
||||
// ========== Execution (all in memory) ==========
|
||||
// Note: request_id = ctx.RequestID (from OpenAPI middleware)
|
||||
|
||||
// hook_create
|
||||
currentStep = stepBuffer.Add(createStep(ctx, "hook_create", "running", input, nil))
|
||||
createResponse := ast.HookScript.Create(...)
|
||||
currentStep.Output = createResponse
|
||||
currentStep.Status = "completed"
|
||||
|
||||
// llm
|
||||
currentStep = stepBuffer.Add(createStep(ctx, "llm", "running", messages, nil))
|
||||
completionResponse := ast.executeLLMStream(...)
|
||||
currentStep.Output = completionResponse
|
||||
currentStep.Status = "completed"
|
||||
|
||||
// tool (if any)
|
||||
for _, toolCall := range completionResponse.ToolCalls {
|
||||
currentStep = stepBuffer.Add(createStep(ctx, "tool", "running", toolCall, nil))
|
||||
result := executeToolCall(toolCall)
|
||||
currentStep.Output = result
|
||||
currentStep.Status = "completed"
|
||||
}
|
||||
|
||||
// hook_next
|
||||
currentStep = stepBuffer.Add(createStep(ctx, "hook_next", "running", payload, nil))
|
||||
nextResponse := ast.HookScript.Next(...)
|
||||
currentStep.Output = nextResponse
|
||||
currentStep.Status = "completed"
|
||||
currentStep = nil // All done
|
||||
|
||||
// Messages are automatically buffered via ctx.Send()
|
||||
}
|
||||
|
||||
// createStep creates a step with context information
|
||||
func createStep(ctx *Context, stepType, status string, input, output interface{}) *Step {
|
||||
return &Step{
|
||||
StepID: generateID(),
|
||||
ConversationID: ctx.ChatID, // ChatID = conversation_id
|
||||
RequestID: ctx.RequestID, // From OpenAPI middleware
|
||||
AssistantID: ctx.AssistantID,
|
||||
StackID: ctx.Stack.ID,
|
||||
StackParentID: ctx.Stack.ParentID,
|
||||
StackDepth: ctx.Stack.Depth,
|
||||
Type: stepType,
|
||||
Status: status,
|
||||
Input: input,
|
||||
Output: output,
|
||||
Sequence: nextSequence(),
|
||||
}
|
||||
}
|
||||
|
||||
## API Interface
|
||||
|
||||
### ChatStore Interface
|
||||
|
||||
```go
|
||||
// ChatStore defines the chat storage interface
|
||||
type ChatStore interface {
|
||||
// Conversation Management
|
||||
CreateConversation(conv *Conversation) error
|
||||
GetConversation(conversationID string) (*Conversation, error)
|
||||
UpdateConversation(conversationID string, updates map[string]interface{}) error
|
||||
DeleteConversation(conversationID string) error
|
||||
ListConversations(filter ConversationFilter) (*ConversationList, error)
|
||||
|
||||
// Message Management
|
||||
SaveMessages(conversationID string, messages []*Message) error
|
||||
GetMessages(conversationID string, filter MessageFilter) ([]*Message, error)
|
||||
UpdateMessage(messageID string, updates map[string]interface{}) error
|
||||
DeleteMessages(conversationID string, messageIDs []string) error
|
||||
|
||||
// Step Management
|
||||
SaveStep(step *Step) error
|
||||
UpdateStep(stepID string, updates map[string]interface{}) error
|
||||
GetSteps(requestID string) ([]*Step, error)
|
||||
GetLastIncompleteStep(conversationID string) (*Step, error)
|
||||
}
|
||||
````
|
||||
|
||||
### Data Structures
|
||||
|
||||
```go
|
||||
// Conversation represents a chat conversation
|
||||
type Conversation struct {
|
||||
ConversationID string `json:"conversation_id"`
|
||||
Title string `json:"title,omitempty"`
|
||||
AssistantID string `json:"assistant_id"`
|
||||
UserID string `json:"user_id"`
|
||||
TeamID string `json:"team_id,omitempty"`
|
||||
Mode string `json:"mode"`
|
||||
Status string `json:"status"`
|
||||
LastMessageAt *time.Time `json:"last_message_at,omitempty"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Message represents a chat message
|
||||
type Message struct {
|
||||
MessageID string `json:"message_id"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
RequestID string `json:"request_id,omitempty"`
|
||||
Role string `json:"role"`
|
||||
Type string `json:"type"`
|
||||
Props map[string]interface{} `json:"props"`
|
||||
BlockID string `json:"block_id,omitempty"`
|
||||
ThreadID string `json:"thread_id,omitempty"`
|
||||
AssistantID string `json:"assistant_id,omitempty"`
|
||||
Sequence int `json:"sequence"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Step represents an execution step
|
||||
type Step struct {
|
||||
StepID string `json:"step_id"`
|
||||
ConversationID string `json:"conversation_id"`
|
||||
RequestID string `json:"request_id"`
|
||||
AssistantID string `json:"assistant_id"`
|
||||
StackID string `json:"stack_id"`
|
||||
StackParentID string `json:"stack_parent_id,omitempty"`
|
||||
StackDepth int `json:"stack_depth"`
|
||||
Type string `json:"type"`
|
||||
Status string `json:"status"`
|
||||
Input map[string]interface{} `json:"input,omitempty"`
|
||||
Output map[string]interface{} `json:"output,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Sequence int `json:"sequence"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
```
|
||||
|
||||
### Filter Structures
|
||||
|
||||
```go
|
||||
// ConversationFilter for listing conversations
|
||||
type ConversationFilter struct {
|
||||
UserID string `json:"user_id,omitempty"`
|
||||
TeamID string `json:"team_id,omitempty"`
|
||||
AssistantID string `json:"assistant_id,omitempty"`
|
||||
Status string `json:"status,omitempty"`
|
||||
Keywords string `json:"keywords,omitempty"`
|
||||
Page int `json:"page,omitempty"`
|
||||
PageSize int `json:"pagesize,omitempty"`
|
||||
}
|
||||
|
||||
// MessageFilter for listing messages
|
||||
type MessageFilter struct {
|
||||
RequestID string `json:"request_id,omitempty"`
|
||||
Role string `json:"role,omitempty"`
|
||||
BlockID string `json:"block_id,omitempty"`
|
||||
Limit int `json:"limit,omitempty"`
|
||||
Offset int `json:"offset,omitempty"`
|
||||
}
|
||||
|
||||
// ConversationList paginated response
|
||||
type ConversationList struct {
|
||||
Data []*Conversation `json:"data"`
|
||||
Page int `json:"page"`
|
||||
PageSize int `json:"pagesize"`
|
||||
PageCount int `json:"pagecount"`
|
||||
Total int `json:"total"`
|
||||
}
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### 1. Normal Request Flow
|
||||
|
||||
See [Write Strategy - Implementation](#implementation) for the complete flow with two-write strategy.
|
||||
|
||||
### 2. Load Chat History
|
||||
|
||||
```go
|
||||
// Get conversation list
|
||||
convs, _ := chatStore.ListConversations(ConversationFilter{
|
||||
UserID: "user123",
|
||||
Status: "active",
|
||||
Page: 1,
|
||||
PageSize: 20,
|
||||
})
|
||||
|
||||
// Get messages for a conversation
|
||||
messages, _ := chatStore.GetMessages("conv_123", MessageFilter{
|
||||
Limit: 100,
|
||||
})
|
||||
|
||||
// Return to frontend
|
||||
return map[string]interface{}{
|
||||
"conversation": conv,
|
||||
"messages": messages,
|
||||
}
|
||||
```
|
||||
|
||||
### 3. Resume from Interruption
|
||||
|
||||
```go
|
||||
func (ast *Assistant) Resume(ctx *Context) error {
|
||||
// 1. Find last incomplete step
|
||||
step, _ := chatStore.GetLastIncompleteStep(ctx.ConversationID)
|
||||
if step == nil {
|
||||
return nil // Nothing to resume
|
||||
}
|
||||
|
||||
// 2. Check if this is an A2A nested call
|
||||
if step.StackDepth > 0 {
|
||||
// Need to rebuild the call stack
|
||||
return ast.ResumeNestedCall(ctx, step)
|
||||
}
|
||||
|
||||
// 3. Resume based on step type
|
||||
switch step.Type {
|
||||
case "llm":
|
||||
// Re-execute LLM call with saved input
|
||||
messages := step.Input["messages"].([]Message)
|
||||
return ast.executeLLMStream(ctx, messages, ...)
|
||||
|
||||
case "tool":
|
||||
// Retry tool call
|
||||
return ast.retryToolCall(ctx, step)
|
||||
|
||||
case "hook_next":
|
||||
// Re-execute hook
|
||||
return ast.executeHookNext(ctx, step.Input)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
```
|
||||
|
||||
### 4. Resume A2A Nested Calls
|
||||
|
||||
For agent-to-agent (A2A) recursive calls, the stack information is essential for proper recovery.
|
||||
|
||||
```go
|
||||
func (ast *Assistant) ResumeNestedCall(ctx *Context, step *Step) error {
|
||||
// 1. Rebuild the call stack from root to interrupted point
|
||||
stackPath, _ := chatStore.GetStackPath(step.StackID)
|
||||
// stackPath: [root_stack_id, parent_stack_id, ..., current_stack_id]
|
||||
|
||||
// 2. Get all steps for each stack level
|
||||
for _, stackID := range stackPath {
|
||||
steps, _ := chatStore.GetStepsByStackID(stackID)
|
||||
// Restore context for each level
|
||||
}
|
||||
|
||||
// 3. Resume from the interrupted assistant
|
||||
targetAssistant := assistant.Select(step.AssistantID)
|
||||
return targetAssistant.Stream(ctx, step.Input["messages"], ...)
|
||||
}
|
||||
```
|
||||
|
||||
### 4. Handle Interruption
|
||||
|
||||
Interruption is handled automatically by the `defer` block in the two-write strategy. When `ctx.IsInterrupted()` returns true, the status is set to `interrupted` and all buffered data is saved.
|
||||
|
||||
```go
|
||||
// Inside the defer block (see Write Strategy - Implementation)
|
||||
if ctx.IsInterrupted() {
|
||||
status = "interrupted"
|
||||
}
|
||||
// Then batch write all buffered messages and steps
|
||||
```
|
||||
|
||||
## A2A (Agent-to-Agent) Call Example
|
||||
|
||||
When Assistant A delegates to Assistant B, the step records look like:
|
||||
|
||||
```
|
||||
Request: User asks "analyze this data and visualize it"
|
||||
|
||||
Step Records:
|
||||
┌─────┬─────────────┬─────────────┬──────────┬────────────┬───────┬─────────────┐
|
||||
│ seq │ assistant │ stack_id │ parent │ depth │ type │ status │
|
||||
├─────┼─────────────┼─────────────┼──────────┼────────────┼───────┼─────────────┤
|
||||
│ 1 │ analyzer │ stk_001 │ null │ 0 │ input │ completed │
|
||||
│ 2 │ analyzer │ stk_001 │ null │ 0 │ llm │ completed │
|
||||
│ 3 │ analyzer │ stk_001 │ null │ 0 │ delegate │ running │ ← delegating
|
||||
│ 4 │ visualizer │ stk_002 │ stk_001 │ 1 │ input │ completed │
|
||||
│ 5 │ visualizer │ stk_002 │ stk_001 │ 1 │ llm │ interrupted │ ← interrupted here
|
||||
└─────┴─────────────┴─────────────┴──────────┴────────────┴───────┴─────────────┘
|
||||
|
||||
Resume Flow:
|
||||
1. Find step with status="interrupted" → step 5
|
||||
2. Check stack_depth=1 → nested call
|
||||
3. Get stack path: [stk_001, stk_002]
|
||||
4. Resume visualizer assistant with step 5's input
|
||||
5. When visualizer completes, update step 3 (delegate) to completed
|
||||
```
|
||||
|
||||
## Migration Notes
|
||||
|
||||
### From Old Schema
|
||||
|
||||
The old `agent_history` and `agent_chat` tables are replaced by:
|
||||
|
||||
| Old Table | New Table | Notes |
|
||||
| --------------- | -------------------- | ------------------------------------------------------ |
|
||||
| `agent_chat` | `agent_conversation` | Similar structure, added `mode`, `metadata` |
|
||||
| `agent_history` | `agent_message` | Changed to store `type`/`props` instead of raw content |
|
||||
| - | `agent_step` | New table for execution tracking |
|
||||
|
||||
### Data Migration
|
||||
|
||||
```sql
|
||||
-- Migrate conversations
|
||||
INSERT INTO agent_conversation (conversation_id, title, assistant_id, ...)
|
||||
SELECT chat_id, title, assistant_id, ...
|
||||
FROM agent_chat;
|
||||
|
||||
-- Migrate messages (simplified, actual migration needs content transformation)
|
||||
INSERT INTO agent_message (message_id, conversation_id, role, type, props, ...)
|
||||
SELECT id, chat_id, role, 'text', JSON_OBJECT('content', content), ...
|
||||
FROM agent_history;
|
||||
```
|
||||
|
||||
## Related Documents
|
||||
|
||||
- [OpenAPI Request Design](../../openapi/request/REQUEST_DESIGN.md) - Global request tracking, billing, rate limiting
|
||||
- [Trace Module](../../trace/README.md) - Detailed execution tracing for debugging
|
||||
- [Agent Context](../context/README.md) - Context and message handling
|
||||
874
openapi/request/REQUEST_DESIGN.md
Normal file
874
openapi/request/REQUEST_DESIGN.md
Normal file
|
|
@ -0,0 +1,874 @@
|
|||
# OpenAPI Request Design
|
||||
|
||||
This document describes the design for global request tracking, billing, rate limiting, and auditing in the YAO OpenAPI layer.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Architecture](#architecture)
|
||||
- [Storage Strategy](#storage-strategy)
|
||||
- [Data Model](#data-model)
|
||||
- [Middleware Design](#middleware-design)
|
||||
- [Rate Limiting](#rate-limiting)
|
||||
- [Billing Integration](#billing-integration)
|
||||
- [API Interface](#api-interface)
|
||||
- [Integration with Services](#integration-with-services)
|
||||
|
||||
## Overview
|
||||
|
||||
The Request module provides a unified layer for:
|
||||
|
||||
1. **Request Tracking** - Record all API requests with unique IDs
|
||||
2. **Billing** - Track token usage and API calls for billing
|
||||
3. **Rate Limiting** - Enforce request limits per user/team
|
||||
4. **Auditing** - Provide audit trail for compliance
|
||||
|
||||
### Design Goals
|
||||
|
||||
| Goal | Solution |
|
||||
| ------------------- | ------------------------------------------------ |
|
||||
| Unified tracking | Single middleware for all API endpoints |
|
||||
| Accurate billing | Token usage updated by services after completion |
|
||||
| Flexible rate limit | Configurable limits per user/team/endpoint |
|
||||
| Low overhead | KV for real-time, SQL for archive |
|
||||
|
||||
### Scope
|
||||
|
||||
| In Scope | Out of Scope |
|
||||
| ------------------------ | ------------------------------ |
|
||||
| All `/api/*` endpoints | Static file serving |
|
||||
| Token usage tracking | Detailed request/response logs |
|
||||
| Rate limiting | Request body storage |
|
||||
| Request duration metrics | Response caching |
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ HTTP Request │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ OAuth Guard │
|
||||
│ - Token validation │
|
||||
│ - Set AuthorizedInfo in context │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Request Middleware │
|
||||
│ - Generate request_id │
|
||||
│ - KV: Rate limit check │
|
||||
│ - KV: Quota check │
|
||||
│ - KV: Request status tracking │
|
||||
│ - Async: Archive to SQL │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ Service Handlers │
|
||||
│ ┌─────────┐ ┌─────────┐ ┌─────────┐ ┌─────────┐ │
|
||||
│ │ Agent │ │ KB │ │ LLM │ │ File │ ... │
|
||||
│ └─────────┘ └─────────┘ └─────────┘ └─────────┘ │
|
||||
│ │ │
|
||||
│ └── Update token usage via request_id │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## Storage Strategy
|
||||
|
||||
### Two-Layer Storage
|
||||
|
||||
| Layer | Storage | Purpose | TTL |
|
||||
| -------------- | -------- | ---------------------------- | --------- |
|
||||
| **Real-time** | KV/Redis | Rate limiting, quota, status | 1h - 7d |
|
||||
| **Persistent** | SQL | Archive, billing, audit | Permanent |
|
||||
|
||||
### Why Hybrid?
|
||||
|
||||
| Scenario | KV (Redis) | SQL |
|
||||
| ----------------- | ----------------- | ---------------- |
|
||||
| Rate limit check | ⚡ < 1ms | ❌ Too slow |
|
||||
| Quota check | ⚡ < 1ms | ❌ Too slow |
|
||||
| Request status | ⚡ Fast update | ❌ Too slow |
|
||||
| Billing report | ❌ No aggregation | ✅ SUM/GROUP BY |
|
||||
| Audit query | ❌ No persistence | ✅ Full history |
|
||||
| Complex filtering | ❌ Key-only | ✅ WHERE clauses |
|
||||
|
||||
### KV Keys Design
|
||||
|
||||
```
|
||||
# Rate Limiting (TTL: 60s)
|
||||
ratelimit:user:{user_id}:{service} → count
|
||||
ratelimit:team:{team_id}:{service} → count
|
||||
ratelimit:ip:{ip} → count
|
||||
|
||||
# Request Status (TTL: 1h)
|
||||
request:{request_id} → {status, service, created_at, ...}
|
||||
|
||||
# Token Usage - Daily (TTL: 7d)
|
||||
tokens:user:{user_id}:{YYYY-MM-DD} → {input, output, total}
|
||||
tokens:team:{team_id}:{YYYY-MM-DD} → {input, output, total}
|
||||
|
||||
# Quota (TTL: 24h for daily, 30d for monthly)
|
||||
quota:user:{user_id}:daily → remaining_tokens
|
||||
quota:team:{team_id}:monthly → remaining_tokens
|
||||
```
|
||||
|
||||
### Data Flow
|
||||
|
||||
```
|
||||
Request arrives
|
||||
│
|
||||
├── 1. KV: Rate limit check
|
||||
│ INCR ratelimit:user:{id}:{service}
|
||||
│ if > limit → 429 Too Many Requests
|
||||
│
|
||||
├── 2. KV: Quota check
|
||||
│ GET quota:user:{id}:daily
|
||||
│ if <= 0 → 429 Quota Exceeded
|
||||
│
|
||||
├── 3. KV: Record request status
|
||||
│ SET request:{id} {status: "running", ...} EX 3600
|
||||
│
|
||||
├── 4. Execute request...
|
||||
│
|
||||
├── 5. KV: Update tokens
|
||||
│ HINCRBY tokens:user:{id}:{date} input {n}
|
||||
│ HINCRBY tokens:user:{id}:{date} output {n}
|
||||
│ DECRBY quota:user:{id}:daily {total}
|
||||
│
|
||||
├── 6. KV: Update request status
|
||||
│ SET request:{id} {status: "completed", duration_ms: ...}
|
||||
│
|
||||
└── 7. Async: Archive to SQL
|
||||
INSERT INTO openapi_request ...
|
||||
|
||||
```
|
||||
|
||||
## Data Model
|
||||
|
||||
### KV Data Structures
|
||||
|
||||
#### Rate Limit Counter
|
||||
|
||||
```go
|
||||
// Key: ratelimit:{type}:{id}:{service}
|
||||
// Value: integer count
|
||||
// TTL: 60 seconds (sliding window)
|
||||
|
||||
type RateLimitKey struct {
|
||||
Type string // "user", "team", "ip"
|
||||
ID string // user_id, team_id, or IP
|
||||
Service string // "agent", "kb", "llm", etc.
|
||||
}
|
||||
|
||||
func (k RateLimitKey) String() string {
|
||||
return fmt.Sprintf("ratelimit:%s:%s:%s", k.Type, k.ID, k.Service)
|
||||
}
|
||||
```
|
||||
|
||||
#### Request Status
|
||||
|
||||
```go
|
||||
// Key: request:{request_id}
|
||||
// Value: JSON object
|
||||
// TTL: 1 hour
|
||||
|
||||
type RequestStatus struct {
|
||||
RequestID string `json:"request_id"`
|
||||
UserID string `json:"user_id"`
|
||||
TeamID string `json:"team_id,omitempty"`
|
||||
Service string `json:"service"`
|
||||
ResourceID string `json:"resource_id,omitempty"`
|
||||
Status string `json:"status"` // running, completed, failed
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
CompletedAt time.Time `json:"completed_at,omitempty"`
|
||||
DurationMs int64 `json:"duration_ms,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
```
|
||||
|
||||
#### Token Usage (Daily)
|
||||
|
||||
```go
|
||||
// Key: tokens:{type}:{id}:{date}
|
||||
// Value: Hash {input, output, total}
|
||||
// TTL: 7 days
|
||||
|
||||
type TokenUsage struct {
|
||||
Input int64 `json:"input"`
|
||||
Output int64 `json:"output"`
|
||||
Total int64 `json:"total"`
|
||||
}
|
||||
```
|
||||
|
||||
#### Quota
|
||||
|
||||
```go
|
||||
// Key: quota:{type}:{id}:{period}
|
||||
// Value: remaining tokens (integer)
|
||||
// TTL: 24h (daily) or 30d (monthly)
|
||||
|
||||
type QuotaKey struct {
|
||||
Type string // "user", "team"
|
||||
ID string
|
||||
Period string // "daily", "monthly"
|
||||
}
|
||||
```
|
||||
|
||||
### SQL Table (Archive)
|
||||
|
||||
**Table Name:** `openapi_request`
|
||||
|
||||
**Purpose:** Long-term storage for billing reports, audit logs, and analytics.
|
||||
|
||||
| Column | Type | Nullable | Index | Description |
|
||||
| --------------- | ----------- | -------- | ------ | ------------------------------------------------ |
|
||||
| `id` | ID | No | PK | Auto-increment primary key |
|
||||
| `request_id` | string(64) | No | Unique | Unique request identifier |
|
||||
| `user_id` | string(200) | No | Yes | User ID from auth |
|
||||
| `team_id` | string(200) | Yes | Yes | Team ID from auth |
|
||||
| `session_id` | string(200) | Yes | Yes | Session ID |
|
||||
| `endpoint` | string(200) | No | Yes | API endpoint path |
|
||||
| `method` | string(10) | No | - | HTTP method (GET, POST, etc.) |
|
||||
| `service` | string(50) | No | Yes | Service type: `agent`, `kb`, `llm`, `file`, etc. |
|
||||
| `resource_id` | string(200) | Yes | Yes | Resource ID (assistant_id, collection_id, etc.) |
|
||||
| `status` | enum | No | Yes | `pending`, `running`, `completed`, `failed` |
|
||||
| `status_code` | integer | Yes | - | HTTP response status code |
|
||||
| `referer` | string(50) | Yes | - | Request source (api, jssdk, agent, etc.) |
|
||||
| `client_type` | string(50) | Yes | - | Client type (web, ios, android, etc.) |
|
||||
| `client_ip` | string(50) | Yes | Yes | Client IP address |
|
||||
| `input_tokens` | integer | Yes | - | Input token count (LLM calls) |
|
||||
| `output_tokens` | integer | Yes | - | Output token count (LLM calls) |
|
||||
| `total_tokens` | integer | Yes | Yes | Total token count |
|
||||
| `duration_ms` | integer | Yes | Yes | Request duration in milliseconds |
|
||||
| `error` | text | Yes | - | Error message if failed |
|
||||
| `metadata` | json | Yes | - | Additional metadata |
|
||||
| `created_at` | timestamp | No | Yes | Request start time |
|
||||
| `completed_at` | timestamp | Yes | Yes | Request completion time |
|
||||
|
||||
**Indexes:**
|
||||
|
||||
| Name | Columns | Type | Purpose |
|
||||
| ------------------ | --------------------------------------- | ----- | ------------------------ |
|
||||
| `idx_req_user` | `user_id`, `created_at` | index | User request history |
|
||||
| `idx_req_team` | `team_id`, `created_at` | index | Team request history |
|
||||
| `idx_req_endpoint` | `endpoint`, `created_at` | index | Endpoint analytics |
|
||||
| `idx_req_service` | `service`, `created_at` | index | Service analytics |
|
||||
| `idx_req_status` | `status` | index | Find incomplete requests |
|
||||
| `idx_req_billing` | `team_id`, `created_at`, `total_tokens` | index | Billing queries |
|
||||
| `idx_req_ip` | `client_ip`, `created_at` | index | IP-based rate limiting |
|
||||
|
||||
### Service Types
|
||||
|
||||
| Service | Description | Resource ID Example |
|
||||
| ------- | -------------------- | ------------------- |
|
||||
| `agent` | Chat/Agent API | `assistant_id` |
|
||||
| `kb` | Knowledge Base API | `collection_id` |
|
||||
| `llm` | Direct LLM API | `connector_id` |
|
||||
| `file` | File upload/download | `file_id` |
|
||||
| `user` | User management | `user_id` |
|
||||
| `team` | Team management | `team_id` |
|
||||
| `mcp` | MCP server calls | `server_id` |
|
||||
|
||||
### Status Values
|
||||
|
||||
| Status | Description | Set By |
|
||||
| ----------- | ------------------------------ | ---------- |
|
||||
| `pending` | Request received, not started | Middleware |
|
||||
| `running` | Request being processed | Middleware |
|
||||
| `completed` | Request completed successfully | Middleware |
|
||||
| `failed` | Request failed with error | Middleware |
|
||||
|
||||
## Middleware Design
|
||||
|
||||
### Request Flow
|
||||
|
||||
```
|
||||
Request arrives
|
||||
│
|
||||
├── 1. Generate request_id (uuid or nanoid)
|
||||
│
|
||||
├── 2. Set request_id in context and response header
|
||||
│ c.Set("request_id", requestID)
|
||||
│ c.Header("X-Request-ID", requestID)
|
||||
│
|
||||
├── 3. Get auth info from context (set by OAuth Guard)
|
||||
│ authInfo := authorized.GetInfo(c)
|
||||
│
|
||||
├── 4. Detect service type from endpoint
|
||||
│ service := detectService(c.FullPath())
|
||||
│
|
||||
├── 5. Create request record (async)
|
||||
│ status = "running"
|
||||
│
|
||||
├── 6. Check rate limits
|
||||
│ if exceeded → return 429, update status = "failed"
|
||||
│
|
||||
├── 7. Execute handler
|
||||
│ c.Next()
|
||||
│
|
||||
└── 8. Update request record (async)
|
||||
status = "completed" or "failed"
|
||||
duration_ms = time.Since(start)
|
||||
status_code = c.Writer.Status()
|
||||
```
|
||||
|
||||
### Implementation
|
||||
|
||||
```go
|
||||
package request
|
||||
|
||||
import (
|
||||
"time"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/yaoapp/yao/openapi/oauth/authorized"
|
||||
)
|
||||
|
||||
// Middleware creates the request tracking middleware
|
||||
func Middleware(kv KVStore, sql SQLStore) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
startTime := time.Now()
|
||||
|
||||
// 1. Generate request ID
|
||||
requestID := generateRequestID()
|
||||
c.Set("request_id", requestID)
|
||||
c.Header("X-Request-ID", requestID)
|
||||
|
||||
// 2. Get auth info
|
||||
authInfo := authorized.GetInfo(c)
|
||||
|
||||
// 3. Detect service and resource
|
||||
service := detectService(c.FullPath())
|
||||
resourceID := extractResourceID(c, service)
|
||||
|
||||
// 4. KV: Check rate limits (synchronous, must be fast)
|
||||
if err := checkRateLimit(kv, authInfo, service, c.ClientIP()); err != nil {
|
||||
c.AbortWithStatusJSON(429, gin.H{
|
||||
"error": "rate_limit_exceeded",
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// 5. KV: Check quota (synchronous)
|
||||
if err := checkQuota(kv, authInfo); err != nil {
|
||||
c.AbortWithStatusJSON(429, gin.H{
|
||||
"error": "quota_exceeded",
|
||||
"message": err.Error(),
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
// 6. KV: Record request status
|
||||
reqStatus := &RequestStatus{
|
||||
RequestID: requestID,
|
||||
UserID: authInfo.UserID,
|
||||
TeamID: authInfo.TeamID,
|
||||
Service: service,
|
||||
ResourceID: resourceID,
|
||||
Status: "running",
|
||||
CreatedAt: startTime,
|
||||
}
|
||||
kv.SetRequestStatus(requestID, reqStatus, time.Hour)
|
||||
|
||||
// 7. Execute handler
|
||||
c.Next()
|
||||
|
||||
// 8. KV: Update request status
|
||||
reqStatus.Status = "completed"
|
||||
reqStatus.CompletedAt = time.Now()
|
||||
reqStatus.DurationMs = time.Since(startTime).Milliseconds()
|
||||
if errMsg := getErrorFromContext(c); errMsg != "" {
|
||||
reqStatus.Status = "failed"
|
||||
reqStatus.Error = errMsg
|
||||
}
|
||||
kv.SetRequestStatus(requestID, reqStatus, time.Hour)
|
||||
|
||||
// 9. Async: Archive to SQL
|
||||
go func() {
|
||||
sql.Archive(&Request{
|
||||
RequestID: requestID,
|
||||
UserID: authInfo.UserID,
|
||||
TeamID: authInfo.TeamID,
|
||||
SessionID: authInfo.SessionID,
|
||||
Endpoint: c.FullPath(),
|
||||
Method: c.Request.Method,
|
||||
Service: service,
|
||||
ResourceID: resourceID,
|
||||
Status: reqStatus.Status,
|
||||
StatusCode: c.Writer.Status(),
|
||||
Referer: c.GetHeader("X-Yao-Referer"),
|
||||
ClientType: getClientType(c.GetHeader("User-Agent")),
|
||||
ClientIP: c.ClientIP(),
|
||||
DurationMs: reqStatus.DurationMs,
|
||||
Error: reqStatus.Error,
|
||||
CreatedAt: startTime,
|
||||
CompletedAt: &reqStatus.CompletedAt,
|
||||
})
|
||||
}()
|
||||
}
|
||||
}
|
||||
|
||||
// detectService determines the service type from endpoint
|
||||
func detectService(endpoint string) string {
|
||||
switch {
|
||||
case strings.HasPrefix(endpoint, "/api/chat"):
|
||||
return ServiceAgent
|
||||
case strings.HasPrefix(endpoint, "/api/agent"):
|
||||
return ServiceAgent
|
||||
case strings.HasPrefix(endpoint, "/api/kb"):
|
||||
return ServiceKB
|
||||
case strings.HasPrefix(endpoint, "/api/llm"):
|
||||
return ServiceLLM
|
||||
case strings.HasPrefix(endpoint, "/api/file"):
|
||||
return ServiceFile
|
||||
case strings.HasPrefix(endpoint, "/api/user"):
|
||||
return ServiceUser
|
||||
case strings.HasPrefix(endpoint, "/api/team"):
|
||||
return ServiceTeam
|
||||
case strings.HasPrefix(endpoint, "/api/mcp"):
|
||||
return ServiceMCP
|
||||
default:
|
||||
return ServiceOther
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Rate Limiting
|
||||
|
||||
### Configuration
|
||||
|
||||
```yaml
|
||||
# openapi.yml
|
||||
rate_limit:
|
||||
enabled: true
|
||||
|
||||
# Default limits (requests per minute)
|
||||
default:
|
||||
per_user: 60
|
||||
per_team: 300
|
||||
per_ip: 100
|
||||
|
||||
# Service-specific limits
|
||||
services:
|
||||
agent:
|
||||
per_user: 30
|
||||
per_team: 150
|
||||
llm:
|
||||
per_user: 20
|
||||
per_team: 100
|
||||
kb:
|
||||
per_user: 60
|
||||
per_team: 300
|
||||
|
||||
# Token limits (per day)
|
||||
tokens:
|
||||
per_user: 100000
|
||||
per_team: 1000000
|
||||
|
||||
# quota configuration
|
||||
quota:
|
||||
enabled: true
|
||||
|
||||
# Default quotas
|
||||
default:
|
||||
user_daily: 100000 # tokens per day
|
||||
team_monthly: 10000000 # tokens per month
|
||||
|
||||
|
||||
# Can be overridden per user/team in database
|
||||
```
|
||||
|
||||
### Rate Limit Check (KV-based)
|
||||
|
||||
```go
|
||||
func checkRateLimit(kv KVStore, authInfo *types.AuthorizedInfo, service, clientIP string) error {
|
||||
config := GetRateLimitConfig()
|
||||
if !config.Enabled {
|
||||
return nil
|
||||
}
|
||||
|
||||
// 1. Check per-user limit (INCR with TTL)
|
||||
userKey := fmt.Sprintf("ratelimit:user:%s:%s", authInfo.UserID, service)
|
||||
userCount, _ := kv.Incr(userKey, 60*time.Second) // TTL 60s
|
||||
if userCount > int64(config.GetUserLimit(service)) {
|
||||
return fmt.Errorf("user rate limit exceeded: %d requests per minute", config.GetUserLimit(service))
|
||||
}
|
||||
|
||||
// 2. Check per-team limit
|
||||
if authInfo.TeamID != "" {
|
||||
teamKey := fmt.Sprintf("ratelimit:team:%s:%s", authInfo.TeamID, service)
|
||||
teamCount, _ := kv.Incr(teamKey, 60*time.Second)
|
||||
if teamCount > int64(config.GetTeamLimit(service)) {
|
||||
return fmt.Errorf("team rate limit exceeded")
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Check per-IP limit
|
||||
ipKey := fmt.Sprintf("ratelimit:ip:%s", clientIP)
|
||||
ipCount, _ := kv.Incr(ipKey, 60*time.Second)
|
||||
if ipCount > int64(config.GetIPLimit()) {
|
||||
return fmt.Errorf("IP rate limit exceeded")
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
```
|
||||
|
||||
### Quota Check (KV-based)
|
||||
|
||||
```go
|
||||
func checkQuota(kv KVStore, authInfo *types.AuthorizedInfo) error {
|
||||
config := GetQuotaConfig()
|
||||
if !config.Enabled {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Check user daily quota
|
||||
userQuotaKey := fmt.Sprintf("quota:user:%s:daily", authInfo.UserID)
|
||||
remaining, exists := kv.Get(userQuotaKey)
|
||||
|
||||
if !exists {
|
||||
// Initialize quota for the day
|
||||
limit := config.GetUserDailyLimit(authInfo.UserID)
|
||||
kv.Set(userQuotaKey, limit, 24*time.Hour)
|
||||
remaining = limit
|
||||
}
|
||||
|
||||
if remaining <= 0 {
|
||||
return fmt.Errorf("daily token quota exceeded")
|
||||
}
|
||||
|
||||
// Check team monthly quota if applicable
|
||||
if authInfo.TeamID != "" {
|
||||
teamQuotaKey := fmt.Sprintf("quota:team:%s:monthly", authInfo.TeamID)
|
||||
teamRemaining, exists := kv.Get(teamQuotaKey)
|
||||
|
||||
if !exists {
|
||||
limit := config.GetTeamMonthlyLimit(authInfo.TeamID)
|
||||
kv.Set(teamQuotaKey, limit, 30*24*time.Hour)
|
||||
teamRemaining = limit
|
||||
}
|
||||
|
||||
if teamRemaining <= 0 {
|
||||
return fmt.Errorf("team monthly token quota exceeded")
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
```
|
||||
|
||||
## Billing Integration
|
||||
|
||||
### Token Usage Update
|
||||
|
||||
Services update token usage after completion. This updates both KV (real-time) and SQL (archive).
|
||||
|
||||
```go
|
||||
// Called by Agent/LLM services after completion
|
||||
func UpdateTokenUsage(kv KVStore, sql SQLStore, requestID string, userID, teamID string, input, output int) error {
|
||||
total := input + output
|
||||
date := time.Now().Format("2006-01-02")
|
||||
|
||||
// 1. KV: Update daily token usage
|
||||
userTokenKey := fmt.Sprintf("tokens:user:%s:%s", userID, date)
|
||||
kv.HIncrBy(userTokenKey, "input", int64(input))
|
||||
kv.HIncrBy(userTokenKey, "output", int64(output))
|
||||
kv.HIncrBy(userTokenKey, "total", int64(total))
|
||||
kv.Expire(userTokenKey, 7*24*time.Hour) // Keep for 7 days
|
||||
|
||||
if teamID != "" {
|
||||
teamTokenKey := fmt.Sprintf("tokens:team:%s:%s", teamID, date)
|
||||
kv.HIncrBy(teamTokenKey, "input", int64(input))
|
||||
kv.HIncrBy(teamTokenKey, "output", int64(output))
|
||||
kv.HIncrBy(teamTokenKey, "total", int64(total))
|
||||
kv.Expire(teamTokenKey, 7*24*time.Hour)
|
||||
}
|
||||
|
||||
// 2. KV: Deduct from quota
|
||||
userQuotaKey := fmt.Sprintf("quota:user:%s:daily", userID)
|
||||
kv.DecrBy(userQuotaKey, int64(total))
|
||||
|
||||
if teamID != "" {
|
||||
teamQuotaKey := fmt.Sprintf("quota:team:%s:monthly", teamID)
|
||||
kv.DecrBy(teamQuotaKey, int64(total))
|
||||
}
|
||||
|
||||
// 3. SQL: Update request record (async)
|
||||
go sql.UpdateTokens(requestID, input, output)
|
||||
|
||||
return nil
|
||||
}
|
||||
```
|
||||
|
||||
### Billing Queries
|
||||
|
||||
```sql
|
||||
-- Daily token usage by team
|
||||
SELECT
|
||||
DATE(created_at) as date,
|
||||
team_id,
|
||||
service,
|
||||
SUM(total_tokens) as tokens,
|
||||
COUNT(*) as requests
|
||||
FROM openapi_request
|
||||
WHERE team_id = ?
|
||||
AND created_at >= ? AND created_at < ?
|
||||
AND status = 'completed'
|
||||
GROUP BY DATE(created_at), team_id, service
|
||||
|
||||
-- Monthly billing summary
|
||||
SELECT
|
||||
team_id,
|
||||
service,
|
||||
SUM(total_tokens) as total_tokens,
|
||||
SUM(input_tokens) as input_tokens,
|
||||
SUM(output_tokens) as output_tokens,
|
||||
COUNT(*) as request_count,
|
||||
AVG(duration_ms) as avg_duration
|
||||
FROM openapi_request
|
||||
WHERE created_at >= ? AND created_at < ?
|
||||
AND status = 'completed'
|
||||
GROUP BY team_id, service
|
||||
|
||||
-- User quota check
|
||||
SELECT SUM(total_tokens) as used
|
||||
FROM openapi_request
|
||||
WHERE user_id = ?
|
||||
AND created_at >= CURDATE()
|
||||
AND status = 'completed'
|
||||
```
|
||||
|
||||
## API Interface
|
||||
|
||||
### KV Store Interface
|
||||
|
||||
```go
|
||||
// KVStore defines the KV storage interface for real-time operations
|
||||
type KVStore interface {
|
||||
// Basic operations
|
||||
Get(key string) (int64, bool)
|
||||
Set(key string, value int64, ttl time.Duration) error
|
||||
Incr(key string, ttl time.Duration) (int64, error)
|
||||
DecrBy(key string, delta int64) (int64, error)
|
||||
Expire(key string, ttl time.Duration) error
|
||||
Del(key string) error
|
||||
|
||||
// Hash operations (for token usage)
|
||||
HGet(key, field string) (int64, error)
|
||||
HSet(key, field string, value int64) error
|
||||
HIncrBy(key, field string, delta int64) (int64, error)
|
||||
HGetAll(key string) (map[string]int64, error)
|
||||
|
||||
// Request status (JSON)
|
||||
SetRequestStatus(requestID string, status *RequestStatus, ttl time.Duration) error
|
||||
GetRequestStatus(requestID string) (*RequestStatus, error)
|
||||
}
|
||||
```
|
||||
|
||||
### SQL Store Interface
|
||||
|
||||
```go
|
||||
// SQLStore defines the SQL storage interface for archiving and analytics
|
||||
type SQLStore interface {
|
||||
// Archive stores a completed request
|
||||
Archive(req *Request) error
|
||||
|
||||
// UpdateTokens updates token usage for a request
|
||||
UpdateTokens(requestID string, input, output int) error
|
||||
|
||||
// Get retrieves a request by ID
|
||||
Get(requestID string) (*Request, error)
|
||||
|
||||
// List lists requests with filters
|
||||
List(filter *RequestFilter) (*RequestList, error)
|
||||
|
||||
// GetUsage gets usage statistics
|
||||
GetUsage(filter *UsageFilter) (*UsageStats, error)
|
||||
}
|
||||
```
|
||||
|
||||
### Data Structures
|
||||
|
||||
```go
|
||||
// Request represents an API request record
|
||||
type Request struct {
|
||||
RequestID string `json:"request_id"`
|
||||
UserID string `json:"user_id"`
|
||||
TeamID string `json:"team_id,omitempty"`
|
||||
SessionID string `json:"session_id,omitempty"`
|
||||
Endpoint string `json:"endpoint"`
|
||||
Method string `json:"method"`
|
||||
Service string `json:"service"`
|
||||
ResourceID string `json:"resource_id,omitempty"`
|
||||
Status Status `json:"status"`
|
||||
StatusCode int `json:"status_code,omitempty"`
|
||||
Referer string `json:"referer,omitempty"`
|
||||
ClientType string `json:"client_type,omitempty"`
|
||||
ClientIP string `json:"client_ip,omitempty"`
|
||||
InputTokens int `json:"input_tokens,omitempty"`
|
||||
OutputTokens int `json:"output_tokens,omitempty"`
|
||||
TotalTokens int `json:"total_tokens,omitempty"`
|
||||
DurationMs int64 `json:"duration_ms,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
CompletedAt *time.Time `json:"completed_at,omitempty"`
|
||||
}
|
||||
|
||||
// CompletionInfo contains info for completing a request
|
||||
type CompletionInfo struct {
|
||||
StatusCode int
|
||||
DurationMs int64
|
||||
Error string
|
||||
}
|
||||
|
||||
// RequestFilter for listing requests
|
||||
type RequestFilter struct {
|
||||
UserID string `json:"user_id,omitempty"`
|
||||
TeamID string `json:"team_id,omitempty"`
|
||||
Service string `json:"service,omitempty"`
|
||||
Status Status `json:"status,omitempty"`
|
||||
StartTime time.Time `json:"start_time,omitempty"`
|
||||
EndTime time.Time `json:"end_time,omitempty"`
|
||||
Page int `json:"page,omitempty"`
|
||||
PageSize int `json:"pagesize,omitempty"`
|
||||
}
|
||||
|
||||
// UsageFilter for usage statistics
|
||||
type UsageFilter struct {
|
||||
UserID string `json:"user_id,omitempty"`
|
||||
TeamID string `json:"team_id,omitempty"`
|
||||
Service string `json:"service,omitempty"`
|
||||
StartTime time.Time `json:"start_time"`
|
||||
EndTime time.Time `json:"end_time"`
|
||||
GroupBy string `json:"group_by,omitempty"` // day, week, month
|
||||
}
|
||||
|
||||
// UsageStats contains usage statistics
|
||||
type UsageStats struct {
|
||||
TotalRequests int64 `json:"total_requests"`
|
||||
TotalTokens int64 `json:"total_tokens"`
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
AvgDurationMs float64 `json:"avg_duration_ms"`
|
||||
ByService map[string]int64 `json:"by_service,omitempty"`
|
||||
ByDay []DailyUsage `json:"by_day,omitempty"`
|
||||
}
|
||||
|
||||
type DailyUsage struct {
|
||||
Date string `json:"date"`
|
||||
Requests int64 `json:"requests"`
|
||||
Tokens int64 `json:"tokens"`
|
||||
}
|
||||
```
|
||||
|
||||
## Integration with Services
|
||||
|
||||
### Agent Service
|
||||
|
||||
```go
|
||||
// agent/context/openapi.go
|
||||
func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest, *Context, *Options, error) {
|
||||
// Get request ID from middleware
|
||||
requestID := c.GetString("request_id")
|
||||
|
||||
// Create context with request ID
|
||||
ctx := New(c.Request.Context(), authInfo, chatID)
|
||||
ctx.RequestID = requestID // Use global request_id
|
||||
|
||||
// ...
|
||||
}
|
||||
|
||||
// agent/assistant/agent.go
|
||||
func (ast *Assistant) Stream(ctx, inputMessages, options) {
|
||||
defer func() {
|
||||
// Update token usage in global request record
|
||||
if ctx.RequestID != "" && completionResponse != nil && completionResponse.Usage != nil {
|
||||
request.UpdateTokenUsage(
|
||||
ctx.RequestID,
|
||||
completionResponse.Usage.PromptTokens,
|
||||
completionResponse.Usage.CompletionTokens,
|
||||
)
|
||||
}
|
||||
}()
|
||||
|
||||
// ...
|
||||
}
|
||||
```
|
||||
|
||||
### KB Service
|
||||
|
||||
```go
|
||||
// kb/api/search.go
|
||||
func (api *API) Search(c *gin.Context) {
|
||||
requestID := c.GetString("request_id")
|
||||
|
||||
// Perform search...
|
||||
|
||||
// Update metadata if needed
|
||||
if requestID != "" {
|
||||
request.UpdateMetadata(requestID, map[string]interface{}{
|
||||
"results_count": len(results),
|
||||
"collection_id": collectionID,
|
||||
})
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Middleware Registration
|
||||
|
||||
```go
|
||||
// openapi/openapi.go
|
||||
func (s *OpenAPI) RegisterRoutes(r *gin.Engine) {
|
||||
api := r.Group("/api")
|
||||
|
||||
// 1. OAuth Guard (authentication)
|
||||
api.Use(oauth.Guard)
|
||||
|
||||
// 2. Request Middleware (tracking, rate limiting)
|
||||
api.Use(request.Middleware(requestStore))
|
||||
|
||||
// 3. Service routes
|
||||
s.registerAgentRoutes(api)
|
||||
s.registerKBRoutes(api)
|
||||
s.registerLLMRoutes(api)
|
||||
// ...
|
||||
}
|
||||
```
|
||||
|
||||
## Summary
|
||||
|
||||
### Components
|
||||
|
||||
| Component | Location | Responsibility |
|
||||
| ------------ | ------------------------------- | ---------------------------- |
|
||||
| KV Store | `openapi/request/kv.go` | Real-time: rate limit, quota |
|
||||
| SQL Store | `openapi/request/sql.go` | Archive: billing, audit |
|
||||
| Middleware | `openapi/request/middleware.go` | Track requests, orchestrate |
|
||||
| Rate Limiter | `openapi/request/ratelimit.go` | Enforce rate limits |
|
||||
| Types | `openapi/request/types.go` | Data structures |
|
||||
|
||||
### Storage Comparison
|
||||
|
||||
| Operation | KV (Redis) | SQL (Archive) |
|
||||
| ---------------- | -------------- | ------------- |
|
||||
| Rate limit check | ✅ Synchronous | ❌ Not used |
|
||||
| Quota check | ✅ Synchronous | ❌ Not used |
|
||||
| Request status | ✅ Synchronous | ❌ Not used |
|
||||
| Token update | ✅ Synchronous | ✅ Async |
|
||||
| Billing report | ❌ Not used | ✅ Query |
|
||||
| Audit log | ❌ Not used | ✅ Query |
|
||||
|
||||
### Key Points
|
||||
|
||||
1. **Two-layer storage**: KV for real-time, SQL for archive
|
||||
2. **KV operations are synchronous**: Rate limit and quota checks must be fast
|
||||
3. **SQL writes are async**: Archive happens in background goroutine
|
||||
4. **Services update tokens via `request_id`**: Updates both KV and SQL
|
||||
5. **KV data has TTL**: Auto-expires to prevent memory bloat
|
||||
6. **SQL data is permanent**: For billing and compliance
|
||||
Loading…
Add table
Reference in a new issue