pro/olons_server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: async ai workbench tasks

Browse Source
This commit is contained in:
richjohnrambo
2026-09-30 11:37:49 +08:00
parent 47a10c93be
commit d2e70efd9d
4 changed files with 493 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/store/ai.go
+74 -4
View File
@@ -22,6 +22,7 @@ type AICapability struct {
type AITask struct {
TaskID string `db:"task_id"`
TenantID int64 `db:"tenant_id"`
TenantCode string `db:"tenant_code"`
UserID string `db:"user_id"`
Feature string `db:"feature"`
@@ -30,6 +31,7 @@ type AITask struct {
SourceType string `db:"source_type"`
SourceRefID string `db:"source_ref_id"`
InputHash string `db:"input_hash"`
InputSnapshot sql.NullString `db:"input_snapshot"`
PromptVersion string `db:"prompt_version"`
Model string `db:"model"`
ResultRefID string `db:"result_ref_id"`
@@ -181,17 +183,59 @@ func (s *Store) CreateAITask(ctx context.Context, task AITask, result any) error
if err != nil {
return err
}
var inputSnapshot any
if task.InputSnapshot.Valid {
inputSnapshot = task.InputSnapshot.String
}
_, err = s.db.ExecContext(ctx, `INSERT INTO wb_ai_task
(task_id,tenant_id,tenant_code,user_id,feature,status,progress,source_type,source_ref_id,input_hash,
input_snapshot,prompt_version,model,result_ref_id,result_json,error_reason,error_message,retry_count,
started_at,finished_at,created_at,updated_at)
VALUES(?,?,?,?,?,?,?,?,?,?,NULL,?,?,?,?,?,?,0,?,?,?,?)`,
VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,0,?,?,?,?)`,
task.TaskID, s.tenant(), task.TenantCode, task.UserID, task.Feature, task.Status, task.Progress,
task.SourceType, task.SourceRefID, task.InputHash, task.PromptVersion, task.Model, task.ResultRefID,
task.SourceType, task.SourceRefID, task.InputHash, inputSnapshot, task.PromptVersion, task.Model, task.ResultRefID,
resultJSON, task.ErrorReason, task.ErrorMessage, task.StartedAt, task.FinishedAt, now, now)
return err
}
func (s *Store) ClaimQueuedAITask(ctx context.Context) (*AITask, error) {
tx, err := s.db.BeginTxx(ctx, nil)
if err != nil {
return nil, err
}
defer tx.Rollback()
var task AITask
err = tx.GetContext(ctx, &task, `SELECT task_id,tenant_id,tenant_code,user_id,feature,status,progress,source_type,source_ref_id,
input_hash,input_snapshot,prompt_version,model,result_ref_id,result_json,error_reason,error_message,retry_count,
started_at,finished_at,created_at,updated_at
FROM wb_ai_task
WHERE status='queued' AND feature IN ('smart_document','smart_table','meeting_summary')
ORDER BY created_at ASC,id ASC LIMIT 1 FOR UPDATE SKIP LOCKED`)
if err == sql.ErrNoRows {
return nil, ErrNotFound
}
if err != nil {
return nil, err
}
now := time.Now().UTC().Round(time.Millisecond)
res, err := tx.ExecContext(ctx, `UPDATE wb_ai_task SET status='processing',progress=?,started_at=COALESCE(started_at,?),updated_at=?
WHERE tenant_id=? AND task_id=? AND status='queued'`, maxInt(task.Progress, 10), now, now, task.TenantID, task.TaskID)
if err != nil {
return nil, err
}
if rows, _ := res.RowsAffected(); rows == 0 {
return nil, ErrNotFound
}
if err := tx.Commit(); err != nil {
return nil, err
}
task.Status = "processing"
task.Progress = maxInt(task.Progress, 10)
task.StartedAt = &now
task.UpdatedAt = now
return &task, nil
}
func (s *Store) UpdateAITaskResult(ctx context.Context, taskID, status string, progress int, resultRefID string, result any, reason, message string) error {
now := time.Now().UTC().Round(time.Millisecond)
resultJSON, err := marshalNullable(result)
@@ -210,6 +254,25 @@ func (s *Store) UpdateAITaskResult(ctx context.Context, taskID, status string, p
return nil
}
func (s *Store) FinishAITaskResult(ctx context.Context, taskID, status string, progress int, resultRefID string, result any, reason, message string) (bool, error) {
now := time.Now().UTC().Round(time.Millisecond)
resultJSON, err := marshalNullable(result)
if err != nil {
return false, err
}
res, err := s.db.ExecContext(ctx, `UPDATE wb_ai_task SET status=?,progress=?,result_ref_id=?,result_json=?,
error_reason=?,error_message=?,finished_at=?,updated_at=? WHERE tenant_id=? AND task_id=? AND status='processing'`,
status, progress, resultRefID, resultJSON, reason, message, now, now, s.tenant(), taskID)
if err != nil {
return false, err
}
rows, err := res.RowsAffected()
if err != nil {
return false, err
}
return rows > 0, nil
}
func (s *Store) CancelAITask(ctx context.Context, userID int64, taskID string) error {
now := time.Now().UTC().Round(time.Millisecond)
res, err := s.db.ExecContext(ctx, `UPDATE wb_ai_task SET status='canceled',progress=100,finished_at=?,updated_at=?
@@ -226,8 +289,8 @@ func (s *Store) CancelAITask(ctx context.Context, userID int64, taskID string) e
func (s *Store) AITask(ctx context.Context, userID int64, taskID string) (*AITask, error) {
var task AITask
err := s.db.GetContext(ctx, &task, `SELECT task_id,tenant_code,user_id,feature,status,progress,source_type,source_ref_id,
input_hash,prompt_version,model,result_ref_id,result_json,error_reason,error_message,retry_count,
err := s.db.GetContext(ctx, &task, `SELECT task_id,tenant_id,tenant_code,user_id,feature,status,progress,source_type,source_ref_id,
input_hash,input_snapshot,prompt_version,model,result_ref_id,result_json,error_reason,error_message,retry_count,
started_at,finished_at,created_at,updated_at
FROM wb_ai_task WHERE tenant_id=? AND user_id=? AND task_id=?`, s.tenant(), formatInt64(userID), taskID)
if err == sql.ErrNoRows {
@@ -527,3 +590,10 @@ func marshalNullable(value any) (any, error) {
func formatInt64(value int64) string {
return strconv.FormatInt(value, 10)
}
func maxInt(a, b int) int {
if a > b {
return a
}
return b
}