Add cross-model global budget rate limiting with Valkey provider

config/config.yaml

@@ -1235,6 +1235,7 @@ global:
       providers:
         - type: redis
           address: redis:6379
+          db: 0
           domain: api
           rules:
             - name: premium-per-minute

src/semantic-router/pkg/cache/valkey_cache.go

@@ -322,15 +322,24 @@ func (c *ValkeyCache) AddPendingRequest(requestID string, model string, query st
 		return nil
 	}
 
-	err := c.addEntry("", requestID, model, query, requestBody, nil, ttlSeconds)
-
+	id := fmt.Sprintf("%x", md5.Sum(fmt.Appendf(nil, "%s_%s_%d", model, query, time.Now().UnixNano())))
+	err := c.addEntry(id, requestID, model, query, requestBody, nil, ttlSeconds)
 	if err != nil {
 		metrics.RecordCacheOperation("valkey", "add_pending", "error", time.Since(start).Seconds())
-	} else {
-		metrics.RecordCacheOperation("valkey", "add_pending", "success", time.Since(start).Seconds())
+		return err
 	}
 
-	return err
+	// Store reverse lookup: request_id -> doc key for direct retrieval in UpdateWithResponse
+	ctx := context.Background()
+	docKey := c.config.Index.Prefix + id
+	pendingKey := c.config.Index.Prefix + "pending:" + requestID
+	setCmd := []string{"SET", pendingKey, docKey, "EX", fmt.Sprintf("%d", ttlSeconds)}
+	if _, err := c.client.CustomCommand(ctx, setCmd); err != nil {
+		logging.Warnf("ValkeyCache.AddPendingRequest: failed to store pending lookup key: %v", err)
+	}
+
+	metrics.RecordCacheOperation("valkey", "add_pending", "success", time.Since(start).Seconds())
+	return nil
 }
 
 func (c *ValkeyCache) UpdateWithResponse(requestID string, responseBody []byte, ttlSeconds int) error {
@@ -345,34 +354,47 @@ func (c *ValkeyCache) UpdateWithResponse(requestID string, responseBody []byte,
 
 	ctx := context.Background()
 
-	query := fmt.Sprintf("@request_id:{%s}", escapeTagValue(requestID))
-	logging.Debugf("UpdateWithResponse: searching with TAG query: %s", query)
-
-	searchCmd := []string{
-		"FT.SEARCH", c.indexName, query,
-		"RETURN", "3", "model", "query", "request_body",
-		"LIMIT", "0", "1",
+	// Look up and atomically delete the reverse lookup key set in AddPendingRequest
+	pendingKey := c.config.Index.Prefix + "pending:" + requestID
+	docKeyResult, err := c.client.CustomCommand(ctx, []string{"GETDEL", pendingKey})
+	if err != nil || docKeyResult == nil {
+		logging.Infof("ValkeyCache.UpdateWithResponse: no pending lookup key for request_id=%s: %v", requestID, err)
+		metrics.RecordCacheOperation("valkey", "update_response", "error", time.Since(start).Seconds())
+		return fmt.Errorf("no pending entry found for request_id=%s", requestID)
 	}
 
-	results, err := c.client.CustomCommand(ctx, searchCmd)
+	docKey := fmt.Sprint(docKeyResult)
+	logging.Debugf("UpdateWithResponse: resolved pending key %s -> %s", pendingKey, docKey)
+
+	// Fetch the pending entry fields directly
+	hgetCmd := []string{"HMGET", docKey, "model", "query", "request_body"}
+	fieldsResult, err := c.client.CustomCommand(ctx, hgetCmd)
 	if err != nil {
-		logging.Infof("ValkeyCache.UpdateWithResponse: TAG search failed: %v", err)
+		logging.Infof("ValkeyCache.UpdateWithResponse: HMGET failed for %s: %v", docKey, err)
+		metrics.RecordCacheOperation("valkey", "update_response", "error", time.Since(start).Seconds())
+		return fmt.Errorf("failed to fetch pending entry: %w", err)
+	}
+
+	fields, ok := fieldsResult.([]interface{})
+	if !ok || len(fields) < 3 || fields[0] == nil {
 		metrics.RecordCacheOperation("valkey", "update_response", "error", time.Since(start).Seconds())
-		return fmt.Errorf("failed to search pending entry: %w", err)
+		return fmt.Errorf("pending entry not found or expired for request_id=%s", requestID)
 	}
 
-	logging.Debugf("UpdateWithResponse: search results type=%T, value=%v", results, results)
+	model := fmt.Sprint(fields[0])
+	query := fmt.Sprint(fields[1])
+	requestBodyStr := fmt.Sprint(fields[2])
 
-	entry, err := parsePendingSearchResult(results, requestID, c.config.Index.Prefix)
-	if err != nil {
+	if model == "" || query == "" {
+		logging.Debugf("ValkeyCache.UpdateWithResponse: missing required fields (key: %s, model: %q, query_len: %d)", docKey, model, len(query))
 		metrics.RecordCacheOperation("valkey", "update_response", "error", time.Since(start).Seconds())
-		return err
+		return fmt.Errorf("incomplete pending entry for request_id=%s: missing required fields", requestID)
 	}
 
-	logging.Debugf("ValkeyCache.UpdateWithResponse: found pending entry, updating (id: %s, model: %s)", entry.docID, entry.model)
+	logging.Debugf("ValkeyCache.UpdateWithResponse: found pending entry, updating (key: %s, model: %s)", docKey, model)
 
 	// Update the document with response body and TTL
-	err = c.addEntry(entry.docID, requestID, entry.model, entry.query, []byte(entry.requestBodyStr), responseBody, ttlSeconds)
+	err = c.addEntry(docKey, requestID, model, query, []byte(requestBodyStr), responseBody, ttlSeconds)
 	if err != nil {
 		metrics.RecordCacheOperation("valkey", "update_response", "error", time.Since(start).Seconds())
 		return fmt.Errorf("failed to update entry: %w", err)

src/semantic-router/pkg/cache/valkey_cache_helpers.go

@@ -7,93 +7,6 @@ import (
 	"github.com/vllm-project/semantic-router/src/semantic-router/pkg/observability/logging"
 )
 
-// pendingEntry holds the parsed fields from a pending cache entry search result.
-type pendingEntry struct {
-	docID          string
-	model          string
-	query          string
-	requestBodyStr string
-}
-
-// extractPendingFields extracts a pendingEntry from a Valkey GLIDE doc map (the second element of FT.SEARCH results).
-func extractPendingFields(docMap map[string]interface{}) (*pendingEntry, error) {
-	entry := &pendingEntry{}
-	for docKey, docValue := range docMap {
-		entry.docID = docKey
-
-		fieldsMap, mapOk := docValue.(map[string]interface{})
-		if !mapOk {
-			logging.Warnf("UpdateWithResponse: document fields is not a map, type=%T", docValue)
-			return nil, fmt.Errorf("invalid search result: expected fields map")
-		}
-
-		if v, exists := fieldsMap["model"]; exists {
-			entry.model = fmt.Sprint(v)
-		}
-		if v, exists := fieldsMap["query"]; exists {
-			entry.query = fmt.Sprint(v)
-		}
-		if v, exists := fieldsMap["request_body"]; exists {
-			entry.requestBodyStr = fmt.Sprint(v)
-		}
-
-		break // Only process the first document
-	}
-	return entry, nil
-}
-
-// parsePendingSearchResult extracts a pendingEntry from a Valkey FT.SEARCH result.
-// Returns the entry or an error describing the parse failure.
-func parsePendingSearchResult(results interface{}, requestID string, prefix string) (*pendingEntry, error) {
-	resultsArray, ok := results.([]interface{})
-	if !ok || len(resultsArray) < 1 {
-		logging.Infof("ValkeyCache.UpdateWithResponse: invalid result format for request_id=%s", requestID)
-		return nil, fmt.Errorf("invalid search result format")
-	}
-
-	totalResults, ok := resultsArray[0].(int64)
-	if !ok {
-		logging.Infof("ValkeyCache.UpdateWithResponse: invalid count type for request_id=%s (got %T)", requestID, resultsArray[0])
-		return nil, fmt.Errorf("invalid search result count type")
-	}
-
-	if totalResults == 0 {
-		logging.Infof("ValkeyCache.UpdateWithResponse: no pending entry found with request_id=%s (count=0, may still be indexing)", requestID)
-		return nil, fmt.Errorf("no pending entry found (indexing may still be in progress)")
-	}
-
-	logging.Infof("UpdateWithResponse: found %d result(s) for request_id=%s", totalResults, requestID)
-
-	if len(resultsArray) < 2 {
-		logging.Warnf("UpdateWithResponse: resultsArray only has %d elements", len(resultsArray))
-		return nil, fmt.Errorf("invalid search result: expected at least 2 elements")
-	}
-
-	docMap, ok := resultsArray[1].(map[string]interface{})
-	if !ok {
-		logging.Warnf("UpdateWithResponse: resultsArray[1] is not a map, type=%T", resultsArray[1])
-		return nil, fmt.Errorf("invalid search result: expected map at index 1")
-	}
-
-	entry, err := extractPendingFields(docMap)
-	if err != nil {
-		return nil, err
-	}
-
-	if !strings.HasPrefix(entry.docID, prefix) {
-		logging.Warnf("UpdateWithResponse: docID '%s' doesn't have expected prefix '%s'", entry.docID, prefix)
-	}
-
-	logging.Debugf("UpdateWithResponse: extracted docID='%s', model='%s', query='%s'", entry.docID, entry.model, entry.query)
-
-	if entry.model == "" || entry.query == "" {
-		logging.Warnf("UpdateWithResponse: missing required fields (model='%s', query='%s')", entry.model, entry.query)
-		return nil, fmt.Errorf("missing required fields in pending entry")
-	}
-
-	return entry, nil
-}
-
 // searchMatch holds the parsed fields from a vector search result.
 type searchMatch struct {
 	distance     float64

src/semantic-router/pkg/cache/valkey_cache_integration_test.go

@@ -223,64 +223,69 @@ func TestValkeyCacheIntegration_AddPendingRequest(t *testing.T) {
 	time.Sleep(100 * time.Millisecond)
 }
 
-func TestValkeyCacheIntegration_UpdateWithResponse(t *testing.T) {
-	cache := setupValkeyCacheIntegration(t)
-	defer func() { _ = cache.Close() }()
-
-	requestID := "req_update_456"
-	model := "gpt-4"
-	query := "What is Python?"
-	requestBody := []byte(`{"model":"gpt-4"}`)
-	ttlSeconds := 300
-
-	err := cache.AddPendingRequest(requestID, model, query, requestBody, ttlSeconds)
-	require.NoError(t, err)
-
-	// Wait for indexing with retry logic
-	time.Sleep(200 * time.Millisecond)
-
-	responseBody := []byte(`{"choices":[{"message":{"content":"Python is a programming language."}}]}`)
-
-	// Retry UpdateWithResponse to handle indexing delays
+// retryUpdateWithResponse retries UpdateWithResponse up to maxRetries times with backoff.
+func retryUpdateWithResponse(t *testing.T, cache *ValkeyCache, requestID string, responseBody []byte, ttlSeconds int) {
+	t.Helper()
 	var updateErr error
-	maxRetries := 5
-	for i := 0; i < maxRetries; i++ {
+	for i := 0; i < 5; i++ {
 		if i > 0 {
-			t.Logf("Retry %d/%d for UpdateWithResponse", i+1, maxRetries)
+			t.Logf("Retry %d/5 for UpdateWithResponse", i+1)
 			time.Sleep(time.Duration(200*(i+1)) * time.Millisecond)
 		}
-
 		updateErr = cache.UpdateWithResponse(requestID, responseBody, ttlSeconds)
 		if updateErr == nil {
-			break
+			return
 		}
 		t.Logf("UpdateWithResponse attempt %d failed: %v", i+1, updateErr)
 	}
-
 	require.NoError(t, updateErr, "UpdateWithResponse should succeed after retries")
+}
 
-	// Wait longer for vector index to update with the new response and retry search
-	var foundResponse []byte
-	var hit bool
-	searchRetries := 5
-
-	for i := 0; i < searchRetries; i++ {
+// retryFindSimilar retries FindSimilar until a hit is found or retries are exhausted.
+func retryFindSimilar(t *testing.T, cache *ValkeyCache, model, query string) ([]byte, bool) {
+	t.Helper()
+	for i := 0; i < 5; i++ {
 		if i > 0 {
-			t.Logf("Search retry %d/%d for updated entry", i+1, searchRetries)
+			t.Logf("Search retry %d/5 for updated entry", i+1)
 		}
 		time.Sleep(time.Duration(200*(i+1)) * time.Millisecond)
-
-		foundResponse, hit, err = cache.FindSimilar(model, query)
+		foundResponse, hit, err := cache.FindSimilar(model, query)
 		require.NoError(t, err)
-
 		if hit {
-			break
+			return foundResponse, true
 		}
 	}
+	return nil, false
+}
+
+func TestValkeyCacheIntegration_UpdateWithResponse(t *testing.T) {
+	cache := setupValkeyCacheIntegration(t)
+	defer func() { _ = cache.Close() }()
+
+	requestID := "req_update_456"
+	model := "gpt-4"
+	query := "What is Python?"
+	requestBody := []byte(`{"model":"gpt-4"}`)
+	ttlSeconds := 300
 
+	err := cache.AddPendingRequest(requestID, model, query, requestBody, ttlSeconds)
+	require.NoError(t, err)
+
+	time.Sleep(200 * time.Millisecond)
+
+	responseBody := []byte(`{"choices":[{"message":{"content":"Python is a programming language."}}]}`)
+	retryUpdateWithResponse(t, cache, requestID, responseBody, ttlSeconds)
+
+	foundResponse, hit := retryFindSimilar(t, cache, model, query)
 	require.True(t, hit, "Should find the updated entry after retries")
 	assert.NotNil(t, foundResponse)
 	assert.Contains(t, string(foundResponse), "Python", "Updated response should be findable")
+
+	// Verify the pending lookup key was cleaned up
+	pendingKey := cache.config.Index.Prefix + "pending:" + requestID
+	result, err := cache.client.CustomCommand(context.Background(), []string{"GET", pendingKey})
+	assert.NoError(t, err)
+	assert.Nil(t, result, "Pending lookup key should be deleted after UpdateWithResponse")
 }
 
 func TestValkeyCacheIntegration_UpdateWithResponseSpecialChars(t *testing.T) {
@@ -334,6 +339,15 @@ func TestValkeyCacheIntegration_UpdateWithResponseSpecialChars(t *testing.T) {
 	}
 }
 
+func TestValkeyCacheIntegration_UpdateWithResponseUnknownRequestID(t *testing.T) {
+	cache := setupValkeyCacheIntegration(t)
+	defer func() { _ = cache.Close() }()
+
+	err := cache.UpdateWithResponse("nonexistent-request-id", []byte(`{"response":"test"}`), 300)
+	assert.Error(t, err, "UpdateWithResponse should fail for unknown request ID")
+	assert.Contains(t, err.Error(), "no pending entry found")
+}
+
 func TestValkeyCacheIntegration_TTLExpiration(t *testing.T) {
 	cache := setupValkeyCacheIntegration(t)
 	defer func() { _ = cache.Close() }()

src/semantic-router/pkg/config/config.go

@@ -118,6 +118,7 @@ type RateLimitConfig struct {
 type RateLimitProviderConfig struct {
 	Type    string          `yaml:"type"`
 	Address string          `yaml:"address,omitempty"`
+	DB      int             `yaml:"db,omitempty"`
 	Domain  string          `yaml:"domain,omitempty"`
 	Rules   []RateLimitRule `yaml:"rules,omitempty"`
 }