Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
1802 changed files with 503346 additions and 2 deletions
414
backend/internal/cache/xai_reasoning_replay_cache.go
vendored
Normal file
414
backend/internal/cache/xai_reasoning_replay_cache.go
vendored
Normal file
|
|
@ -0,0 +1,414 @@
|
|||
package cache
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
homekv "github.com/router-for-me/CLIProxyAPI/v7/internal/home"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/signature"
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
const (
|
||||
// XAIReasoningReplayCacheTTL limits how long encrypted reasoning replay
|
||||
// items stay in process memory.
|
||||
XAIReasoningReplayCacheTTL = 1 * time.Hour
|
||||
|
||||
// XAIReasoningReplayCacheMaxEntries bounds process memory for replay
|
||||
// continuity. Oldest entries are evicted first.
|
||||
XAIReasoningReplayCacheMaxEntries = 10240
|
||||
|
||||
// XAIReasoningReplayCacheEvictBatchSize leaves headroom after the cache
|
||||
// reaches capacity so high write volume does not rescan the map every turn.
|
||||
XAIReasoningReplayCacheEvictBatchSize = 128
|
||||
)
|
||||
|
||||
type xaiReasoningReplayEntry struct {
|
||||
Items [][]byte
|
||||
Timestamp time.Time
|
||||
}
|
||||
|
||||
var (
|
||||
xaiReasoningReplayMu sync.Mutex
|
||||
xaiReasoningReplayEntries = make(map[string]xaiReasoningReplayEntry)
|
||||
)
|
||||
|
||||
type xaiReasoningReplayKVClient interface {
|
||||
KVGet(ctx context.Context, key string) ([]byte, bool, error)
|
||||
KVSet(ctx context.Context, key string, value []byte, opts homekv.KVSetOptions) (bool, error)
|
||||
KVDel(ctx context.Context, keys ...string) (int64, error)
|
||||
KVExpire(ctx context.Context, key string, ttl time.Duration) (bool, error)
|
||||
}
|
||||
|
||||
var currentXAIReasoningReplayKVClient = func() (xaiReasoningReplayKVClient, bool, error) {
|
||||
return homekv.CurrentKVClient()
|
||||
}
|
||||
|
||||
// CacheXAIReasoningReplayItem stores a final Grok reasoning item for stateless
|
||||
// replay. The stored item is normalized to the minimal shape accepted by
|
||||
// Responses input replay.
|
||||
func CacheXAIReasoningReplayItem(modelName, sessionKey string, item []byte) bool {
|
||||
return CacheXAIReasoningReplayItems(modelName, sessionKey, [][]byte{item})
|
||||
}
|
||||
|
||||
// CacheXAIReasoningReplayItems stores the final Grok assistant output items
|
||||
// needed to replay a stateless next turn.
|
||||
func CacheXAIReasoningReplayItems(modelName, sessionKey string, items [][]byte) bool {
|
||||
return CacheXAIReasoningReplayItemsBestEffort(context.Background(), modelName, sessionKey, items)
|
||||
}
|
||||
|
||||
// XAIReasoningReplayStoreStatus reports why a completed-turn cache write
|
||||
// succeeded or failed so callers can decide whether to keep prior entries.
|
||||
type XAIReasoningReplayStoreStatus int
|
||||
|
||||
const (
|
||||
// XAIReasoningReplayStoreInvalidArgs means model/session were empty.
|
||||
XAIReasoningReplayStoreInvalidArgs XAIReasoningReplayStoreStatus = iota
|
||||
// XAIReasoningReplayStored means a valid reasoning batch was written.
|
||||
XAIReasoningReplayStored
|
||||
// XAIReasoningReplayNoReplayableState means the completed output had no
|
||||
// cacheable reasoning batch (for example reasoning disabled).
|
||||
XAIReasoningReplayNoReplayableState
|
||||
// XAIReasoningReplayStoreBackendError means normalize succeeded but the
|
||||
// storage backend failed; previous entries should be retained.
|
||||
XAIReasoningReplayStoreBackendError
|
||||
)
|
||||
|
||||
// CacheXAIReasoningReplayItemsBestEffort stores replay items for completed response paths.
|
||||
func CacheXAIReasoningReplayItemsBestEffort(ctx context.Context, modelName, sessionKey string, items [][]byte) bool {
|
||||
return StoreXAIReasoningReplayItems(ctx, modelName, sessionKey, items) == XAIReasoningReplayStored
|
||||
}
|
||||
|
||||
// StoreXAIReasoningReplayItems stores replay items and distinguishes empty
|
||||
// completed state from backend failures.
|
||||
func StoreXAIReasoningReplayItems(ctx context.Context, modelName, sessionKey string, items [][]byte) XAIReasoningReplayStoreStatus {
|
||||
key := xaiReasoningReplayCacheKey(modelName, sessionKey)
|
||||
if key == "" {
|
||||
return XAIReasoningReplayStoreInvalidArgs
|
||||
}
|
||||
normalized, ok := normalizeXAIReasoningReplayItems(items)
|
||||
if !ok {
|
||||
return XAIReasoningReplayNoReplayableState
|
||||
}
|
||||
if client, homeMode, errClient := currentXAIReasoningReplayKVClient(); homeMode {
|
||||
if errClient != nil {
|
||||
log.Errorf("home kv best-effort xai reasoning replay set failed prefix=cpa:xai:*: %v", errClient)
|
||||
return XAIReasoningReplayStoreBackendError
|
||||
}
|
||||
raw, errMarshal := json.Marshal(normalized)
|
||||
if errMarshal != nil {
|
||||
log.Errorf("home kv best-effort xai reasoning replay set failed prefix=cpa:xai:*: %v", errMarshal)
|
||||
return XAIReasoningReplayStoreBackendError
|
||||
}
|
||||
written, errSet := client.KVSet(ctx, xaiReasoningReplayKVKey(modelName, sessionKey), raw, homekv.KVSetOptions{EX: XAIReasoningReplayCacheTTL})
|
||||
if errSet != nil {
|
||||
log.Errorf("home kv best-effort xai reasoning replay set failed prefix=cpa:xai:*: %v", errSet)
|
||||
return XAIReasoningReplayStoreBackendError
|
||||
}
|
||||
if !written {
|
||||
return XAIReasoningReplayStoreBackendError
|
||||
}
|
||||
return XAIReasoningReplayStored
|
||||
}
|
||||
|
||||
cacheCleanupOnce.Do(startCacheCleanup)
|
||||
now := time.Now()
|
||||
xaiReasoningReplayMu.Lock()
|
||||
defer xaiReasoningReplayMu.Unlock()
|
||||
xaiReasoningReplayEntries[key] = xaiReasoningReplayEntry{
|
||||
Items: normalized,
|
||||
Timestamp: now,
|
||||
}
|
||||
if len(xaiReasoningReplayEntries) > XAIReasoningReplayCacheMaxEntries {
|
||||
evictOldestXAIReasoningReplayEntriesLocked(XAIReasoningReplayCacheEvictBatchSize)
|
||||
}
|
||||
return XAIReasoningReplayStored
|
||||
}
|
||||
|
||||
// GetXAIReasoningReplayItem retrieves a normalized reasoning replay item.
|
||||
func GetXAIReasoningReplayItem(modelName, sessionKey string) ([]byte, bool) {
|
||||
items, ok := GetXAIReasoningReplayItems(modelName, sessionKey)
|
||||
if !ok || len(items) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return items[0], true
|
||||
}
|
||||
|
||||
// GetXAIReasoningReplayItems retrieves normalized assistant output items.
|
||||
func GetXAIReasoningReplayItems(modelName, sessionKey string) ([][]byte, bool) {
|
||||
items, ok, err := GetXAIReasoningReplayItemsRequired(context.Background(), modelName, sessionKey)
|
||||
if err == nil {
|
||||
return items, ok
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// GetXAIReasoningReplayItemsRequired retrieves replay items for request-time paths.
|
||||
func GetXAIReasoningReplayItemsRequired(ctx context.Context, modelName, sessionKey string) ([][]byte, bool, error) {
|
||||
key := xaiReasoningReplayCacheKey(modelName, sessionKey)
|
||||
if key == "" {
|
||||
return nil, false, nil
|
||||
}
|
||||
client, homeMode, errClient := currentXAIReasoningReplayKVClient()
|
||||
if homeMode {
|
||||
if errClient != nil {
|
||||
return nil, false, errClient
|
||||
}
|
||||
raw, found, errGet := client.KVGet(ctx, xaiReasoningReplayKVKey(modelName, sessionKey))
|
||||
if errGet != nil || !found {
|
||||
return nil, false, errGet
|
||||
}
|
||||
var homeItems [][]byte
|
||||
if errUnmarshal := json.Unmarshal(raw, &homeItems); errUnmarshal != nil {
|
||||
return nil, false, errUnmarshal
|
||||
}
|
||||
if _, errExpire := client.KVExpire(ctx, xaiReasoningReplayKVKey(modelName, sessionKey), XAIReasoningReplayCacheTTL); errExpire != nil {
|
||||
log.Warnf("home kv xai reasoning replay expire failed prefix=cpa:xai:*: %v", errExpire)
|
||||
}
|
||||
return cloneXAIReasoningReplayItems(homeItems), true, nil
|
||||
}
|
||||
|
||||
cacheCleanupOnce.Do(startCacheCleanup)
|
||||
now := time.Now()
|
||||
xaiReasoningReplayMu.Lock()
|
||||
defer xaiReasoningReplayMu.Unlock()
|
||||
entry, ok := xaiReasoningReplayEntries[key]
|
||||
if !ok {
|
||||
return nil, false, nil
|
||||
}
|
||||
if now.Sub(entry.Timestamp) > XAIReasoningReplayCacheTTL {
|
||||
delete(xaiReasoningReplayEntries, key)
|
||||
return nil, false, nil
|
||||
}
|
||||
entry.Timestamp = now
|
||||
xaiReasoningReplayEntries[key] = entry
|
||||
return cloneXAIReasoningReplayItems(entry.Items), true, nil
|
||||
}
|
||||
|
||||
// DeleteXAIReasoningReplayItem removes one replay item after upstream rejects
|
||||
// it or the caller otherwise knows it is stale.
|
||||
func DeleteXAIReasoningReplayItem(modelName, sessionKey string) {
|
||||
if errDelete := DeleteXAIReasoningReplayItemRequired(context.Background(), modelName, sessionKey); errDelete != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// DeleteXAIReasoningReplayItemRequired removes one replay item for request-time paths.
|
||||
func DeleteXAIReasoningReplayItemRequired(ctx context.Context, modelName, sessionKey string) error {
|
||||
key := xaiReasoningReplayCacheKey(modelName, sessionKey)
|
||||
if key == "" {
|
||||
return nil
|
||||
}
|
||||
client, homeMode, errClient := currentXAIReasoningReplayKVClient()
|
||||
if homeMode {
|
||||
if errClient != nil {
|
||||
return errClient
|
||||
}
|
||||
_, errDel := client.KVDel(ctx, xaiReasoningReplayKVKey(modelName, sessionKey))
|
||||
return errDel
|
||||
}
|
||||
xaiReasoningReplayMu.Lock()
|
||||
delete(xaiReasoningReplayEntries, key)
|
||||
xaiReasoningReplayMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// ClearXAIReasoningReplayCache clears all xAI reasoning replay state.
|
||||
func ClearXAIReasoningReplayCache() {
|
||||
xaiReasoningReplayMu.Lock()
|
||||
xaiReasoningReplayEntries = make(map[string]xaiReasoningReplayEntry)
|
||||
xaiReasoningReplayMu.Unlock()
|
||||
}
|
||||
|
||||
func xaiReasoningReplayCacheKey(modelName, sessionKey string) string {
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
sessionKey = strings.TrimSpace(sessionKey)
|
||||
if modelName == "" || sessionKey == "" {
|
||||
return ""
|
||||
}
|
||||
// The session key is the continuity boundary. Keep this independent from
|
||||
// the selected upstream xAI credential so auth failover can preserve replay.
|
||||
return strings.Join([]string{"xai-reasoning-replay", modelName, sessionKey}, "\x00")
|
||||
}
|
||||
|
||||
func xaiReasoningReplayKVKey(modelName, sessionKey string) string {
|
||||
return "cpa:xai:reasoning-replay:" + homekv.HashKeyPart(strings.TrimSpace(modelName)) + ":" + homekv.HashKeyPart(strings.TrimSpace(sessionKey))
|
||||
}
|
||||
|
||||
func normalizeXAIReasoningReplayItems(items [][]byte) ([][]byte, bool) {
|
||||
normalized := make([][]byte, 0, len(items))
|
||||
hasReplayAnchor := false
|
||||
for _, item := range items {
|
||||
normalizedItem, ok := normalizeXAIReasoningReplayItem(item)
|
||||
if ok {
|
||||
normalized = append(normalized, normalizedItem)
|
||||
switch strings.TrimSpace(gjson.GetBytes(normalizedItem, "type").String()) {
|
||||
case "reasoning", "function_call", "custom_tool_call":
|
||||
hasReplayAnchor = true
|
||||
}
|
||||
}
|
||||
}
|
||||
return normalized, hasReplayAnchor
|
||||
}
|
||||
|
||||
func normalizeXAIReasoningReplayItem(item []byte) ([]byte, bool) {
|
||||
itemResult := gjson.ParseBytes(item)
|
||||
switch strings.TrimSpace(itemResult.Get("type").String()) {
|
||||
case "reasoning":
|
||||
return normalizeXAIReasoningReplayReasoningItem(itemResult)
|
||||
case "message":
|
||||
return normalizeXAIReasoningReplayMessageItem(itemResult)
|
||||
case "function_call":
|
||||
return normalizeXAIReasoningReplayFunctionCallItem(itemResult)
|
||||
case "custom_tool_call":
|
||||
return normalizeXAIReasoningReplayCustomToolCallItem(itemResult)
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeXAIReasoningReplayReasoningItem(itemResult gjson.Result) ([]byte, bool) {
|
||||
encryptedContentResult := itemResult.Get("encrypted_content")
|
||||
if encryptedContentResult.Type != gjson.String {
|
||||
return nil, false
|
||||
}
|
||||
encryptedContent := encryptedContentResult.String()
|
||||
if encryptedContent != strings.TrimSpace(encryptedContent) {
|
||||
return nil, false
|
||||
}
|
||||
if _, err := signature.InspectGrokEncryptedContent(encryptedContent); err != nil {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
normalized := []byte(`{"type":"reasoning","summary":[],"content":null}`)
|
||||
normalized, _ = sjson.SetBytes(normalized, "encrypted_content", encryptedContent)
|
||||
return normalized, true
|
||||
}
|
||||
|
||||
func normalizeXAIReasoningReplayMessageItem(itemResult gjson.Result) ([]byte, bool) {
|
||||
if !strings.EqualFold(strings.TrimSpace(itemResult.Get("role").String()), "assistant") {
|
||||
return nil, false
|
||||
}
|
||||
content := itemResult.Get("content")
|
||||
if !content.IsArray() || len(content.Array()) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
normalized := []byte(`{"type":"message","role":"assistant","content":[]}`)
|
||||
for _, part := range content.Array() {
|
||||
partType := strings.TrimSpace(part.Get("type").String())
|
||||
var nextPart []byte
|
||||
switch partType {
|
||||
case "output_text":
|
||||
textValue := part.Get("text")
|
||||
if textValue.Type != gjson.String {
|
||||
continue
|
||||
}
|
||||
nextPart = []byte(`{"type":"output_text","text":""}`)
|
||||
nextPart, _ = sjson.SetBytes(nextPart, "text", textValue.String())
|
||||
case "refusal":
|
||||
// Responses API refusal parts use the "refusal" field, not "text".
|
||||
refusalValue := part.Get("refusal")
|
||||
if refusalValue.Type != gjson.String {
|
||||
continue
|
||||
}
|
||||
nextPart = []byte(`{"type":"refusal","refusal":""}`)
|
||||
nextPart, _ = sjson.SetBytes(nextPart, "refusal", refusalValue.String())
|
||||
default:
|
||||
continue
|
||||
}
|
||||
updated, errSet := sjson.SetRawBytes(normalized, "content.-1", nextPart)
|
||||
if errSet != nil {
|
||||
return nil, false
|
||||
}
|
||||
normalized = updated
|
||||
}
|
||||
if len(gjson.GetBytes(normalized, "content").Array()) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return normalized, true
|
||||
}
|
||||
|
||||
func normalizeXAIReasoningReplayFunctionCallItem(itemResult gjson.Result) ([]byte, bool) {
|
||||
callID := strings.TrimSpace(itemResult.Get("call_id").String())
|
||||
name := strings.TrimSpace(itemResult.Get("name").String())
|
||||
arguments := itemResult.Get("arguments")
|
||||
if callID == "" || name == "" || arguments.Type != gjson.String {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
normalized := []byte(`{"type":"function_call"}`)
|
||||
normalized, _ = sjson.SetBytes(normalized, "call_id", callID)
|
||||
normalized, _ = sjson.SetBytes(normalized, "name", name)
|
||||
normalized, _ = sjson.SetBytes(normalized, "arguments", arguments.String())
|
||||
return normalized, true
|
||||
}
|
||||
|
||||
func normalizeXAIReasoningReplayCustomToolCallItem(itemResult gjson.Result) ([]byte, bool) {
|
||||
callID := strings.TrimSpace(itemResult.Get("call_id").String())
|
||||
name := strings.TrimSpace(itemResult.Get("name").String())
|
||||
input := itemResult.Get("input")
|
||||
if callID == "" || name == "" || !input.Exists() {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
normalized := []byte(`{"type":"custom_tool_call","status":"completed"}`)
|
||||
if status := strings.TrimSpace(itemResult.Get("status").String()); status != "" {
|
||||
normalized, _ = sjson.SetBytes(normalized, "status", status)
|
||||
}
|
||||
normalized, _ = sjson.SetBytes(normalized, "call_id", callID)
|
||||
normalized, _ = sjson.SetBytes(normalized, "name", name)
|
||||
if input.Type == gjson.String {
|
||||
normalized, _ = sjson.SetBytes(normalized, "input", input.String())
|
||||
} else {
|
||||
normalized, _ = sjson.SetRawBytes(normalized, "input", []byte(input.Raw))
|
||||
}
|
||||
return normalized, true
|
||||
}
|
||||
|
||||
func cloneXAIReasoningReplayItems(items [][]byte) [][]byte {
|
||||
cloned := make([][]byte, 0, len(items))
|
||||
for _, item := range items {
|
||||
cloned = append(cloned, append([]byte(nil), item...))
|
||||
}
|
||||
return cloned
|
||||
}
|
||||
|
||||
func evictOldestXAIReasoningReplayEntriesLocked(count int) {
|
||||
if count <= 0 || len(xaiReasoningReplayEntries) == 0 {
|
||||
return
|
||||
}
|
||||
type candidate struct {
|
||||
key string
|
||||
timestamp time.Time
|
||||
}
|
||||
candidates := make([]candidate, 0, len(xaiReasoningReplayEntries))
|
||||
for key, entry := range xaiReasoningReplayEntries {
|
||||
candidates = append(candidates, candidate{key: key, timestamp: entry.Timestamp})
|
||||
}
|
||||
sort.Slice(candidates, func(i, j int) bool {
|
||||
return candidates[i].timestamp.Before(candidates[j].timestamp)
|
||||
})
|
||||
if count > len(candidates) {
|
||||
count = len(candidates)
|
||||
}
|
||||
for i := 0; i < count; i++ {
|
||||
delete(xaiReasoningReplayEntries, candidates[i].key)
|
||||
}
|
||||
}
|
||||
|
||||
func purgeExpiredXAIReasoningReplayCache(now time.Time) {
|
||||
xaiReasoningReplayMu.Lock()
|
||||
for key, entry := range xaiReasoningReplayEntries {
|
||||
if now.Sub(entry.Timestamp) > XAIReasoningReplayCacheTTL {
|
||||
delete(xaiReasoningReplayEntries, key)
|
||||
}
|
||||
}
|
||||
xaiReasoningReplayMu.Unlock()
|
||||
}
|
||||
Loading…
Reference in a new issue