semantic_cache.hpp

Go to the documentation of this file.

#pragma once
 
#include <continuum/runtime/hit_verifier.hpp>
 
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace continuum::runtime {
 
struct EmbeddingProvider {
  virtual ~EmbeddingProvider() = default;
  virtual std::vector<float> embed(const std::string& text) const = 0;
  virtual std::size_t dimension() const = 0;
  virtual std::string identity() const = 0;
};
 
struct SemanticCacheEntry {
  std::vector<float> embedding;
  std::vector<std::uint8_t> cached_output;
  std::string model_id;
  std::int64_t last_access_ns = 0;
  std::string cache_namespace;
  std::string embedder_id;
  std::string prompt;
};
 
class SemanticCacheIndex {
 public:
  explicit SemanticCacheIndex(std::size_t max_entries = 2048,
                              float similarity_threshold = 0.85f);
 
  struct LookupResult {
    std::vector<std::uint8_t> output;
    float similarity = 0.0f;
    bool above_threshold = false;  
    std::string prompt;            
    std::int32_t verifier_rejections = 0;  
  };
 
  LookupResult lookup(const std::vector<float>& query_embedding,
                      const std::string& model_id,
                      const std::string& cache_namespace = {},
                      const std::string& embedder_id = {},
                      const std::string& query_prompt = {}) const;
 
  void insert(const std::vector<float>& embedding,
              const std::string& model_id,
              std::vector<std::uint8_t> output,
              const std::string& cache_namespace = {},
              const std::string& embedder_id = {},
              const std::string& prompt = {});
 
  void set_verifier(std::shared_ptr<const HitVerifier> verifier);
  std::shared_ptr<const HitVerifier> verifier() const;
  std::int64_t verifier_rejections() const;
 
  void clear();
  std::size_t size() const;
  std::size_t max_entries() const { return max_entries_; }
  std::size_t estimated_bytes() const;
  float similarity_threshold() const;
  void set_similarity_threshold(float t);
 
  static float cosine_similarity(const std::vector<float>& a,
                                  const std::vector<float>& b);
 
 private:
  mutable std::mutex mu_;
  // Mutable so a const lookup can refresh LRU recency on a hit.
  mutable std::vector<SemanticCacheEntry> entries_;
  std::size_t max_entries_;
  float similarity_threshold_;
  mutable std::uint64_t clock_ = 0;
  std::shared_ptr<const HitVerifier> verifier_ = std::make_shared<LexicalNearMissVerifier>();
  mutable std::int64_t verifier_rejections_ = 0;
  // Verdicts per (cached prompt, query prompt): a Session's metrics pass and
  // its interpreter look up the same prompt, and a verifier may be costly.
  mutable std::unordered_map<std::string, bool> verdicts_;
  static constexpr std::size_t kMaxVerdicts = 4096;
};
 
class BruteForceEmbeddingProvider : public EmbeddingProvider {
 public:
  explicit BruteForceEmbeddingProvider(std::size_t dim = 64);
 
  std::vector<float> embed(const std::string& text) const override;
  std::size_t dimension() const override;
  std::string identity() const override;
 
 private:
  std::size_t dim_;
};
 
}  // namespace continuum::runtime