姚章浩
326 lines
12 KiB
Go
326 lines
12 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgtype"
|
|
|
|
"gitea.zheromiku.top/Weknora/dobotai-platform/internal/reasoningmode"
|
|
)
|
|
|
|
func (s *Store) ResolveState(ctx context.Context, scope reasoningmode.Scope, now time.Time) (reasoningmode.State, bool, error) {
|
|
if err := validateReasoningScope(scope); err != nil {
|
|
return reasoningmode.State{}, false, err
|
|
}
|
|
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
|
|
if err != nil {
|
|
return reasoningmode.State{}, false, fmt.Errorf("begin reasoning state resolve: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
if err = ensureReasoningState(ctx, tx, scope); err != nil {
|
|
return reasoningmode.State{}, false, err
|
|
}
|
|
state, err := selectReasoningStateForUpdate(ctx, tx, scope)
|
|
if err != nil {
|
|
return reasoningmode.State{}, false, err
|
|
}
|
|
effective, expired := reasoningmode.EffectiveState(state, now)
|
|
if expired {
|
|
if err = updateReasoningState(ctx, tx, scope, effective, now); err != nil {
|
|
return reasoningmode.State{}, false, err
|
|
}
|
|
}
|
|
if err = tx.Commit(ctx); err != nil {
|
|
return reasoningmode.State{}, false, fmt.Errorf("commit reasoning state resolve: %w", err)
|
|
}
|
|
return effective, expired, nil
|
|
}
|
|
|
|
func (s *Store) Enable(ctx context.Context, scope reasoningmode.Scope, now time.Time, idle time.Duration) (reasoningmode.State, error) {
|
|
if err := validateReasoningScope(scope); err != nil || idle <= 0 {
|
|
if err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
return reasoningmode.State{}, reasoningmode.ErrInvalidState
|
|
}
|
|
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
|
|
if err != nil {
|
|
return reasoningmode.State{}, fmt.Errorf("begin reasoning enable: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
if err = ensureReasoningState(ctx, tx, scope); err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
state, err := selectReasoningStateForUpdate(ctx, tx, scope)
|
|
if err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
state.Mode = reasoningmode.ModeReasoning
|
|
state.RoundsUsed = 0
|
|
state.LastActivityAt = now
|
|
state.ExpiresAt = now.Add(idle)
|
|
if err = updateReasoningState(ctx, tx, scope, state, now); err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
if err = tx.Commit(ctx); err != nil {
|
|
return reasoningmode.State{}, fmt.Errorf("commit reasoning enable: %w", err)
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func (s *Store) Disable(ctx context.Context, scope reasoningmode.Scope, now time.Time) (reasoningmode.State, bool, error) {
|
|
if err := validateReasoningScope(scope); err != nil {
|
|
return reasoningmode.State{}, false, err
|
|
}
|
|
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
|
|
if err != nil {
|
|
return reasoningmode.State{}, false, fmt.Errorf("begin reasoning disable: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
if err = ensureReasoningState(ctx, tx, scope); err != nil {
|
|
return reasoningmode.State{}, false, err
|
|
}
|
|
state, err := selectReasoningStateForUpdate(ctx, tx, scope)
|
|
if err != nil {
|
|
return reasoningmode.State{}, false, err
|
|
}
|
|
state, _ = reasoningmode.EffectiveState(state, now)
|
|
wasReasoning := state.Mode == reasoningmode.ModeReasoning
|
|
state.Mode = reasoningmode.ModeQuick
|
|
state.RoundsUsed = 0
|
|
state.ExpiresAt = time.Time{}
|
|
if err = updateReasoningState(ctx, tx, scope, state, now); err != nil {
|
|
return reasoningmode.State{}, false, err
|
|
}
|
|
if err = tx.Commit(ctx); err != nil {
|
|
return reasoningmode.State{}, false, fmt.Errorf("commit reasoning disable: %w", err)
|
|
}
|
|
return state, wasReasoning, nil
|
|
}
|
|
|
|
func (s *Store) SetMaxRounds(ctx context.Context, scope reasoningmode.Scope, maxRounds int, now time.Time) (reasoningmode.State, error) {
|
|
if err := validateReasoningScope(scope); err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
if maxRounds < 1 || maxRounds > 50 {
|
|
return reasoningmode.State{}, reasoningmode.ErrInvalidState
|
|
}
|
|
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
|
|
if err != nil {
|
|
return reasoningmode.State{}, fmt.Errorf("begin reasoning rounds update: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
if err = ensureReasoningState(ctx, tx, scope); err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
state, err := selectReasoningStateForUpdate(ctx, tx, scope)
|
|
if err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
state.MaxRounds = maxRounds
|
|
if state.RoundsUsed >= maxRounds {
|
|
state.Mode = reasoningmode.ModeQuick
|
|
state.RoundsUsed = 0
|
|
state.ExpiresAt = time.Time{}
|
|
}
|
|
if err = updateReasoningState(ctx, tx, scope, state, now); err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
if err = tx.Commit(ctx); err != nil {
|
|
return reasoningmode.State{}, fmt.Errorf("commit reasoning rounds update: %w", err)
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func (s *Store) SetHintEnabled(ctx context.Context, scope reasoningmode.Scope, enabled bool, now time.Time) (reasoningmode.State, error) {
|
|
if err := validateReasoningScope(scope); err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
|
|
if err != nil {
|
|
return reasoningmode.State{}, fmt.Errorf("begin reasoning hint update: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
if err = ensureReasoningState(ctx, tx, scope); err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
state, err := selectReasoningStateForUpdate(ctx, tx, scope)
|
|
if err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
state.HintEnabled = enabled
|
|
if err = updateReasoningState(ctx, tx, scope, state, now); err != nil {
|
|
return reasoningmode.State{}, err
|
|
}
|
|
if err = tx.Commit(ctx); err != nil {
|
|
return reasoningmode.State{}, fmt.Errorf("commit reasoning hint update: %w", err)
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func (s *Store) CompleteTurn(ctx context.Context, input reasoningmode.TurnCompletion) (reasoningmode.Completion, error) {
|
|
if err := validateReasoningScope(input.Scope); err != nil {
|
|
return reasoningmode.Completion{}, err
|
|
}
|
|
if strings.TrimSpace(input.TurnID) == "" || input.At.IsZero() || input.Idle <= 0 {
|
|
return reasoningmode.Completion{}, reasoningmode.ErrInvalidState
|
|
}
|
|
tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{})
|
|
if err != nil {
|
|
return reasoningmode.Completion{}, fmt.Errorf("begin reasoning turn completion: %w", err)
|
|
}
|
|
defer func() { _ = tx.Rollback(ctx) }()
|
|
|
|
var existing reasoningmode.Completion
|
|
err = tx.QueryRow(ctx, `SELECT round_number,showed_hint,auto_closed,showed_onboarding
|
|
FROM reasoning_mode_turns
|
|
WHERE tenant_key=$1 AND channel=$2 AND turn_id=$3 AND external_actor_id=$4
|
|
FOR UPDATE`, input.Scope.TenantKey, input.Scope.Channel, input.TurnID, input.Scope.ExternalActorID).Scan(
|
|
&existing.RoundNumber, &existing.ShowHint, &existing.AutoClosed, &existing.ShowOnboarding,
|
|
)
|
|
if err == nil {
|
|
if existing.Next, err = selectReasoningState(ctx, tx, input.Scope); err != nil {
|
|
return reasoningmode.Completion{}, err
|
|
}
|
|
if err = tx.Commit(ctx); err != nil {
|
|
return reasoningmode.Completion{}, fmt.Errorf("commit replayed reasoning completion: %w", err)
|
|
}
|
|
return existing, nil
|
|
}
|
|
if !errors.Is(err, pgx.ErrNoRows) {
|
|
return reasoningmode.Completion{}, fmt.Errorf("find reasoning turn completion: %w", err)
|
|
}
|
|
var status string
|
|
if err = tx.QueryRow(ctx, `SELECT t.status FROM turns t
|
|
JOIN messages m ON m.id=t.request_message_id
|
|
JOIN channel_actors a ON a.id=m.actor_id
|
|
WHERE t.id=$1 AND t.tenant_key=$2 AND t.channel=$3
|
|
AND a.external_actor_id=$4
|
|
FOR UPDATE OF t`, input.TurnID, input.Scope.TenantKey, input.Scope.Channel, input.Scope.ExternalActorID).Scan(&status); err != nil {
|
|
return reasoningmode.Completion{}, fmt.Errorf("lock reasoning completion turn: %w", err)
|
|
}
|
|
if status != "succeeded" {
|
|
return reasoningmode.Completion{}, errors.New("reasoning completion requires a succeeded turn")
|
|
}
|
|
|
|
if err = ensureReasoningState(ctx, tx, input.Scope); err != nil {
|
|
return reasoningmode.Completion{}, err
|
|
}
|
|
state, err := selectReasoningStateForUpdate(ctx, tx, input.Scope)
|
|
if err != nil {
|
|
return reasoningmode.Completion{}, err
|
|
}
|
|
completion, err := reasoningmode.ApplySuccessfulTurn(state, input.Execution, input.At, input.Idle)
|
|
if err != nil {
|
|
return reasoningmode.Completion{}, err
|
|
}
|
|
if err = updateReasoningState(ctx, tx, input.Scope, completion.Next, input.At); err != nil {
|
|
return reasoningmode.Completion{}, err
|
|
}
|
|
tag, err := tx.Exec(ctx, `INSERT INTO reasoning_mode_turns
|
|
(id,tenant_key,channel,turn_id,external_actor_id,execution_mode,round_number,showed_hint,auto_closed,showed_onboarding)
|
|
VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9,$10)
|
|
ON CONFLICT (tenant_key,channel,turn_id) DO NOTHING`,
|
|
newSettingsID(), input.Scope.TenantKey, input.Scope.Channel, input.TurnID,
|
|
input.Scope.ExternalActorID, string(input.Execution), completion.RoundNumber,
|
|
completion.ShowHint, completion.AutoClosed, completion.ShowOnboarding,
|
|
)
|
|
if err != nil {
|
|
return reasoningmode.Completion{}, fmt.Errorf("insert reasoning turn completion: %w", err)
|
|
}
|
|
if tag.RowsAffected() != 1 {
|
|
return reasoningmode.Completion{}, errors.New("insert reasoning turn completion affected unexpected rows")
|
|
}
|
|
if err = tx.Commit(ctx); err != nil {
|
|
return reasoningmode.Completion{}, fmt.Errorf("commit reasoning turn completion: %w", err)
|
|
}
|
|
return completion, nil
|
|
}
|
|
|
|
func ensureReasoningState(ctx context.Context, tx pgx.Tx, scope reasoningmode.Scope) error {
|
|
_, err := tx.Exec(ctx, `INSERT INTO reasoning_mode_states
|
|
(tenant_key,channel,external_actor_id,mode,max_rounds,rounds_used,hint_enabled)
|
|
VALUES($1,$2,$3,'quick',10,0,TRUE)
|
|
ON CONFLICT (tenant_key,channel,external_actor_id) DO NOTHING`,
|
|
scope.TenantKey, scope.Channel, scope.ExternalActorID,
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("ensure reasoning state: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func selectReasoningStateForUpdate(ctx context.Context, tx pgx.Tx, scope reasoningmode.Scope) (reasoningmode.State, error) {
|
|
return scanReasoningState(tx.QueryRow(ctx, `SELECT mode,max_rounds,rounds_used,hint_enabled,last_activity_at,expires_at,onboarding_shown_at
|
|
FROM reasoning_mode_states
|
|
WHERE tenant_key=$1 AND channel=$2 AND external_actor_id=$3
|
|
FOR UPDATE`, scope.TenantKey, scope.Channel, scope.ExternalActorID))
|
|
}
|
|
|
|
func selectReasoningState(ctx context.Context, tx pgx.Tx, scope reasoningmode.Scope) (reasoningmode.State, error) {
|
|
return scanReasoningState(tx.QueryRow(ctx, `SELECT mode,max_rounds,rounds_used,hint_enabled,last_activity_at,expires_at,onboarding_shown_at
|
|
FROM reasoning_mode_states
|
|
WHERE tenant_key=$1 AND channel=$2 AND external_actor_id=$3`, scope.TenantKey, scope.Channel, scope.ExternalActorID))
|
|
}
|
|
|
|
func scanReasoningState(row pgx.Row) (reasoningmode.State, error) {
|
|
var (
|
|
state reasoningmode.State
|
|
mode string
|
|
last pgtype.Timestamptz
|
|
expires pgtype.Timestamptz
|
|
onboarding pgtype.Timestamptz
|
|
)
|
|
if err := row.Scan(&mode, &state.MaxRounds, &state.RoundsUsed, &state.HintEnabled, &last, &expires, &onboarding); err != nil {
|
|
return reasoningmode.State{}, fmt.Errorf("scan reasoning state: %w", err)
|
|
}
|
|
state.Mode = reasoningmode.Mode(mode)
|
|
if last.Valid {
|
|
state.LastActivityAt = last.Time
|
|
}
|
|
if expires.Valid {
|
|
state.ExpiresAt = expires.Time
|
|
}
|
|
if onboarding.Valid {
|
|
state.OnboardingShown = true
|
|
state.OnboardingShownAt = onboarding.Time
|
|
}
|
|
return state, nil
|
|
}
|
|
|
|
func updateReasoningState(ctx context.Context, tx pgx.Tx, scope reasoningmode.Scope, state reasoningmode.State, now time.Time) error {
|
|
_, err := tx.Exec(ctx, `UPDATE reasoning_mode_states SET
|
|
mode=$4,max_rounds=$5,rounds_used=$6,hint_enabled=$7,
|
|
last_activity_at=$8,expires_at=$9,onboarding_shown_at=$10,updated_at=$11
|
|
WHERE tenant_key=$1 AND channel=$2 AND external_actor_id=$3`,
|
|
scope.TenantKey, scope.Channel, scope.ExternalActorID,
|
|
string(state.Mode), state.MaxRounds, state.RoundsUsed, state.HintEnabled,
|
|
nullableTime(state.LastActivityAt), nullableTime(state.ExpiresAt), nullableTime(state.OnboardingShownAt), now,
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("update reasoning state: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func nullableTime(value time.Time) any {
|
|
if value.IsZero() {
|
|
return nil
|
|
}
|
|
return value
|
|
}
|
|
|
|
func validateReasoningScope(scope reasoningmode.Scope) error {
|
|
if strings.TrimSpace(scope.TenantKey) == "" || strings.TrimSpace(scope.Channel) == "" || strings.TrimSpace(scope.ExternalActorID) == "" {
|
|
return errors.New("reasoning mode scope is required")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
var _ reasoningmode.Repository = (*Store)(nil)
|