Refactor chat storage design and enhance data models
- Updated the chat storage design to reflect a shift from "Conversation" to "Chat" terminology, improving clarity in the data structure. - Revised the data models for Chat and Message tables, including new fields for enhanced metadata management and permissions. - Introduced a `space_snapshot` field in the Step model to facilitate recovery during execution, improving resilience in chat interactions. - Enhanced middleware documentation to reflect the modular architecture and provide clearer usage examples for different API routes. - Updated tests to ensure the integrity of new data structures and functionalities, reinforcing the robustness of the chat storage system.
This commit is contained in:
Max
parent
467d4e2398
commit
42d13ec1cb
2 changed files with 716 additions and 321 deletions
|
|
@ -39,14 +39,14 @@ The chat storage system is designed to:
|
||||||
|
|
||||||
The Agent storage focuses on **chat content and execution state**, while request tracking (billing, rate limiting, auditing) is handled globally by the OpenAPI layer:
|
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 |
|
| Concern | Module | Table |
|
||||||
| ------------------ | ----------------- | -------------------- |
|
| ---------------- | ----------------- | ----------------- |
|
||||||
| Request tracking | `openapi/request` | `openapi_request` |
|
| Request tracking | `openapi/request` | `openapi_request` |
|
||||||
| Billing (tokens) | `openapi/request` | `openapi_request` |
|
| Billing (tokens) | `openapi/request` | `openapi_request` |
|
||||||
| Rate limiting | `openapi/request` | - |
|
| Rate limiting | `openapi/request` | - |
|
||||||
| Chat conversations | `agent/store` | `agent_conversation` |
|
| Chat sessions | `agent/store` | `agent_chat` |
|
||||||
| Chat messages | `agent/store` | `agent_message` |
|
| Chat messages | `agent/store` | `agent_message` |
|
||||||
| Execution steps | `agent/store` | `agent_step` |
|
| Execution steps | `agent/store` | `agent_step` |
|
||||||
|
|
||||||
The `request_id` from OpenAPI middleware is passed to Agent and stored in messages/steps for correlation.
|
The `request_id` from OpenAPI middleware is passed to Agent and stored in messages/steps for correlation.
|
||||||
|
|
||||||
|
|
@ -58,7 +58,7 @@ The `request_id` from OpenAPI middleware is passed to Agent and stored in messag
|
||||||
├─────────────────────────────────────────────────────────────┤
|
├─────────────────────────────────────────────────────────────┤
|
||||||
│ │
|
│ │
|
||||||
│ ┌─────────────────┐ │
|
│ ┌─────────────────┐ │
|
||||||
│ │ Conversation │ Metadata: title, assistant, user │
|
│ │ Chat │ Metadata: title, assistant, user │
|
||||||
│ └────────┬────────┘ │
|
│ └────────┬────────┘ │
|
||||||
│ │ │
|
│ │ │
|
||||||
│ │ 1:N │
|
│ │ 1:N │
|
||||||
|
|
@ -78,35 +78,60 @@ The `request_id` from OpenAPI middleware is passed to Agent and stored in messag
|
||||||
|
|
||||||
## Data Models
|
## Data Models
|
||||||
|
|
||||||
### 1. Conversation Table
|
### 1. Chat Table
|
||||||
|
|
||||||
Stores conversation metadata and session information.
|
Stores chat metadata and session information.
|
||||||
|
|
||||||
**Table Name:** `agent_conversation`
|
**Table Name:** `agent_chat`
|
||||||
|
|
||||||
| Column | Type | Nullable | Index | Description |
|
| Column | Type | Nullable | Index | Description |
|
||||||
| ----------------- | ----------- | -------- | ------ | ----------------------------------- |
|
| ----------------- | ----------- | -------- | ------ | -------------------------------- |
|
||||||
| `id` | ID | No | PK | Auto-increment primary key |
|
| `id` | ID | No | PK | Auto-increment primary key |
|
||||||
| `conversation_id` | string(64) | No | Unique | Unique conversation identifier |
|
| `chat_id` | string(64) | No | Unique | Unique chat identifier |
|
||||||
| `title` | string(500) | Yes | - | Conversation title |
|
| `title` | string(500) | Yes | - | Chat title |
|
||||||
| `assistant_id` | string(200) | No | Yes | Associated assistant ID |
|
| `assistant_id` | string(200) | No | Yes | Associated assistant ID |
|
||||||
| `user_id` | string(200) | No | Yes | Owner user ID |
|
| `mode` | string(50) | No | - | Chat mode (default: "chat") |
|
||||||
| `team_id` | string(200) | Yes | Yes | Team ID for access control |
|
| `status` | enum | No | Yes | Status: `active`, `archived` |
|
||||||
| `mode` | string(50) | No | - | Conversation mode (default: "chat") |
|
| `preset` | boolean | No | - | Whether this is a preset chat |
|
||||||
| `status` | enum | No | Yes | Status: `active`, `archived` |
|
| `public` | boolean | No | - | Whether shared across all teams |
|
||||||
| `last_message_at` | timestamp | Yes | Yes | Timestamp of last message |
|
| `share` | enum | No | Yes | Sharing scope: `private`, `team` |
|
||||||
| `metadata` | json | Yes | - | Additional metadata |
|
| `sort` | integer | No | - | Sort order for display |
|
||||||
| `created_at` | timestamp | No | Yes | Creation timestamp |
|
| `last_message_at` | timestamp | Yes | Yes | Timestamp of last message |
|
||||||
| `updated_at` | timestamp | No | - | Last update timestamp |
|
| `metadata` | json | Yes | - | Additional metadata |
|
||||||
|
| `created_at` | timestamp | No | Yes | Creation timestamp |
|
||||||
|
| `updated_at` | timestamp | No | - | Last update timestamp |
|
||||||
|
|
||||||
|
**Model Options:**
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"option": {
|
||||||
|
"soft_deletes": true,
|
||||||
|
"permission": true,
|
||||||
|
"timestamps": true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
**Note:** `permission: true` enables Yao's built-in permission management, which automatically adds the following fields:
|
||||||
|
|
||||||
|
| Field | Type | Description |
|
||||||
|
| ------------------ | ----------- | ------------------------------ |
|
||||||
|
| `__yao_created_by` | string(200) | User ID who created the record |
|
||||||
|
| `__yao_updated_by` | string(200) | User ID who last updated |
|
||||||
|
| `__yao_team_id` | string(200) | Team ID for team-level access |
|
||||||
|
| `__yao_tenant_id` | string(200) | Tenant ID for multi-tenancy |
|
||||||
|
|
||||||
|
These fields are automatically managed by the framework and used for access control filtering.
|
||||||
|
|
||||||
**Indexes:**
|
**Indexes:**
|
||||||
|
|
||||||
| Name | Columns | Type |
|
| Name | Columns | Type |
|
||||||
| -------------------- | ------------------- | ----- |
|
| -------------------- | ----------------- | ----- |
|
||||||
| `idx_conv_user` | `user_id`, `status` | index |
|
| `idx_chat_assistant` | `assistant_id` | index |
|
||||||
| `idx_conv_team` | `team_id`, `status` | index |
|
| `idx_chat_status` | `status` | index |
|
||||||
| `idx_conv_assistant` | `assistant_id` | index |
|
| `idx_chat_share` | `share` | index |
|
||||||
| `idx_conv_last_msg` | `last_message_at` | index |
|
| `idx_chat_last_msg` | `last_message_at` | index |
|
||||||
|
|
||||||
### 2. Message Table
|
### 2. Message Table
|
||||||
|
|
||||||
|
|
@ -114,31 +139,31 @@ Stores user-visible messages (both user input and assistant responses).
|
||||||
|
|
||||||
**Table Name:** `agent_message`
|
**Table Name:** `agent_message`
|
||||||
|
|
||||||
| Column | Type | Nullable | Index | Description |
|
| Column | Type | Nullable | Index | Description |
|
||||||
| ----------------- | ----------- | -------- | ------ | ----------------------------------------- |
|
| -------------- | ----------- | -------- | ------ | ----------------------------------------- |
|
||||||
| `id` | ID | No | PK | Auto-increment primary key |
|
| `id` | ID | No | PK | Auto-increment primary key |
|
||||||
| `message_id` | string(64) | No | Unique | Unique message identifier |
|
| `message_id` | string(64) | No | Unique | Unique message identifier |
|
||||||
| `conversation_id` | string(64) | No | Yes | Parent conversation ID |
|
| `chat_id` | string(64) | No | Yes | Parent chat ID |
|
||||||
| `request_id` | string(64) | Yes | Yes | Request ID for grouping |
|
| `request_id` | string(64) | Yes | Yes | Request ID for grouping |
|
||||||
| `role` | enum | No | Yes | Role: `user`, `assistant` |
|
| `role` | enum | No | Yes | Role: `user`, `assistant` |
|
||||||
| `type` | string(50) | No | - | Message type (text, image, loading, etc.) |
|
| `type` | string(50) | No | - | Message type (text, image, loading, etc.) |
|
||||||
| `props` | json | No | - | Message properties (content, url, etc.) |
|
| `props` | json | No | - | Message properties (content, url, etc.) |
|
||||||
| `block_id` | string(64) | Yes | Yes | Block grouping ID |
|
| `block_id` | string(64) | Yes | Yes | Block grouping ID |
|
||||||
| `thread_id` | string(64) | Yes | Yes | Thread grouping ID |
|
| `thread_id` | string(64) | Yes | Yes | Thread grouping ID |
|
||||||
| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) |
|
| `assistant_id` | string(200) | Yes | Yes | Assistant ID (join to get name/avatar) |
|
||||||
| `sequence` | integer | No | Yes | Message order within conversation |
|
| `sequence` | integer | No | Yes | Message order within chat |
|
||||||
| `metadata` | json | Yes | - | Additional metadata |
|
| `metadata` | json | Yes | - | Additional metadata |
|
||||||
| `created_at` | timestamp | No | Yes | Creation timestamp |
|
| `created_at` | timestamp | No | Yes | Creation timestamp |
|
||||||
| `updated_at` | timestamp | No | - | Last update timestamp |
|
| `updated_at` | timestamp | No | - | Last update timestamp |
|
||||||
|
|
||||||
**Indexes:**
|
**Indexes:**
|
||||||
|
|
||||||
| Name | Columns | Type |
|
| Name | Columns | Type |
|
||||||
| ------------------- | ----------------------------- | ----- |
|
| ------------------- | --------------------- | ----- |
|
||||||
| `idx_msg_conv_seq` | `conversation_id`, `sequence` | index |
|
| `idx_msg_chat_seq` | `chat_id`, `sequence` | index |
|
||||||
| `idx_msg_request` | `request_id` | index |
|
| `idx_msg_request` | `request_id` | index |
|
||||||
| `idx_msg_block` | `block_id` | index |
|
| `idx_msg_block` | `block_id` | index |
|
||||||
| `idx_msg_assistant` | `assistant_id` | index |
|
| `idx_msg_assistant` | `assistant_id` | index |
|
||||||
|
|
||||||
**Message Types:**
|
**Message Types:**
|
||||||
|
|
||||||
|
|
@ -161,7 +186,7 @@ Stores execution steps for resume/retry functionality.
|
||||||
| ----------------- | ----------- | -------- | ------ | -------------------------------- |
|
| ----------------- | ----------- | -------- | ------ | -------------------------------- |
|
||||||
| `id` | ID | No | PK | Auto-increment primary key |
|
| `id` | ID | No | PK | Auto-increment primary key |
|
||||||
| `step_id` | string(64) | No | Unique | Unique step identifier |
|
| `step_id` | string(64) | No | Unique | Unique step identifier |
|
||||||
| `conversation_id` | string(64) | No | Yes | Parent conversation ID |
|
| `chat_id` | string(64) | No | Yes | Parent chat ID |
|
||||||
| `request_id` | string(64) | No | Yes | Request ID |
|
| `request_id` | string(64) | No | Yes | Request ID |
|
||||||
| `assistant_id` | string(200) | No | Yes | Assistant executing this step |
|
| `assistant_id` | string(200) | No | Yes | Assistant executing this step |
|
||||||
| `stack_id` | string(64) | No | Yes | Stack node ID for this execution |
|
| `stack_id` | string(64) | No | Yes | Stack node ID for this execution |
|
||||||
|
|
@ -171,12 +196,34 @@ Stores execution steps for resume/retry functionality.
|
||||||
| `status` | enum | No | Yes | Step status |
|
| `status` | enum | No | Yes | Step status |
|
||||||
| `input` | json | Yes | - | Step input data |
|
| `input` | json | Yes | - | Step input data |
|
||||||
| `output` | json | Yes | - | Step output data |
|
| `output` | json | Yes | - | Step output data |
|
||||||
|
| `space_snapshot` | json | Yes | - | Space data snapshot for recovery |
|
||||||
| `error` | text | Yes | - | Error message if failed |
|
| `error` | text | Yes | - | Error message if failed |
|
||||||
| `sequence` | integer | No | Yes | Step order within request |
|
| `sequence` | integer | No | Yes | Step order within request |
|
||||||
| `metadata` | json | Yes | - | Additional metadata |
|
| `metadata` | json | Yes | - | Additional metadata |
|
||||||
| `created_at` | timestamp | No | Yes | Creation timestamp |
|
| `created_at` | timestamp | No | Yes | Creation timestamp |
|
||||||
| `updated_at` | timestamp | No | - | Last update timestamp |
|
| `updated_at` | timestamp | No | - | Last update timestamp |
|
||||||
|
|
||||||
|
**Space Snapshot:**
|
||||||
|
|
||||||
|
The `space_snapshot` field stores the shared data space (`ctx.Space`) at each step for recovery purposes.
|
||||||
|
|
||||||
|
```typescript
|
||||||
|
// Example: In Next hook, set data to Space before delegate
|
||||||
|
ctx.space.Set("choose_prompt", "query");
|
||||||
|
return {
|
||||||
|
delegate: { agent_id: "expense", messages: payload.messages },
|
||||||
|
};
|
||||||
|
```
|
||||||
|
|
||||||
|
If interrupted during delegate, the `space_snapshot` allows restoring `ctx.Space` state:
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"choose_prompt": "query",
|
||||||
|
"user_preferences": { "currency": "USD" }
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
**Step Types:**
|
**Step Types:**
|
||||||
|
|
||||||
| Type | Description | Input | Output |
|
| Type | Description | Input | Output |
|
||||||
|
|
@ -202,7 +249,7 @@ Stores execution steps for resume/retry functionality.
|
||||||
|
|
||||||
| Name | Columns | Type |
|
| Name | Columns | Type |
|
||||||
| -------------------- | ------------------------ | ----- |
|
| -------------------- | ------------------------ | ----- |
|
||||||
| `idx_step_conv` | `conversation_id` | index |
|
| `idx_step_chat` | `chat_id` | index |
|
||||||
| `idx_step_request` | `request_id`, `sequence` | index |
|
| `idx_step_request` | `request_id`, `sequence` | index |
|
||||||
| `idx_step_status` | `status` | index |
|
| `idx_step_status` | `status` | index |
|
||||||
| `idx_step_stack` | `stack_id` | index |
|
| `idx_step_stack` | `stack_id` | index |
|
||||||
|
|
@ -350,19 +397,26 @@ func (ast *Assistant) Stream(ctx, inputMessages, options) {
|
||||||
|
|
||||||
// createStep creates a step with context information
|
// createStep creates a step with context information
|
||||||
func createStep(ctx *Context, stepType, status string, input, output interface{}) *Step {
|
func createStep(ctx *Context, stepType, status string, input, output interface{}) *Step {
|
||||||
|
// Capture Space snapshot for recovery
|
||||||
|
var spaceSnapshot map[string]interface{}
|
||||||
|
if ctx.Space != nil {
|
||||||
|
spaceSnapshot = ctx.Space.Snapshot() // Get all key-value pairs
|
||||||
|
}
|
||||||
|
|
||||||
return &Step{
|
return &Step{
|
||||||
StepID: generateID(),
|
StepID: generateID(),
|
||||||
ConversationID: ctx.ChatID, // ChatID = conversation_id
|
ChatID: ctx.ChatID, // ChatID
|
||||||
RequestID: ctx.RequestID, // From OpenAPI middleware
|
RequestID: ctx.RequestID, // From OpenAPI middleware
|
||||||
AssistantID: ctx.AssistantID,
|
AssistantID: ctx.AssistantID,
|
||||||
StackID: ctx.Stack.ID,
|
StackID: ctx.Stack.ID,
|
||||||
StackParentID: ctx.Stack.ParentID,
|
StackParentID: ctx.Stack.ParentID,
|
||||||
StackDepth: ctx.Stack.Depth,
|
StackDepth: ctx.Stack.Depth,
|
||||||
Type: stepType,
|
Type: stepType,
|
||||||
Status: status,
|
Status: status,
|
||||||
Input: input,
|
Input: input,
|
||||||
Output: output,
|
Output: output,
|
||||||
Sequence: nextSequence(),
|
SpaceSnapshot: spaceSnapshot, // Shared space data for recovery
|
||||||
|
Sequence: nextSequence(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
@ -373,88 +427,103 @@ func createStep(ctx *Context, stepType, status string, input, output interface{}
|
||||||
```go
|
```go
|
||||||
// ChatStore defines the chat storage interface
|
// ChatStore defines the chat storage interface
|
||||||
type ChatStore interface {
|
type ChatStore interface {
|
||||||
// Conversation Management
|
// Chat Management
|
||||||
CreateConversation(conv *Conversation) error
|
CreateChat(chat *Chat) error
|
||||||
GetConversation(conversationID string) (*Conversation, error)
|
GetChat(chatID string) (*Chat, error)
|
||||||
UpdateConversation(conversationID string, updates map[string]interface{}) error
|
UpdateChat(chatID string, updates map[string]interface{}) error
|
||||||
DeleteConversation(conversationID string) error
|
DeleteChat(chatID string) error
|
||||||
ListConversations(filter ConversationFilter) (*ConversationList, error)
|
ListChats(filter ChatFilter) (*ChatList, error)
|
||||||
|
|
||||||
// Message Management
|
// Message Management
|
||||||
SaveMessages(conversationID string, messages []*Message) error
|
SaveMessages(chatID string, messages []*Message) error
|
||||||
GetMessages(conversationID string, filter MessageFilter) ([]*Message, error)
|
GetMessages(chatID string, filter MessageFilter) ([]*Message, error)
|
||||||
UpdateMessage(messageID string, updates map[string]interface{}) error
|
UpdateMessage(messageID string, updates map[string]interface{}) error
|
||||||
DeleteMessages(conversationID string, messageIDs []string) error
|
DeleteMessages(chatID string, messageIDs []string) error
|
||||||
|
|
||||||
// Step Management
|
// Step Management
|
||||||
SaveStep(step *Step) error
|
SaveSteps(steps []*Step) error
|
||||||
UpdateStep(stepID string, updates map[string]interface{}) error
|
UpdateStep(stepID string, updates map[string]interface{}) error
|
||||||
GetSteps(requestID string) ([]*Step, error)
|
GetSteps(requestID string) ([]*Step, error)
|
||||||
GetLastIncompleteStep(conversationID string) (*Step, error)
|
GetLastIncompleteStep(chatID string) (*Step, error)
|
||||||
|
GetStepsByStackID(stackID string) ([]*Step, error)
|
||||||
|
GetStackPath(stackID string) ([]string, error) // Returns [root_stack_id, ..., current_stack_id]
|
||||||
|
}
|
||||||
|
|
||||||
|
// SpaceStore defines the interface for Space snapshot operations
|
||||||
|
// Note: Space itself uses plan.Space interface, this is for persistence
|
||||||
|
type SpaceStore interface {
|
||||||
|
// Snapshot returns all key-value pairs in the space
|
||||||
|
Snapshot() map[string]interface{}
|
||||||
|
|
||||||
|
// Restore sets multiple key-value pairs from a snapshot
|
||||||
|
Restore(data map[string]interface{}) error
|
||||||
}
|
}
|
||||||
````
|
````
|
||||||
|
|
||||||
### Data Structures
|
### Data Structures
|
||||||
|
|
||||||
```go
|
```go
|
||||||
// Conversation represents a chat conversation
|
// Chat represents a chat session
|
||||||
type Conversation struct {
|
type Chat struct {
|
||||||
ConversationID string `json:"conversation_id"`
|
ChatID string `json:"chat_id"`
|
||||||
Title string `json:"title,omitempty"`
|
Title string `json:"title,omitempty"`
|
||||||
AssistantID string `json:"assistant_id"`
|
AssistantID string `json:"assistant_id"`
|
||||||
UserID string `json:"user_id"`
|
Mode string `json:"mode"`
|
||||||
TeamID string `json:"team_id,omitempty"`
|
Status string `json:"status"`
|
||||||
Mode string `json:"mode"`
|
Preset bool `json:"preset"`
|
||||||
Status string `json:"status"`
|
Public bool `json:"public"`
|
||||||
LastMessageAt *time.Time `json:"last_message_at,omitempty"`
|
Share string `json:"share"` // "private" or "team"
|
||||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
Sort int `json:"sort"`
|
||||||
CreatedAt time.Time `json:"created_at"`
|
LastMessageAt *time.Time `json:"last_message_at,omitempty"`
|
||||||
UpdatedAt time.Time `json:"updated_at"`
|
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||||
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// Message represents a chat message
|
// Message represents a chat message
|
||||||
type Message struct {
|
type Message struct {
|
||||||
MessageID string `json:"message_id"`
|
MessageID string `json:"message_id"`
|
||||||
ConversationID string `json:"conversation_id"`
|
ChatID string `json:"chat_id"`
|
||||||
RequestID string `json:"request_id,omitempty"`
|
RequestID string `json:"request_id,omitempty"`
|
||||||
Role string `json:"role"`
|
Role string `json:"role"`
|
||||||
Type string `json:"type"`
|
Type string `json:"type"`
|
||||||
Props map[string]interface{} `json:"props"`
|
Props map[string]interface{} `json:"props"`
|
||||||
BlockID string `json:"block_id,omitempty"`
|
BlockID string `json:"block_id,omitempty"`
|
||||||
ThreadID string `json:"thread_id,omitempty"`
|
ThreadID string `json:"thread_id,omitempty"`
|
||||||
AssistantID string `json:"assistant_id,omitempty"`
|
AssistantID string `json:"assistant_id,omitempty"`
|
||||||
Sequence int `json:"sequence"`
|
Sequence int `json:"sequence"`
|
||||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||||
CreatedAt time.Time `json:"created_at"`
|
CreatedAt time.Time `json:"created_at"`
|
||||||
UpdatedAt time.Time `json:"updated_at"`
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// Step represents an execution step
|
// Step represents an execution step
|
||||||
type Step struct {
|
type Step struct {
|
||||||
StepID string `json:"step_id"`
|
StepID string `json:"step_id"`
|
||||||
ConversationID string `json:"conversation_id"`
|
ChatID string `json:"chat_id"`
|
||||||
RequestID string `json:"request_id"`
|
RequestID string `json:"request_id"`
|
||||||
AssistantID string `json:"assistant_id"`
|
AssistantID string `json:"assistant_id"`
|
||||||
StackID string `json:"stack_id"`
|
StackID string `json:"stack_id"`
|
||||||
StackParentID string `json:"stack_parent_id,omitempty"`
|
StackParentID string `json:"stack_parent_id,omitempty"`
|
||||||
StackDepth int `json:"stack_depth"`
|
StackDepth int `json:"stack_depth"`
|
||||||
Type string `json:"type"`
|
Type string `json:"type"`
|
||||||
Status string `json:"status"`
|
Status string `json:"status"`
|
||||||
Input map[string]interface{} `json:"input,omitempty"`
|
Input map[string]interface{} `json:"input,omitempty"`
|
||||||
Output map[string]interface{} `json:"output,omitempty"`
|
Output map[string]interface{} `json:"output,omitempty"`
|
||||||
Error string `json:"error,omitempty"`
|
SpaceSnapshot map[string]interface{} `json:"space_snapshot,omitempty"` // Shared space data for recovery
|
||||||
Sequence int `json:"sequence"`
|
Error string `json:"error,omitempty"`
|
||||||
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
Sequence int `json:"sequence"`
|
||||||
CreatedAt time.Time `json:"created_at"`
|
Metadata map[string]interface{} `json:"metadata,omitempty"`
|
||||||
UpdatedAt time.Time `json:"updated_at"`
|
CreatedAt time.Time `json:"created_at"`
|
||||||
|
UpdatedAt time.Time `json:"updated_at"`
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
### Filter Structures
|
### Filter Structures
|
||||||
|
|
||||||
```go
|
```go
|
||||||
// ConversationFilter for listing conversations
|
// ChatFilter for listing chats
|
||||||
type ConversationFilter struct {
|
type ChatFilter struct {
|
||||||
UserID string `json:"user_id,omitempty"`
|
UserID string `json:"user_id,omitempty"`
|
||||||
TeamID string `json:"team_id,omitempty"`
|
TeamID string `json:"team_id,omitempty"`
|
||||||
AssistantID string `json:"assistant_id,omitempty"`
|
AssistantID string `json:"assistant_id,omitempty"`
|
||||||
|
|
@ -473,13 +542,13 @@ type MessageFilter struct {
|
||||||
Offset int `json:"offset,omitempty"`
|
Offset int `json:"offset,omitempty"`
|
||||||
}
|
}
|
||||||
|
|
||||||
// ConversationList paginated response
|
// ChatList paginated response
|
||||||
type ConversationList struct {
|
type ChatList struct {
|
||||||
Data []*Conversation `json:"data"`
|
Data []*Chat `json:"data"`
|
||||||
Page int `json:"page"`
|
Page int `json:"page"`
|
||||||
PageSize int `json:"pagesize"`
|
PageSize int `json:"pagesize"`
|
||||||
PageCount int `json:"pagecount"`
|
PageCount int `json:"pagecount"`
|
||||||
Total int `json:"total"`
|
Total int `json:"total"`
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|
@ -492,23 +561,23 @@ See [Write Strategy - Implementation](#implementation) for the complete flow wit
|
||||||
### 2. Load Chat History
|
### 2. Load Chat History
|
||||||
|
|
||||||
```go
|
```go
|
||||||
// Get conversation list
|
// Get chat list
|
||||||
convs, _ := chatStore.ListConversations(ConversationFilter{
|
chats, _ := chatStore.ListChats(ChatFilter{
|
||||||
UserID: "user123",
|
UserID: "user123",
|
||||||
Status: "active",
|
Status: "active",
|
||||||
Page: 1,
|
Page: 1,
|
||||||
PageSize: 20,
|
PageSize: 20,
|
||||||
})
|
})
|
||||||
|
|
||||||
// Get messages for a conversation
|
// Get messages for a chat
|
||||||
messages, _ := chatStore.GetMessages("conv_123", MessageFilter{
|
messages, _ := chatStore.GetMessages("chat_123", MessageFilter{
|
||||||
Limit: 100,
|
Limit: 100,
|
||||||
})
|
})
|
||||||
|
|
||||||
// Return to frontend
|
// Return to frontend
|
||||||
return map[string]interface{}{
|
return map[string]interface{}{
|
||||||
"conversation": conv,
|
"chat": chat,
|
||||||
"messages": messages,
|
"messages": messages,
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|
@ -517,18 +586,25 @@ return map[string]interface{}{
|
||||||
```go
|
```go
|
||||||
func (ast *Assistant) Resume(ctx *Context) error {
|
func (ast *Assistant) Resume(ctx *Context) error {
|
||||||
// 1. Find last incomplete step
|
// 1. Find last incomplete step
|
||||||
step, _ := chatStore.GetLastIncompleteStep(ctx.ConversationID)
|
step, _ := chatStore.GetLastIncompleteStep(ctx.ChatID)
|
||||||
if step == nil {
|
if step == nil {
|
||||||
return nil // Nothing to resume
|
return nil // Nothing to resume
|
||||||
}
|
}
|
||||||
|
|
||||||
// 2. Check if this is an A2A nested call
|
// 2. Restore Space data from snapshot
|
||||||
|
if step.SpaceSnapshot != nil && ctx.Space != nil {
|
||||||
|
for key, value := range step.SpaceSnapshot {
|
||||||
|
ctx.Space.Set(key, value)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Check if this is an A2A nested call
|
||||||
if step.StackDepth > 0 {
|
if step.StackDepth > 0 {
|
||||||
// Need to rebuild the call stack
|
// Need to rebuild the call stack
|
||||||
return ast.ResumeNestedCall(ctx, step)
|
return ast.ResumeNestedCall(ctx, step)
|
||||||
}
|
}
|
||||||
|
|
||||||
// 3. Resume based on step type
|
// 4. Resume based on step type
|
||||||
switch step.Type {
|
switch step.Type {
|
||||||
case "llm":
|
case "llm":
|
||||||
// Re-execute LLM call with saved input
|
// Re-execute LLM call with saved input
|
||||||
|
|
@ -542,6 +618,12 @@ func (ast *Assistant) Resume(ctx *Context) error {
|
||||||
case "hook_next":
|
case "hook_next":
|
||||||
// Re-execute hook
|
// Re-execute hook
|
||||||
return ast.executeHookNext(ctx, step.Input)
|
return ast.executeHookNext(ctx, step.Input)
|
||||||
|
|
||||||
|
case "delegate":
|
||||||
|
// Resume delegated agent call
|
||||||
|
agentID := step.Input["agent_id"].(string)
|
||||||
|
messages := step.Input["messages"].([]Message)
|
||||||
|
return ast.delegateToAgent(ctx, agentID, messages)
|
||||||
}
|
}
|
||||||
|
|
||||||
return nil
|
return nil
|
||||||
|
|
@ -590,22 +672,37 @@ When Assistant A delegates to Assistant B, the step records look like:
|
||||||
Request: User asks "analyze this data and visualize it"
|
Request: User asks "analyze this data and visualize it"
|
||||||
|
|
||||||
Step Records:
|
Step Records:
|
||||||
┌─────┬─────────────┬─────────────┬──────────┬────────────┬───────┬─────────────┐
|
┌─────┬─────────────┬─────────────┬──────────┬───────┬───────┬─────────────┬─────────────────────────────┐
|
||||||
│ seq │ assistant │ stack_id │ parent │ depth │ type │ status │
|
│ seq │ assistant │ stack_id │ parent │ depth │ type │ status │ space_snapshot │
|
||||||
├─────┼─────────────┼─────────────┼──────────┼────────────┼───────┼─────────────┤
|
├─────┼─────────────┼─────────────┼──────────┼───────┼───────┼─────────────┼─────────────────────────────┤
|
||||||
│ 1 │ analyzer │ stk_001 │ null │ 0 │ input │ completed │
|
│ 1 │ analyzer │ stk_001 │ null │ 0 │ input │ completed │ {} │
|
||||||
│ 2 │ analyzer │ stk_001 │ null │ 0 │ llm │ completed │
|
│ 2 │ analyzer │ stk_001 │ null │ 0 │ llm │ completed │ {} │
|
||||||
│ 3 │ analyzer │ stk_001 │ null │ 0 │ delegate │ running │ ← delegating
|
│ 3 │ analyzer │ stk_001 │ null │ 0 │ delegate │ running │ {"choose_prompt": "query"} │ ← Space data set before delegate
|
||||||
│ 4 │ visualizer │ stk_002 │ stk_001 │ 1 │ input │ completed │
|
│ 4 │ visualizer │ stk_002 │ stk_001 │ 1 │ input │ completed │ {"choose_prompt": "query"} │
|
||||||
│ 5 │ visualizer │ stk_002 │ stk_001 │ 1 │ llm │ interrupted │ ← interrupted here
|
│ 5 │ visualizer │ stk_002 │ stk_001 │ 1 │ llm │ interrupted │ {"choose_prompt": "query"} │ ← interrupted here
|
||||||
└─────┴─────────────┴─────────────┴──────────┴────────────┴───────┴─────────────┘
|
└─────┴─────────────┴─────────────┴──────────┴───────┴───────┴─────────────┴─────────────────────────────┘
|
||||||
|
|
||||||
Resume Flow:
|
Resume Flow:
|
||||||
1. Find step with status="interrupted" → step 5
|
1. Find step with status="interrupted" → step 5
|
||||||
2. Check stack_depth=1 → nested call
|
2. Restore Space from space_snapshot: {"choose_prompt": "query"}
|
||||||
3. Get stack path: [stk_001, stk_002]
|
3. Check stack_depth=1 → nested call
|
||||||
4. Resume visualizer assistant with step 5's input
|
4. Get stack path: [stk_001, stk_002]
|
||||||
5. When visualizer completes, update step 3 (delegate) to completed
|
5. Resume visualizer assistant with step 5's input
|
||||||
|
6. When visualizer completes, update step 3 (delegate) to completed
|
||||||
|
```
|
||||||
|
|
||||||
|
**Space Snapshot Use Case (from expense assistant):**
|
||||||
|
|
||||||
|
```typescript
|
||||||
|
// In Next hook, before delegating to another agent
|
||||||
|
ctx.space.Set("choose_prompt", "query");
|
||||||
|
return {
|
||||||
|
delegate: { agent_id: "expense", messages: payload.messages },
|
||||||
|
};
|
||||||
|
|
||||||
|
// If interrupted during delegate, Resume will:
|
||||||
|
// 1. Restore space_snapshot → ctx.space now has "choose_prompt": "query"
|
||||||
|
// 2. The delegated agent's Create hook can read: ctx.space.GetDel("choose_prompt")
|
||||||
```
|
```
|
||||||
|
|
||||||
## Migration Notes
|
## Migration Notes
|
||||||
|
|
|
||||||
|
|
@ -284,135 +284,128 @@ type QuotaKey struct {
|
||||||
|
|
||||||
## Middleware Design
|
## Middleware Design
|
||||||
|
|
||||||
### Request Flow
|
### Modular Middleware Architecture
|
||||||
|
|
||||||
|
Each middleware is independent and can be composed based on business needs.
|
||||||
|
|
||||||
```
|
```
|
||||||
Request arrives
|
┌─────────────────────────────────────────────────────────────┐
|
||||||
│
|
│ Available Middlewares │
|
||||||
├── 1. Generate request_id (uuid or nanoid)
|
├─────────────────────────────────────────────────────────────┤
|
||||||
│
|
│ │
|
||||||
├── 2. Set request_id in context and response header
|
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │
|
||||||
│ c.Set("request_id", requestID)
|
│ │ RequestID │ │ RateLimit │ │ Quota │ │
|
||||||
│ c.Header("X-Request-ID", requestID)
|
│ │ (Basic) │ │ (Protect) │ │ (Billing) │ │
|
||||||
│
|
│ └─────────────┘ └─────────────┘ └─────────────┘ │
|
||||||
├── 3. Get auth info from context (set by OAuth Guard)
|
│ │
|
||||||
│ authInfo := authorized.GetInfo(c)
|
│ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │
|
||||||
│
|
│ │ Metrics │ │ Archive │ │ Billing │ │
|
||||||
├── 4. Detect service type from endpoint
|
│ │ (Monitor) │ │ (Audit) │ │ (Charge) │ │
|
||||||
│ 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
|
### Middleware List
|
||||||
|
|
||||||
|
| Middleware | File | Purpose | Dependencies |
|
||||||
|
| ----------- | --------------- | -------------------------------- | ------------------ |
|
||||||
|
| `RequestID` | `request_id.go` | Generate and track request ID | None |
|
||||||
|
| `RateLimit` | `ratelimit.go` | Request frequency limiting | KV, RequestID |
|
||||||
|
| `Quota` | `quota.go` | Token quota enforcement | KV, RequestID |
|
||||||
|
| `Metrics` | `metrics.go` | Request duration, status metrics | RequestID |
|
||||||
|
| `Archive` | `archive.go` | Persist request to SQL | SQL, RequestID |
|
||||||
|
| `Billing` | `billing.go` | Token usage tracking & charging | KV, SQL, RequestID |
|
||||||
|
|
||||||
|
### Usage Examples
|
||||||
|
|
||||||
|
#### Example 1: Full Protection (Agent API)
|
||||||
|
|
||||||
```go
|
```go
|
||||||
|
// Agent API needs all protections
|
||||||
|
agent := api.Group("/chat")
|
||||||
|
agent.Use(
|
||||||
|
request.RequestID(), // Generate request_id
|
||||||
|
request.RateLimit(kv, config), // Rate limiting
|
||||||
|
request.Quota(kv, config), // Token quota
|
||||||
|
request.Metrics(), // Duration tracking
|
||||||
|
request.Archive(sql), // Audit logging
|
||||||
|
request.Billing(kv, sql), // Token billing
|
||||||
|
)
|
||||||
|
agent.POST("/completions", handler.ChatCompletions)
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Example 2: Light Protection (File API)
|
||||||
|
|
||||||
|
```go
|
||||||
|
// File API only needs basic tracking
|
||||||
|
file := api.Group("/file")
|
||||||
|
file.Use(
|
||||||
|
request.RequestID(), // Generate request_id
|
||||||
|
request.RateLimit(kv, config), // Rate limiting
|
||||||
|
request.Metrics(), // Duration tracking
|
||||||
|
)
|
||||||
|
file.POST("/upload", handler.Upload)
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Example 3: Internal API (No Billing)
|
||||||
|
|
||||||
|
```go
|
||||||
|
// Internal API skips billing
|
||||||
|
internal := api.Group("/internal")
|
||||||
|
internal.Use(
|
||||||
|
request.RequestID(), // Generate request_id
|
||||||
|
request.Metrics(), // Duration tracking
|
||||||
|
request.Archive(sql), // Audit logging only
|
||||||
|
)
|
||||||
|
internal.GET("/health", handler.Health)
|
||||||
|
```
|
||||||
|
|
||||||
|
#### Example 4: Public API (Rate Limit Only)
|
||||||
|
|
||||||
|
```go
|
||||||
|
// Public endpoints only need rate limiting
|
||||||
|
public := api.Group("/public")
|
||||||
|
public.Use(
|
||||||
|
request.RequestID(), // Generate request_id
|
||||||
|
request.RateLimit(kv, config), // Rate limiting by IP
|
||||||
|
)
|
||||||
|
public.GET("/models", handler.ListModels)
|
||||||
|
```
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
### Middleware Implementations
|
||||||
|
|
||||||
|
#### 1. RequestID Middleware (Base)
|
||||||
|
|
||||||
|
```go
|
||||||
|
// request_id.go
|
||||||
package request
|
package request
|
||||||
|
|
||||||
import (
|
// RequestID generates and sets request ID
|
||||||
"time"
|
func RequestID() gin.HandlerFunc {
|
||||||
"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) {
|
return func(c *gin.Context) {
|
||||||
startTime := time.Now()
|
|
||||||
|
|
||||||
// 1. Generate request ID
|
|
||||||
requestID := generateRequestID()
|
requestID := generateRequestID()
|
||||||
c.Set("request_id", requestID)
|
c.Set("request_id", requestID)
|
||||||
c.Header("X-Request-ID", requestID)
|
c.Header("X-Request-ID", requestID)
|
||||||
|
|
||||||
// 2. Get auth info
|
// Also set start time for metrics
|
||||||
authInfo := authorized.GetInfo(c)
|
c.Set("request_start_time", time.Now())
|
||||||
|
|
||||||
// 3. Detect service and resource
|
// Detect and set service info
|
||||||
service := detectService(c.FullPath())
|
service := detectService(c.FullPath())
|
||||||
resourceID := extractResourceID(c, service)
|
c.Set("request_service", service)
|
||||||
|
c.Set("request_resource_id", 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()
|
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 generateRequestID() string {
|
||||||
|
return fmt.Sprintf("req_%s", nanoid.New())
|
||||||
|
}
|
||||||
|
|
||||||
func detectService(endpoint string) string {
|
func detectService(endpoint string) string {
|
||||||
switch {
|
switch {
|
||||||
case strings.HasPrefix(endpoint, "/api/chat"):
|
case strings.HasPrefix(endpoint, "/api/chat"):
|
||||||
|
|
@ -437,6 +430,256 @@ func detectService(endpoint string) string {
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
|
#### 2. RateLimit Middleware
|
||||||
|
|
||||||
|
```go
|
||||||
|
// ratelimit.go
|
||||||
|
package request
|
||||||
|
|
||||||
|
// RateLimit enforces request frequency limits
|
||||||
|
func RateLimit(kv KVStore, config *RateLimitConfig) gin.HandlerFunc {
|
||||||
|
return func(c *gin.Context) {
|
||||||
|
if config == nil || !config.Enabled {
|
||||||
|
c.Next()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
authInfo := authorized.GetInfo(c)
|
||||||
|
service := c.GetString("request_service")
|
||||||
|
|
||||||
|
// Check user rate limit
|
||||||
|
userKey := fmt.Sprintf("ratelimit:user:%s:%s", authInfo.UserID, service)
|
||||||
|
userCount, _ := kv.Incr(userKey, 60*time.Second)
|
||||||
|
if userCount > int64(config.GetUserLimit(service)) {
|
||||||
|
c.AbortWithStatusJSON(429, gin.H{
|
||||||
|
"error": "rate_limit_exceeded",
|
||||||
|
"message": fmt.Sprintf("User rate limit exceeded: %d requests per minute", config.GetUserLimit(service)),
|
||||||
|
"retry_after": 60,
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check team rate 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)) {
|
||||||
|
c.AbortWithStatusJSON(429, gin.H{
|
||||||
|
"error": "rate_limit_exceeded",
|
||||||
|
"message": "Team rate limit exceeded",
|
||||||
|
"retry_after": 60,
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check IP rate limit
|
||||||
|
ipKey := fmt.Sprintf("ratelimit:ip:%s", c.ClientIP())
|
||||||
|
ipCount, _ := kv.Incr(ipKey, 60*time.Second)
|
||||||
|
if ipCount > int64(config.GetIPLimit()) {
|
||||||
|
c.AbortWithStatusJSON(429, gin.H{
|
||||||
|
"error": "rate_limit_exceeded",
|
||||||
|
"message": "IP rate limit exceeded",
|
||||||
|
"retry_after": 60,
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
c.Next()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
#### 3. Quota Middleware
|
||||||
|
|
||||||
|
```go
|
||||||
|
// quota.go
|
||||||
|
package request
|
||||||
|
|
||||||
|
// Quota enforces token quota limits
|
||||||
|
func Quota(kv KVStore, config *QuotaConfig) gin.HandlerFunc {
|
||||||
|
return func(c *gin.Context) {
|
||||||
|
if config == nil || !config.Enabled {
|
||||||
|
c.Next()
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
authInfo := authorized.GetInfo(c)
|
||||||
|
|
||||||
|
// 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 {
|
||||||
|
c.AbortWithStatusJSON(429, gin.H{
|
||||||
|
"error": "quota_exceeded",
|
||||||
|
"message": "Daily token quota exceeded",
|
||||||
|
"reset_at": getNextDayStart(),
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check team monthly quota
|
||||||
|
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 {
|
||||||
|
c.AbortWithStatusJSON(429, gin.H{
|
||||||
|
"error": "quota_exceeded",
|
||||||
|
"message": "Team monthly token quota exceeded",
|
||||||
|
"reset_at": getNextMonthStart(),
|
||||||
|
})
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
c.Next()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
#### 4. Metrics Middleware
|
||||||
|
|
||||||
|
```go
|
||||||
|
// metrics.go
|
||||||
|
package request
|
||||||
|
|
||||||
|
// Metrics tracks request duration and status
|
||||||
|
func Metrics() gin.HandlerFunc {
|
||||||
|
return func(c *gin.Context) {
|
||||||
|
startTime := c.GetTime("request_start_time")
|
||||||
|
if startTime.IsZero() {
|
||||||
|
startTime = time.Now()
|
||||||
|
}
|
||||||
|
|
||||||
|
c.Next()
|
||||||
|
|
||||||
|
// Calculate duration
|
||||||
|
duration := time.Since(startTime)
|
||||||
|
c.Set("request_duration_ms", duration.Milliseconds())
|
||||||
|
|
||||||
|
// Determine status
|
||||||
|
status := "completed"
|
||||||
|
if c.Writer.Status() >= 400 {
|
||||||
|
status = "failed"
|
||||||
|
}
|
||||||
|
c.Set("request_status", status)
|
||||||
|
|
||||||
|
// TODO: Export to Prometheus/metrics system
|
||||||
|
// metrics.RequestDuration.WithLabelValues(service, status).Observe(duration.Seconds())
|
||||||
|
// metrics.RequestTotal.WithLabelValues(service, status).Inc()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
#### 5. Archive Middleware
|
||||||
|
|
||||||
|
```go
|
||||||
|
// archive.go
|
||||||
|
package request
|
||||||
|
|
||||||
|
// Archive persists request to SQL for audit
|
||||||
|
func Archive(sql SQLStore) gin.HandlerFunc {
|
||||||
|
return func(c *gin.Context) {
|
||||||
|
c.Next()
|
||||||
|
|
||||||
|
// Get request info from context
|
||||||
|
requestID := c.GetString("request_id")
|
||||||
|
if requestID == "" {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
authInfo := authorized.GetInfo(c)
|
||||||
|
startTime := c.GetTime("request_start_time")
|
||||||
|
durationMs := c.GetInt64("request_duration_ms")
|
||||||
|
status := c.GetString("request_status")
|
||||||
|
if status == "" {
|
||||||
|
status = "completed"
|
||||||
|
}
|
||||||
|
|
||||||
|
completedAt := time.Now()
|
||||||
|
|
||||||
|
// 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: c.GetString("request_service"),
|
||||||
|
ResourceID: c.GetString("request_resource_id"),
|
||||||
|
Status: status,
|
||||||
|
StatusCode: c.Writer.Status(),
|
||||||
|
Referer: c.GetHeader("X-Yao-Referer"),
|
||||||
|
ClientType: getClientType(c.GetHeader("User-Agent")),
|
||||||
|
ClientIP: c.ClientIP(),
|
||||||
|
DurationMs: durationMs,
|
||||||
|
Error: c.GetString("request_error"),
|
||||||
|
CreatedAt: startTime,
|
||||||
|
CompletedAt: &completedAt,
|
||||||
|
})
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
#### 6. Billing Middleware
|
||||||
|
|
||||||
|
```go
|
||||||
|
// billing.go
|
||||||
|
package request
|
||||||
|
|
||||||
|
// Billing tracks token usage (called by services after completion)
|
||||||
|
func Billing(kv KVStore, sql SQLStore) gin.HandlerFunc {
|
||||||
|
return func(c *gin.Context) {
|
||||||
|
c.Next()
|
||||||
|
|
||||||
|
// Token usage is updated by services via UpdateTokenUsage()
|
||||||
|
// This middleware just ensures the billing context is available
|
||||||
|
c.Set("billing_kv", kv)
|
||||||
|
c.Set("billing_sql", sql)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// UpdateTokenUsage is called by services after completion
|
||||||
|
func UpdateTokenUsage(c *gin.Context, input, output int) error {
|
||||||
|
kv, ok := c.Get("billing_kv")
|
||||||
|
if !ok {
|
||||||
|
return nil // Billing not enabled
|
||||||
|
}
|
||||||
|
|
||||||
|
sql, _ := c.Get("billing_sql")
|
||||||
|
requestID := c.GetString("request_id")
|
||||||
|
authInfo := authorized.GetInfo(c)
|
||||||
|
|
||||||
|
return updateTokenUsageInternal(
|
||||||
|
kv.(KVStore),
|
||||||
|
sql.(SQLStore),
|
||||||
|
requestID,
|
||||||
|
authInfo.UserID,
|
||||||
|
authInfo.TeamID,
|
||||||
|
input,
|
||||||
|
output,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
## Rate Limiting
|
## Rate Limiting
|
||||||
|
|
||||||
### Configuration
|
### Configuration
|
||||||
|
|
@ -769,7 +1012,76 @@ type DailyUsage struct {
|
||||||
|
|
||||||
## Integration with Services
|
## Integration with Services
|
||||||
|
|
||||||
### Agent Service
|
### Route Registration Example
|
||||||
|
|
||||||
|
```go
|
||||||
|
// openapi/openapi.go
|
||||||
|
func (s *OpenAPI) RegisterRoutes(r *gin.Engine) {
|
||||||
|
api := r.Group("/api")
|
||||||
|
|
||||||
|
// 1. OAuth Guard (authentication) - for all routes
|
||||||
|
api.Use(oauth.Guard)
|
||||||
|
|
||||||
|
// 2. Register different route groups with different middleware combinations
|
||||||
|
s.registerAgentRoutes(api)
|
||||||
|
s.registerKBRoutes(api)
|
||||||
|
s.registerLLMRoutes(api)
|
||||||
|
s.registerFileRoutes(api)
|
||||||
|
s.registerPublicRoutes(api)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *OpenAPI) registerAgentRoutes(api *gin.RouterGroup) {
|
||||||
|
// Agent API: Full protection + billing
|
||||||
|
agent := api.Group("/chat")
|
||||||
|
agent.Use(
|
||||||
|
request.RequestID(),
|
||||||
|
request.RateLimit(s.kv, s.rateLimitConfig),
|
||||||
|
request.Quota(s.kv, s.quotaConfig),
|
||||||
|
request.Metrics(),
|
||||||
|
request.Archive(s.sql),
|
||||||
|
request.Billing(s.kv, s.sql),
|
||||||
|
)
|
||||||
|
agent.POST("/completions", s.handler.ChatCompletions)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *OpenAPI) registerKBRoutes(api *gin.RouterGroup) {
|
||||||
|
// KB API: Rate limit + archive (no token billing)
|
||||||
|
kb := api.Group("/kb")
|
||||||
|
kb.Use(
|
||||||
|
request.RequestID(),
|
||||||
|
request.RateLimit(s.kv, s.rateLimitConfig),
|
||||||
|
request.Metrics(),
|
||||||
|
request.Archive(s.sql),
|
||||||
|
)
|
||||||
|
kb.POST("/search", s.handler.KBSearch)
|
||||||
|
kb.POST("/upload", s.handler.KBUpload)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *OpenAPI) registerFileRoutes(api *gin.RouterGroup) {
|
||||||
|
// File API: Light protection
|
||||||
|
file := api.Group("/file")
|
||||||
|
file.Use(
|
||||||
|
request.RequestID(),
|
||||||
|
request.RateLimit(s.kv, s.rateLimitConfig),
|
||||||
|
request.Metrics(),
|
||||||
|
)
|
||||||
|
file.POST("/upload", s.handler.FileUpload)
|
||||||
|
file.GET("/download/:id", s.handler.FileDownload)
|
||||||
|
}
|
||||||
|
|
||||||
|
func (s *OpenAPI) registerPublicRoutes(api *gin.RouterGroup) {
|
||||||
|
// Public API: Rate limit only (no auth required)
|
||||||
|
public := api.Group("/public")
|
||||||
|
public.Use(
|
||||||
|
request.RequestID(),
|
||||||
|
request.RateLimit(s.kv, s.rateLimitConfig), // IP-based only
|
||||||
|
)
|
||||||
|
public.GET("/models", s.handler.ListModels)
|
||||||
|
public.GET("/health", s.handler.Health)
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
### Agent Service Integration
|
||||||
|
|
||||||
```go
|
```go
|
||||||
// agent/context/openapi.go
|
// agent/context/openapi.go
|
||||||
|
|
@ -780,6 +1092,7 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest
|
||||||
// Create context with request ID
|
// Create context with request ID
|
||||||
ctx := New(c.Request.Context(), authInfo, chatID)
|
ctx := New(c.Request.Context(), authInfo, chatID)
|
||||||
ctx.RequestID = requestID // Use global request_id
|
ctx.RequestID = requestID // Use global request_id
|
||||||
|
ctx.GinContext = c // Keep gin context for billing
|
||||||
|
|
||||||
// ...
|
// ...
|
||||||
}
|
}
|
||||||
|
|
@ -787,10 +1100,10 @@ func GetCompletionRequest(c *gin.Context, cache store.Store) (*CompletionRequest
|
||||||
// agent/assistant/agent.go
|
// agent/assistant/agent.go
|
||||||
func (ast *Assistant) Stream(ctx, inputMessages, options) {
|
func (ast *Assistant) Stream(ctx, inputMessages, options) {
|
||||||
defer func() {
|
defer func() {
|
||||||
// Update token usage in global request record
|
// Update token usage via billing middleware
|
||||||
if ctx.RequestID != "" && completionResponse != nil && completionResponse.Usage != nil {
|
if ctx.GinContext != nil && completionResponse != nil && completionResponse.Usage != nil {
|
||||||
request.UpdateTokenUsage(
|
request.UpdateTokenUsage(
|
||||||
ctx.RequestID,
|
ctx.GinContext,
|
||||||
completionResponse.Usage.PromptTokens,
|
completionResponse.Usage.PromptTokens,
|
||||||
completionResponse.Usage.CompletionTokens,
|
completionResponse.Usage.CompletionTokens,
|
||||||
)
|
)
|
||||||
|
|
@ -801,74 +1114,59 @@ func (ast *Assistant) Stream(ctx, inputMessages, options) {
|
||||||
}
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
### KB Service
|
### LLM Service Integration
|
||||||
|
|
||||||
```go
|
```go
|
||||||
// kb/api/search.go
|
// llm/api/completion.go
|
||||||
func (api *API) Search(c *gin.Context) {
|
func (api *API) Completion(c *gin.Context) {
|
||||||
requestID := c.GetString("request_id")
|
// ... execute LLM call ...
|
||||||
|
|
||||||
// Perform search...
|
// Update token usage
|
||||||
|
if response.Usage != nil {
|
||||||
// Update metadata if needed
|
request.UpdateTokenUsage(c, response.Usage.PromptTokens, response.Usage.CompletionTokens)
|
||||||
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
|
## Summary
|
||||||
|
|
||||||
### Components
|
### Middleware Components
|
||||||
|
|
||||||
| Component | Location | Responsibility |
|
| Middleware | File | Purpose | Storage |
|
||||||
| ------------ | ------------------------------- | ---------------------------- |
|
| ----------- | --------------- | ------------------------ | -------- |
|
||||||
| KV Store | `openapi/request/kv.go` | Real-time: rate limit, quota |
|
| `RequestID` | `request_id.go` | Generate request ID | - |
|
||||||
| SQL Store | `openapi/request/sql.go` | Archive: billing, audit |
|
| `RateLimit` | `ratelimit.go` | Frequency limiting | KV |
|
||||||
| Middleware | `openapi/request/middleware.go` | Track requests, orchestrate |
|
| `Quota` | `quota.go` | Token quota enforcement | KV |
|
||||||
| Rate Limiter | `openapi/request/ratelimit.go` | Enforce rate limits |
|
| `Metrics` | `metrics.go` | Duration/status tracking | - |
|
||||||
| Types | `openapi/request/types.go` | Data structures |
|
| `Archive` | `archive.go` | Persist to SQL | SQL |
|
||||||
|
| `Billing` | `billing.go` | Token usage tracking | KV + SQL |
|
||||||
|
|
||||||
### Storage Comparison
|
### Storage Components
|
||||||
|
|
||||||
| Operation | KV (Redis) | SQL (Archive) |
|
| Component | File | Purpose |
|
||||||
| ---------------- | -------------- | ------------- |
|
| --------- | ---------- | ---------------------------- |
|
||||||
| Rate limit check | ✅ Synchronous | ❌ Not used |
|
| KV Store | `kv.go` | Real-time: rate limit, quota |
|
||||||
| Quota check | ✅ Synchronous | ❌ Not used |
|
| SQL Store | `sql.go` | Archive: billing, audit |
|
||||||
| Request status | ✅ Synchronous | ❌ Not used |
|
| Types | `types.go` | Data structures |
|
||||||
| Token update | ✅ Synchronous | ✅ Async |
|
|
||||||
| Billing report | ❌ Not used | ✅ Query |
|
### Middleware Combinations by Use Case
|
||||||
| Audit log | ❌ Not used | ✅ Query |
|
|
||||||
|
| Use Case | RequestID | RateLimit | Quota | Metrics | Archive | Billing |
|
||||||
|
| ------------ | --------- | --------- | ----- | ------- | ------- | ------- |
|
||||||
|
| Agent API | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||||
|
| LLM API | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||||
|
| KB API | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
|
||||||
|
| File API | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ |
|
||||||
|
| Public API | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||||
|
| Internal API | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||||
|
|
||||||
### Key Points
|
### Key Points
|
||||||
|
|
||||||
1. **Two-layer storage**: KV for real-time, SQL for archive
|
1. **Modular design**: Each middleware is independent and composable
|
||||||
2. **KV operations are synchronous**: Rate limit and quota checks must be fast
|
2. **Business-driven composition**: Routes choose which middleware to use
|
||||||
3. **SQL writes are async**: Archive happens in background goroutine
|
3. **Two-layer storage**: KV for real-time, SQL for archive
|
||||||
4. **Services update tokens via `request_id`**: Updates both KV and SQL
|
4. **KV operations are synchronous**: Rate limit and quota checks must be fast
|
||||||
5. **KV data has TTL**: Auto-expires to prevent memory bloat
|
5. **SQL writes are async**: Archive happens in background goroutine
|
||||||
6. **SQL data is permanent**: For billing and compliance
|
6. **Services update tokens via gin context**: `request.UpdateTokenUsage(c, input, output)`
|
||||||
|
7. **KV data has TTL**: Auto-expires to prevent memory bloat
|
||||||
|
8. **SQL data is permanent**: For billing and compliance
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue