easyai-ai-gateway/apps/api/internal/runner/cache_affinity.go

70 lines
1.6 KiB
Go

package runner
import (
"crypto/sha256"
"encoding/hex"
"encoding/json"
"strings"
)
func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
if len(body) == 0 {
return ""
}
for _, key := range []string{"cacheAffinityKey", "cache_affinity_key"} {
if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
return "explicit:" + sha256Text(value)
}
}
for _, key := range []string{"sessionId", "session_id", "conversationId", "conversation_id"} {
if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
return key + ":" + sha256Text(value)
}
}
if !cacheAffinityPromptHashSupported(kind, modelType) {
return ""
}
payload := map[string]any{
"kind": kind,
"modelType": modelType,
}
if messages, ok := body["messages"]; ok {
payload["messages"] = messages
}
if tools, ok := body["tools"]; ok {
payload["tools"] = tools
}
if instructions, ok := body["instructions"]; ok {
payload["instructions"] = instructions
}
if input, ok := body["input"]; ok {
payload["input"] = input
}
if len(payload) <= 2 {
return ""
}
raw, err := json.Marshal(payload)
if err != nil || len(raw) == 0 {
return ""
}
return "prompt:" + sha256Text(string(raw))
}
func cacheAffinityPromptHashSupported(kind string, modelType string) bool {
switch strings.TrimSpace(kind) {
case "chat.completions", "responses":
return true
}
switch strings.TrimSpace(modelType) {
case "chat", "text_generate", "responses":
return true
default:
return false
}
}
func sha256Text(value string) string {
sum := sha256.Sum256([]byte(value))
return hex.EncodeToString(sum[:])
}