internal/responsecache/semantic.go
Rama main · réplica de lectura del repositorio original.
package responsecache
import (
"bytes"
"context"
"crypto/sha256"
"encoding/hex"
"fmt"
"hash"
"log/slog"
"net/http"
"slices"
"sort"
"strconv"
"strings"
"sync"
"time"
"github.com/goccy/go-json"
"github.com/enterpilot/gomodel/config"
"github.com/enterpilot/gomodel/internal/core"
"github.com/enterpilot/gomodel/internal/embedding"
)
// semanticCacheWriteJob carries one vector-store insert handed to a background
// worker. Insert arguments are captured by value so the request goroutine can
// return immediately.
type semanticCacheWriteJob struct {
cacheKey string
vec []float32
data []byte
paramsHash string
ttl time.Duration
}
// SemanticCacheMiddleware implements the vector-similarity response cache layer.
// It is the second cache layer, consulted only after an exact-match miss.
type semanticCacheMiddleware struct {
embedder embedding.Embedder
store VecStore
cfg config.SemanticCacheConfig
embedderIdentity string
hitRecorder func(exchange, []byte, string)
partitionByUserPath bool
// Vector-store inserts run on a bounded worker pool so a burst of cache
// misses cannot spawn unbounded goroutines against the vector store. wg
// tracks queued jobs; workers drains them; mu guards closed/close(jobs).
jobs chan semanticCacheWriteJob
wg sync.WaitGroup
workers sync.WaitGroup
mu sync.RWMutex
closed bool
}
func newSemanticCacheMiddleware(emb embedding.Embedder, store VecStore, cfg config.SemanticCacheConfig, hitRecorder func(exchange, []byte, string)) *semanticCacheMiddleware {
m := &semanticCacheMiddleware{
embedder: emb,
store: store,
cfg: cfg,
embedderIdentity: emb.Identity(),
hitRecorder: hitRecorder,
jobs: make(chan semanticCacheWriteJob, cacheWriteQueueSize),
}
m.startWorkers()
return m
}
func (m *semanticCacheMiddleware) startWorkers() {
for range cacheWriteWorkerCount {
m.workers.Go(func() {
for job := range m.jobs {
storeCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
err := m.store.Insert(storeCtx, job.cacheKey, job.vec, job.data, job.paramsHash, job.ttl)
cancel()
if err != nil {
slog.Warn("semantic cache: store failed", "key", job.cacheKey, "err", err)
}
m.wg.Done()
}
})
}
}
// enqueueWrite hands a store insert to the worker pool. It mirrors
// simpleCacheMiddleware.enqueueWrite: the read lock spans Add+send so close()
// cannot race an untracked write, and a full queue rolls back the Add rather
// than blocking the request goroutine.
func (m *semanticCacheMiddleware) enqueueWrite(job semanticCacheWriteJob) {
m.mu.RLock()
if m.closed {
m.mu.RUnlock()
return
}
m.wg.Add(1)
select {
case m.jobs <- job:
m.mu.RUnlock()
default:
m.wg.Done()
m.mu.RUnlock()
slog.Warn("semantic cache write queue full", "key", job.cacheKey)
}
}
// Handle processes a single request/response cycle for semantic caching.
// It must be called after guardrail patching so params_hash reflects the final policy.
func (m *semanticCacheMiddleware) Handle(ex exchange, body []byte, next func() error) error {
if m == nil || m.store == nil {
return next()
}
path := ex.Path()
if !semanticCacheablePath(path) || ex.Method() != http.MethodPost {
return next()
}
if shouldSkipSemanticCacheHeaders(ex.RequestHeader) {
return next()
}
ctx := ex.Context()
plan := core.GetWorkflow(ctx)
embedText, msgCount := extractEmbedText(body, m.cfg.ExcludeSystemPrompt)
if embedText == "" {
return next()
}
// The header may only tighten matching. Entries are shared across callers
// (across every caller of one user path when partitioning is on), so a looser
// client threshold would return answers cached for unrelated prompts,
// including other tenants'.
threshold := m.cfg.SimilarityThreshold
if v := headerFloat64(ex.RequestHeader, "X-Cache-Semantic-Threshold"); v > threshold {
threshold = v
}
if m.cfg.MaxConversationMessages != nil && *m.cfg.MaxConversationMessages > 0 && msgCount > *m.cfg.MaxConversationMessages {
return next()
}
msgFp, fpOK := conversationInvariantFingerprint(body, m.cfg.ExcludeSystemPrompt)
if !fpOK {
return next()
}
baseParams := computeParamsHash(body, path, plan, core.GetGuardrailsHash(ctx), m.embedderIdentity, m.cfg.ExcludeSystemPrompt)
paramsHash := sha256HexOf(baseParams + "\x00" + msgFp)
if m.partitionByUserPath {
// Every VecStore backend filters Search by exact params_hash inside the
// vector query, so this confines matching to the caller's partition. The
// cache key below derives from it too, so partitions cannot overwrite
// each other's points.
paramsHash = userPathPartitionedHash(paramsHash, core.UserPathFromContext(ctx))
}
vec, err := m.embedder.Embed(ctx, embedText)
if err != nil {
slog.Warn("semantic cache: embed failed, bypassing", "err", err)
return next()
}
results, err := m.store.Search(ctx, vec, paramsHash, 1)
if err != nil {
slog.Warn("semantic cache: search failed, bypassing", "err", err)
return next()
}
if len(results) > 0 && float64(results[0].Score) >= threshold {
replayErr := ex.ReplayHit(body, results[0].Response, CacheTypeSemantic)
if replayErr == nil || replayCommitted(replayErr) {
ex.MarkHit(CacheTypeSemantic)
if m.hitRecorder != nil {
m.hitRecorder(ex, results[0].Response, CacheTypeSemantic)
}
slog.Info("semantic cache hit",
"path", path,
"score", results[0].Score,
"request_id", core.GetRequestID(ex.Context()),
)
return replayErr
}
slog.Warn("semantic cache replay failed", "path", path, "err", replayErr)
}
data, ok, err := ex.Capture("semantic cache: failed to capture cacheable response body", next)
if err != nil {
return err
}
if !ok {
return nil
}
ttlSec := 0
if m.cfg.TTL != nil {
ttlSec = *m.cfg.TTL
}
ttl := time.Duration(ttlSec) * time.Second
if v := headerDuration(ex.RequestHeader, "X-Cache-TTL"); v > 0 {
ttl = v
}
cacheKey := sha256HexOf(embedText + "\x00" + paramsHash)
m.enqueueWrite(semanticCacheWriteJob{
cacheKey: cacheKey,
vec: vec,
data: data,
paramsHash: paramsHash,
ttl: ttl,
})
return nil
}
func (m *semanticCacheMiddleware) close() error {
m.mu.Lock()
if !m.closed {
m.closed = true
close(m.jobs)
}
m.mu.Unlock()
m.workers.Wait()
m.wg.Wait()
if m.embedder != nil {
_ = m.embedder.Close() //nolint:errcheck
}
if m.store != nil {
return m.store.Close()
}
return nil
}
type embedMessage struct {
Role string `json:"role"`
Content any `json:"content"`
}
// extractEmbedText returns the text to embed and the total non-system message count.
// When excludeSystem is true, system messages are stripped from counting and embedding.
// Only the last user message is used as the embedding text to maximize cache hit rate.
// Supports chat bodies with "messages" and Responses API bodies with "input" as either a
// string or an array of {role, content} items (OpenAI-style).
func extractEmbedText(body []byte, excludeSystem bool) (text string, nonSystemCount int) {
var envelope struct {
Messages []embedMessage `json:"messages"`
Input json.RawMessage `json:"input"`
}
if err := json.Unmarshal(body, &envelope); err != nil {
return "", 0
}
if len(envelope.Messages) > 0 {
return embedTextFromMessages(envelope.Messages, excludeSystem)
}
if len(envelope.Input) == 0 {
return "", 0
}
var s string
if json.Unmarshal(envelope.Input, &s) == nil {
if s == "" {
return "", 0
}
return s, 1
}
var inputMsgs []embedMessage
if json.Unmarshal(envelope.Input, &inputMsgs) != nil || len(inputMsgs) == 0 {
return "", 0
}
return embedTextFromMessages(inputMsgs, excludeSystem)
}
func embedTextFromMessages(messages []embedMessage, excludeSystem bool) (text string, nonSystemCount int) {
var lastUserText string
for _, m := range messages {
if m.Role == "system" && excludeSystem {
continue
}
nonSystemCount++
if m.Role == "user" {
lastUserText = extractTextFromContent(m.Content)
}
}
return lastUserText, nonSystemCount
}
// conversationInvariantFingerprint hashes structural cache context: every message's
// role and raw content except the last user turn, where only non-text parts (e.g.
// image_url) are included so paraphrases of the final user text share a namespace.
// For Responses API, "input" may be a string (empty fingerprint) or a message array
// (same treatment as "messages"). ok is false if the JSON envelope cannot be parsed or
// there is no usable messages/input payload.
func conversationInvariantFingerprint(body []byte, excludeSystem bool) (fingerprint string, ok bool) {
var envelope struct {
Messages []json.RawMessage `json:"messages"`
Input json.RawMessage `json:"input"`
}
if err := json.Unmarshal(body, &envelope); err != nil {
return "", false
}
msgs, fpOK := messageRawListFromMessagesOrInput(envelope.Messages, envelope.Input)
if !fpOK {
return "", false
}
if len(msgs) == 0 {
return "", true
}
type msgPart struct {
role string
content json.RawMessage
unparseable bool
rawMsg json.RawMessage
}
var included []msgPart
for _, rawMsg := range msgs {
var m struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
}
if err := json.Unmarshal(rawMsg, &m); err != nil {
var roleOnly struct {
Role string `json:"role"`
}
_ = json.Unmarshal(rawMsg, &roleOnly)
if roleOnly.Role == "system" && excludeSystem {
continue
}
included = append(included, msgPart{role: roleOnly.Role, unparseable: true, rawMsg: rawMsg})
continue
}
if m.Role == "system" && excludeSystem {
continue
}
included = append(included, msgPart{role: m.Role, content: m.Content})
}
lastUser := -1
for i, part := range slices.Backward(included) {
if part.role == "user" {
lastUser = i
break
}
}
h := sha256.New()
for i, p := range included {
h.Write([]byte(p.role))
h.Write([]byte{0})
if p.unparseable {
sum := sha256.Sum256(p.rawMsg)
h.Write(sum[:])
h.Write([]byte{0})
continue
}
if i == lastUser && lastUser >= 0 {
writeNonTextUserContentFingerprint(h, p.content)
} else if len(p.content) > 0 {
h.Write(p.content)
}
h.Write([]byte{0})
}
return hex.EncodeToString(h.Sum(nil)), true
}
func messageRawListFromMessagesOrInput(messages []json.RawMessage, input json.RawMessage) (msgs []json.RawMessage, ok bool) {
if len(messages) > 0 {
return messages, true
}
if len(input) == 0 {
return nil, false
}
var s string
if json.Unmarshal(input, &s) == nil {
return nil, true
}
var arr []json.RawMessage
if json.Unmarshal(input, &arr) != nil {
return nil, false
}
if len(arr) == 0 {
return nil, false
}
return arr, true
}
func writeNonTextUserContentFingerprint(h hash.Hash, content json.RawMessage) {
if len(bytes.TrimSpace(content)) == 0 {
return
}
var s string
if json.Unmarshal(content, &s) == nil {
return
}
var parts []json.RawMessage
if json.Unmarshal(content, &parts) != nil {
_, _ = h.Write(content)
return
}
for _, p := range parts {
var obj map[string]json.RawMessage
if json.Unmarshal(p, &obj) != nil {
_, _ = h.Write(p)
_, _ = h.Write([]byte{0})
continue
}
tBytes, hasType := obj["type"]
if !hasType {
_, _ = h.Write(p)
_, _ = h.Write([]byte{0})
continue
}
var typeStr string
_ = json.Unmarshal(tBytes, &typeStr)
if typeStr == "text" {
continue
}
_, _ = h.Write(p)
_, _ = h.Write([]byte{0})
}
}
func extractTextFromContent(content any) string {
switch v := content.(type) {
case string:
return v
case []any:
var parts []string
for _, item := range v {
if m, ok := item.(map[string]any); ok {
if t, ok := m["type"].(string); ok && t == "text" {
if text, ok := m["text"].(string); ok {
parts = append(parts, text)
}
}
}
}
return strings.Join(parts, " ")
default:
return fmt.Sprintf("%v", content)
}
}
// paramsHashPromptFields are top-level request fields never hashed as plain
// fields: messages and input are the prompt, matched by embedding and by
// conversationInvariantFingerprint, and stream and stream_options are hashed
// in normalized form by writeOutputShapingFields. Every field not listed here
// or in paramsHashCallerFields is hashed, so a field GoModel does not know
// about causes a miss, never a hit generated under a different value.
var paramsHashPromptFields = map[string]struct{}{
"messages": {},
"input": {},
"stream": {},
"stream_options": {},
}
// paramsHashCallerFields identify the caller or steer provider-side prompt
// caching. They never change what the model generates, so they are left out
// of the hash, except on responsesPath: a Responses object echoes them back,
// and a hit must not hand one caller's values to another.
var paramsHashCallerFields = map[string]struct{}{
"user": {},
"metadata": {},
"safety_identifier": {},
"prompt_cache_key": {},
"prompt_cache_retention": {},
}
const responsesPath = "/v1/responses"
// computeParamsHash builds a stable SHA-256 hash of all output-shaping parameters.
// This ensures semantically similar prompts with different parameters or guardrail
// policies never share a cache entry. endpointPath is the raw URL path
// (e.g. "/v1/chat/completions") and isolates entries across distinct endpoints.
// excludeSystem leaves out a top-level system prompt (Anthropic Messages), as
// conversationInvariantFingerprint leaves out system messages.
func computeParamsHash(body []byte, endpointPath string, plan *core.Workflow, guardrailsHash, embedderIdentity string, excludeSystem bool) string {
h := sha256.New()
h.Write([]byte(endpointPath))
h.Write([]byte{0})
if plan != nil {
h.Write([]byte(plan.ProviderType))
h.Write([]byte{0})
h.Write([]byte(plan.ResolvedQualifiedModel()))
h.Write([]byte{0})
}
writeOutputShapingFields(h, body, endpointPath, excludeSystem)
h.Write([]byte(guardrailsHash))
h.Write([]byte{0})
h.Write([]byte(embedderIdentity))
return hex.EncodeToString(h.Sum(nil))
}
// writeOutputShapingFields hashes every top-level field of body that is not a
// prompt field or, outside responsesPath, a caller field, in sorted key order
// with canonical JSON values, so member order and whitespace do not matter.
// Keys are written JSON-encoded, so a key containing a NUL byte cannot pass
// for the separator. A body that is not a JSON object, or that repeats a member
// name, is hashed verbatim instead: a provider may read a duplicated member
// differently from the value decoding keeps.
func writeOutputShapingFields(h hash.Hash, body []byte, endpointPath string, excludeSystem bool) {
var fields map[string]json.RawMessage
if hasDuplicateJSONMemberNames(body) || json.Unmarshal(body, &fields) != nil {
h.Write(body)
h.Write([]byte{0})
return
}
// Embedding and the conversation fingerprint read messages when a body has
// both messages and input, while a provider may read input. Hashing both
// keeps such a body from sharing entries with one whose other field differs.
_, hasMessages := fields["messages"]
_, hasInput := fields["input"]
bothPrompts := hasMessages && hasInput
keys := make([]string, 0, len(fields))
for key := range fields {
_, prompt := paramsHashPromptFields[key]
if bothPrompts && (key == "messages" || key == "input") {
prompt = false
}
if prompt {
continue
}
if _, caller := paramsHashCallerFields[key]; caller && endpointPath != responsesPath {
continue
}
if key == "system" && excludeSystem {
continue
}
keys = append(keys, key)
}
sort.Strings(keys)
for _, key := range keys {
encodedKey, _ := json.Marshal(key)
h.Write(encodedKey)
h.Write([]byte{0})
if key == "tools" {
h.Write(canonicalToolsJSON(fields[key]))
} else {
h.Write(canonicalJSON(fields[key]))
}
h.Write([]byte{0})
}
var stream struct {
Stream bool `json:"stream"`
StreamOptions *core.StreamOptions `json:"stream_options"`
}
_ = json.Unmarshal(body, &stream)
if stream.Stream {
h.Write([]byte("stream"))
h.Write([]byte{0})
if streamOptions := normalizeStreamOptionsForCache(stream.StreamOptions); streamOptions != nil {
soJSON, _ := json.Marshal(streamOptions)
h.Write(soJSON)
}
h.Write([]byte{0})
}
}
// canonicalJSON re-encodes raw with sorted object keys and no insignificant
// whitespace. Numbers keep their literal text, so large integers such as a
// seed are never rounded into a colliding value.
func canonicalJSON(raw json.RawMessage) []byte {
dec := json.NewDecoder(bytes.NewReader(raw))
dec.UseNumber()
var v any
if err := dec.Decode(&v); err != nil {
return raw
}
out, err := json.Marshal(v)
if err != nil {
return raw
}
return out
}
// canonicalToolsJSON is canonicalJSON for a tools array, ordered by function
// name so the same tool set in a different order shares entries.
func canonicalToolsJSON(raw json.RawMessage) []byte {
dec := json.NewDecoder(bytes.NewReader(raw))
dec.UseNumber()
var tools []map[string]any
if err := dec.Decode(&tools); err != nil {
return canonicalJSON(raw)
}
out, err := json.Marshal(sortedTools(tools))
if err != nil {
return canonicalJSON(raw)
}
return out
}
// sortedTools orders a list of Chat Completions function tools by name. Any
// other list, one that mixes in another tool shape or repeats a name, keeps its
// order: there is no single order to normalize it to.
func sortedTools(tools []map[string]any) []map[string]any {
names := make([]string, len(tools))
seen := make(map[string]struct{}, len(tools))
for i, tool := range tools {
function, _ := tool["function"].(map[string]any)
name, _ := function["name"].(string)
if _, duplicate := seen[name]; name == "" || duplicate {
return tools
}
seen[name] = struct{}{}
names[i] = name
}
order := make([]int, len(tools))
for i := range order {
order[i] = i
}
sort.Slice(order, func(a, b int) bool { return names[order[a]] < names[order[b]] })
sorted := make([]map[string]any, len(tools))
for i, idx := range order {
sorted[i] = tools[idx]
}
return sorted
}
func shouldSkipSemanticCacheHeaders(header func(string) string) bool {
if shouldSkipCacheControl(header("Cache-Control")) {
return true
}
return strings.EqualFold(header("X-Cache-Type"), "exact")
}
func headerFloat64(header func(string) string, name string) float64 {
s := header(name)
if s == "" {
return 0
}
v, err := strconv.ParseFloat(s, 64)
if err != nil {
return 0
}
return v
}
func headerDuration(header func(string) string, name string) time.Duration {
s := header(name)
if s == "" {
return 0
}
seconds, err := strconv.ParseInt(s, 10, 64)
if err != nil {
return 0
}
return time.Duration(seconds) * time.Second
}
func sha256HexOf(s string) string {
h := sha256.Sum256([]byte(s))
return hex.EncodeToString(h[:])
}
// CacheTypeHeader values for X-Cache-Type.
const (
CacheTypeExact = "exact"
CacheTypeSemantic = "semantic"
CacheHeaderExact = "HIT (exact)"
CacheHeaderSemantic = "HIT (semantic)"
)
// shouldSkipAllCacheHeaders reports whether caching must be bypassed for this
// request, matching the exact-cache middleware semantics for no-cache and no-store.
func shouldSkipAllCacheHeaders(header func(string) string) bool {
if strings.EqualFold(header("X-Cache-Control"), "no-store") {
return true
}
return shouldSkipCacheControl(header("Cache-Control"))
}