131 lines
3.5 KiB
Go
131 lines
3.5 KiB
Go
package runner
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"strings"
|
|
)
|
|
|
|
type cacheAffinityKeys struct {
|
|
Primary string
|
|
Lookup []string
|
|
Record []string
|
|
}
|
|
|
|
func buildCacheAffinityKey(kind string, modelType string, body map[string]any) string {
|
|
return buildCacheAffinityKeys(kind, modelType, body).Primary
|
|
}
|
|
|
|
func buildCacheAffinityKeys(kind string, modelType string, body map[string]any) cacheAffinityKeys {
|
|
if len(body) == 0 {
|
|
return cacheAffinityKeys{}
|
|
}
|
|
for _, key := range []string{"cacheAffinityKey", "cache_affinity_key"} {
|
|
if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
|
|
key := "explicit:" + sha256Text(value)
|
|
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
|
|
}
|
|
}
|
|
for _, key := range []string{"sessionId", "session_id", "conversationId", "conversation_id"} {
|
|
if value := strings.TrimSpace(stringFromAny(body[key])); value != "" {
|
|
key := key + ":" + sha256Text(value)
|
|
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
|
|
}
|
|
}
|
|
if !cacheAffinityPromptHashSupported(kind, modelType) {
|
|
return cacheAffinityKeys{}
|
|
}
|
|
basePayload := map[string]any{
|
|
"kind": kind,
|
|
"modelType": modelType,
|
|
}
|
|
if tools, ok := body["tools"]; ok {
|
|
basePayload["tools"] = tools
|
|
}
|
|
if instructions, ok := body["instructions"]; ok {
|
|
basePayload["instructions"] = instructions
|
|
}
|
|
messageHashes := cacheAffinityMessageHashes(body)
|
|
if len(messageHashes) > 0 {
|
|
baseHash := stableJSONHash(basePayload)
|
|
keys := make([]string, 0, len(messageHashes))
|
|
for index := range messageHashes {
|
|
keyPayload := map[string]any{
|
|
"base": baseHash,
|
|
"messageHashes": messageHashes[:index+1],
|
|
"prefixLength": index + 1,
|
|
}
|
|
keys = append(keys, "prompt_lcp:"+stableJSONHash(keyPayload))
|
|
}
|
|
if len(keys) == 0 {
|
|
return cacheAffinityKeys{}
|
|
}
|
|
lookup := make([]string, 0, len(keys))
|
|
for index := len(keys) - 1; index >= 0; index-- {
|
|
lookup = append(lookup, keys[index])
|
|
}
|
|
return cacheAffinityKeys{
|
|
Primary: keys[len(keys)-1],
|
|
Lookup: lookup,
|
|
Record: keys,
|
|
}
|
|
}
|
|
if input, ok := body["input"]; ok {
|
|
basePayload["input"] = input
|
|
}
|
|
if len(basePayload) <= 2 {
|
|
return cacheAffinityKeys{}
|
|
}
|
|
key := "prompt:" + stableJSONHash(basePayload)
|
|
return cacheAffinityKeys{Primary: key, Lookup: []string{key}, Record: []string{key}}
|
|
}
|
|
|
|
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 cacheAffinityMessageHashes(body map[string]any) []string {
|
|
if hashes := stringListFromAny(body["messageHashes"]); len(hashes) > 0 {
|
|
return hashes
|
|
}
|
|
if hashes := stringListFromAny(body["message_hashes"]); len(hashes) > 0 {
|
|
return hashes
|
|
}
|
|
messages, ok := body["messages"].([]any)
|
|
if !ok || len(messages) == 0 {
|
|
return nil
|
|
}
|
|
hashes := make([]string, 0, len(messages))
|
|
for _, message := range messages {
|
|
raw, err := json.Marshal(message)
|
|
if err != nil || len(raw) == 0 {
|
|
continue
|
|
}
|
|
hashes = append(hashes, sha256Text(string(raw)))
|
|
}
|
|
return hashes
|
|
}
|
|
|
|
func stableJSONHash(value any) string {
|
|
raw, err := json.Marshal(value)
|
|
if err != nil || len(raw) == 0 {
|
|
return ""
|
|
}
|
|
return sha256Text(string(raw))
|
|
}
|
|
|
|
func sha256Text(value string) string {
|
|
sum := sha256.Sum256([]byte(value))
|
|
return hex.EncodeToString(sum[:])
|
|
}
|