llm_optimizer 0.1.5 → 0.1.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b5f6d0b99af3e0801e77df0316ac767e0e10e0d4e7bba9dc19623797681a2961
-  data.tar.gz: d8644df814cb0c7f219a51620d3cd409e1bb5822228278245b2572dbaf666fdc
+  metadata.gz: 3a0ec4bdfa750f16155927a3e00c9fe2c1c39da7e85866eb6c65855ac6eebaef
+  data.tar.gz: 0e5820f0503fbef14dc1ad858dfaa7527e3dba278fbf7640df377d82fbc61ad7
 SHA512:
-  metadata.gz: 1396d95f7e3f498e600cf6e3b99627ee2f746692a1f002be989ce1b13859f5a1af8f50656a82ff0fa853d3e42b0c219de49ac152baa2107d9c9529fc82bf63e4
-  data.tar.gz: 6b45bae664e4d43fd54fe47c2c8c9ebdeea2b4442f78bf22daee56d730095a471f58faf0e6432c7e5ef18a58cfc20a9279dc8706c9f41e4ca55db0cb441df1e8
+  metadata.gz: 8c2f376e324a7678063e66a89b6ad89e476bd699fd3a816c7c91a79b16ba40e09111cfdfacb1206946e2d111122e63cf70babc09a0467821723b2b286eda235a
+  data.tar.gz: 5bba8c343627f230c13f0671cd8b1374ab0405f6c6369457b92e9093ac1cd2f780797797a26fabdc865981ca1e131b6dc80ae4a97342f3de2f3297255d8e13c9

data/CHANGELOG.md

@@ -7,6 +7,22 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.1.6] - 2026-05-04
+
+### Added
+- `with_tools` configuration option (aliased as `tools`) — allows passing function/tool definitions to LLM calls via the `optimize` method
+- Tool support for both `llm_caller` and `messages_caller` — `tools:` keyword argument is now passed to all underlying LLM callers
+- `with_tools` examples in the README and Rails initializer template
+- `cache_scope` configuration option — isolates semantic cache entries into separate namespaces; useful for ensuring cache hits only occur within specific contexts (e.g., user IDs, account types, or dynamic categories)
+
+### Changed
+- `Pipeline#raw_llm_call` refactored to handle global and per-call tools consistently
+- Refactored `Pipeline` to remove duplicate internal method definitions (`semantic_cache_lookup`, `store_in_cache`)
+- `SemanticCache#lookup` return format updated to `[response, token_info]` to support better metadata tracking
+
+### Fixed
+- RuboCop `Metrics/ParameterLists` offense in `OptimizeResult#initialize` by adding targeted override for the necessary result fields
+
 ## [0.1.5] - 2026-04-22
 
 ### Added

data/README.md

@@ -200,6 +200,7 @@ end
 | `messages_caller` | Lambda | `nil` | `(messages, model:) -> String` — used when `conversation_id` is present; receives full history including current user turn |
 | `system_prompt` | String | `nil` | Seeded as the first system message when a new conversation is created via `conversation_id` |
 | `conversation_ttl` | Integer | `86400` | TTL in seconds for Redis-backed conversation history (`0` for no expiry) |
+| `with_tools` | Array | `nil` | Tools (functions) available to the LLM; passed as `tools:` keyword to callers |
 
 ## Per-call configuration
 
@@ -212,19 +213,6 @@ result = LlmOptimizer.optimize(prompt) do |config|
 end
 ```
 
-## Conversation history
-
-Pass a `messages` array to enable history management:
-
-```ruby
-messages = [
-  { role: "user",      content: "Tell me about Redis" },
-  { role: "assistant", content: "Redis is an in-memory data store..." },
-  # ... more messages
-]
-
-result = LlmOptimizer.optimize("What else can it do?", messages: messages)
-
 ## OptimizeResult
 
 Every call returns an `OptimizeResult` struct:

data/lib/generators/llm_optimizer/templates/initializer.rb

@@ -27,6 +27,9 @@ LlmOptimizer.configure do |config|
   config.cache_ttl            = 86_400 # cache entry TTL in seconds (default: 24h)
   config.timeout_seconds      = 5 # timeout for embedding / external API calls
 
+  # --- Tools ---
+  # config.with_tools = [] # Array of tool definitions (OpenAI/Anthropic format)
+
   # --- Logging ---
   config.logger        = Rails.logger
   config.debug_logging = Rails.env.development?
@@ -81,36 +84,39 @@ LlmOptimizer.configure do |config|
   # Messages caller for history manager/conversation summary - Optional
   # config.system_prompt = "You are a helpful person who gives responses in a non harmful way. " \
   #                  "If any serious question is asked, handle it in effectively."
-  # OpenAI implmeentation -
-  # config.messages_caller = ->(messages, model:) {
-  #   response = $openai.chat(
-  #     parameters: {
-  #       model: model,
-  #       messages: messages.map { |m| { role: m[:role], content: m[:content] } }
-  #     }
-  #   )
+  # OpenAI implementation -
+  # config.messages_caller = ->(messages, model:, tools: nil) {
+  #   parameters = {
+  #     model: model,
+  #     messages: messages.map { |m| { role: m[:role], content: m[:content] } }
+  #   }
+  #   parameters[:tools] = tools if tools&.any?
+  #
+  #   response = $openai.chat(parameters: parameters)
   #   response.dig("choices", 0, "message", "content")
   # }
 
   # RubyLLM implementation -
-  # config.messages_caller = ->(messages, model:) {
+  # config.messages_caller = ->(messages, model:, tools: nil) {
   #   chat = RubyLLM.chat(model: model)
+  #   chat.with_tools(*tools) if tools&.any?
   #   messages[0..-2].each { |m| chat.add_message(role: m[:role], content: m[:content]) }
   #   chat.ask(messages.last[:content]).content
   # }
 
   # Anthropic implementation -
-  # config.messages_caller = ->(messages, model:) {
+  # config.messages_caller = ->(messages, model:, tools: nil) {
   #   # Anthropic separates system messages from the messages array
   #   system_msg = messages.find { |m| m[:role] == "system" }&.dig(:content)
   #   chat_msgs  = messages.reject { |m| m[:role] == "system" }
   #                       .map { |m| { role: m[:role], content: m[:content] } }
-
+  #
   #   response = $anthropic.messages(
   #     model: model,
   #     max_tokens: 1024,
   #     system: system_msg,
-  #     messages: chat_msgs
+  #     messages: chat_msgs,
+  #     tools: tools
   #   )
   #   response["content"].first["text"]
   # }

data/lib/llm_optimizer/configuration.rb

@@ -25,6 +25,10 @@ module LlmOptimizer
       conversation_ttl
       system_prompt
       messages_caller
+      cache_scope
+      tools
+      with_tools
+      tools_caller
     ].freeze
 
     # Define readers for all known keys (setters below track explicit sets)
@@ -52,6 +56,7 @@ module LlmOptimizer
       @classifier_caller    = nil
       @conversation_ttl     = 86_400
       @system_prompt        = nil
+      @with_tools           = nil
     end
 
     # Copies only explicitly set keys from other_config without resetting unmentioned keys.

data/lib/llm_optimizer/optimize_result.rb

@@ -1,8 +1,43 @@
 # frozen_string_literal: true
 
 module LlmOptimizer
-  OptimizeResult = Struct.new(
-    :response, :model, :model_tier, :cache_status,
-    :original_tokens, :compressed_tokens, :latency_ms, :messages
-  )
+  class OptimizeResult
+    attr_accessor :response, :model, :model_tier, :cache_status,
+                  :original_tokens, :compressed_tokens, :input_tokens,
+                  :output_tokens, :cached_tokens, :latency_ms, :messages
+
+    # rubocop:disable Metrics/ParameterLists
+    def initialize(response: nil, model: nil, model_tier: nil, cache_status: nil,
+                   original_tokens: 0, compressed_tokens: 0, input_tokens: 0,
+                   output_tokens: 0, cached_tokens: 0, latency_ms: 0, messages: [])
+      @response = response
+      @model = model
+      @model_tier = model_tier
+      @cache_status = cache_status
+      @original_tokens = original_tokens
+      @compressed_tokens = compressed_tokens
+      @input_tokens = input_tokens
+      @output_tokens = output_tokens
+      @cached_tokens = cached_tokens
+      @latency_ms = latency_ms
+      @messages = messages
+    end
+    # rubocop:enable Metrics/ParameterLists
+
+    def to_h
+      {
+        response: @response,
+        model: @model,
+        model_tier: @model_tier,
+        cache_status: @cache_status,
+        original_tokens: @original_tokens,
+        compressed_tokens: @compressed_tokens,
+        input_tokens: @input_tokens,
+        output_tokens: @output_tokens,
+        cached_tokens: @cached_tokens,
+        latency_ms: @latency_ms,
+        messages: @messages
+      }
+    end
+  end
 end

data/lib/llm_optimizer/pipeline.rb

@@ -44,25 +44,6 @@ module LlmOptimizer
       [model_tier, model]
     end
 
-    def semantic_cache_lookup(prompt, model, model_tier, original_tokens,
-                              compressed_tokens, original_prompt, start, config)
-      return [nil, nil] unless config.use_semantic_cache
-
-      emb_client = EmbeddingClient.new(
-        model: config.embedding_model,
-        timeout_seconds: config.timeout_seconds,
-        embedding_caller: config.embedding_caller
-      )
-      embedding = emb_client.embed(prompt)
-      embedding, result = check_cache_hit(embedding, prompt, model, model_tier,
-                                          original_tokens, compressed_tokens,
-                                          original_prompt, start, config)
-      [embedding, result]
-    rescue EmbeddingError => e
-      config.logger.warn("[llm_optimizer] EmbeddingError (treating as cache miss): #{e.message}")
-      [nil, nil]
-    end
-
     def load_conversation(conversation_id, options, config)
       return [options[:messages], nil] unless conversation_id
 
@@ -93,41 +74,42 @@ module LlmOptimizer
       store.save(conversation_id, messages, prompt, response) || messages
     end
 
-    def store_in_cache(embedding, response, config)
-      return unless config.use_semantic_cache && embedding && config.redis_url
-
-      redis = build_redis(config.redis_url)
-      cache = SemanticCache.new(redis, threshold: config.similarity_threshold, ttl: config.cache_ttl)
-      cache.store(embedding, response)
-    rescue StandardError => e
-      config.logger.warn("[llm_optimizer] SemanticCache store failed: #{e.message}")
-    end
-
     def build_result(response, model, model_tier, cache_status,
-                     original_tokens, compressed_tokens, latency_ms, messages)
+                     original_tokens, compressed_tokens, latency_ms, messages, token_info = {})
       OptimizeResult.new(
         response: response, model: model, model_tier: model_tier,
         cache_status: cache_status, original_tokens: original_tokens,
-        compressed_tokens: compressed_tokens, latency_ms: latency_ms,
+        compressed_tokens: compressed_tokens,
+        input_tokens: token_info[:input_tokens] || compressed_tokens || original_tokens,
+        output_tokens: token_info[:output_tokens],
+        cached_tokens: token_info[:cached_tokens],
+        latency_ms: latency_ms,
         messages: messages
       )
     end
 
     def fallback_result(original_prompt, original_tokens, options, start)
       latency_ms = elapsed_ms(start)
-      response   = raw_llm_call(original_prompt, model: nil, config: configuration)
+      response, _token_info = raw_llm_call(original_prompt, model: nil, config: configuration)
       build_result(response, nil, nil, :miss, original_tokens || 0, nil,
                    latency_ms, options[:messages])
     end
 
     def raw_llm_call(prompt, model:, messages: nil, config: nil)
-      if messages && !messages.empty? && config&.messages_caller
-        config.messages_caller.call(messages + [{ role: "user", content: prompt }], model: model)
+      tools = config&.with_tools || config&.tools
+      result = if messages && !messages.empty? && config&.messages_caller
+                 config.messages_caller.call(messages + [{ role: "user", content: prompt }], model: model, tools: tools)
+               else
+                 llm = config&.llm_caller || @_current_llm_caller
+                 raise ConfigurationError, "No llm_caller configured." unless llm
+
+                 llm.call(prompt, model: model, tools: tools)
+               end
+
+      if result.is_a?(Hash)
+        [result[:content], result]
       else
-        llm = config&.llm_caller || @_current_llm_caller
-        raise ConfigurationError, "No llm_caller configured." unless llm
-
-        llm.call(prompt, model: model)
+        [result, {}]
       end
     end
 
@@ -152,22 +134,41 @@ module LlmOptimizer
       Redis.new(url: redis_url)
     end
 
-    def check_cache_hit(embedding, _prompt, model, model_tier, original_tokens,
-                        compressed_tokens, original_prompt, start, config)
-      return [embedding, nil] unless config.redis_url
+    def semantic_cache_lookup(prompt, model, model_tier, original_tokens,
+                              compressed_tokens, original_prompt, start, config)
+      return [nil, nil] unless config.use_semantic_cache
+
+      embedding = config.embedding_caller.call(prompt)
+      cache     = SemanticCache.new(build_redis(config.redis_url),
+                                    threshold: config.similarity_threshold,
+                                    ttl: config.cache_ttl,
+                                    cache_scope: config.cache_scope)
+      cached, token_info = cache.lookup(embedding)
+
+      if cached
+        latency_ms = elapsed_ms(start)
+        emit_log(config.logger, config,
+                 cache_status: :hit, model_tier: model_tier,
+                 original_tokens: original_tokens, compressed_tokens: compressed_tokens,
+                 latency_ms: latency_ms, prompt: original_prompt, response: cached)
+
+        [embedding, build_result(cached, model, model_tier, :hit,
+                                 original_tokens, compressed_tokens, latency_ms, nil, token_info)]
+      else
+        [embedding, nil]
+      end
+    rescue StandardError => e
+      config.logger.warn("[llm_optimizer] semantic_cache_lookup failed: #{e.message}")
+      [nil, nil]
+    end
 
-      redis  = build_redis(config.redis_url)
-      cache  = SemanticCache.new(redis, threshold: config.similarity_threshold, ttl: config.cache_ttl)
-      cached = cache.lookup(embedding)
-      return [embedding, nil] unless cached
+    def store_in_cache(embedding, response, config, token_info = {})
+      return unless config.use_semantic_cache && embedding
 
-      latency_ms = elapsed_ms(start)
-      emit_log(config.logger, config,
-               cache_status: :hit, model_tier: model_tier,
-               original_tokens: original_tokens, compressed_tokens: compressed_tokens,
-               latency_ms: latency_ms, prompt: original_prompt, response: cached)
-      [embedding, build_result(cached, model, model_tier, :hit,
-                               original_tokens, compressed_tokens, latency_ms, nil)]
+      SemanticCache.new(build_redis(config.redis_url),
+                        threshold: config.similarity_threshold,
+                        ttl: config.cache_ttl,
+                        cache_scope: config.cache_scope).store(embedding, response, token_info)
     end
   end
 end

data/lib/llm_optimizer/semantic_cache.rb

@@ -7,20 +7,19 @@ module LlmOptimizer
   class SemanticCache
     KEY_NAMESPACE = "llm_optimizer:cache:"
 
-    def initialize(redis_client, threshold:, ttl:)
-      @redis     = redis_client
-      @threshold = threshold
-      @ttl       = ttl
+    def initialize(redis_client, threshold:, ttl:, cache_scope: nil)
+      @redis       = redis_client
+      @threshold   = threshold
+      @ttl         = ttl
+      @cache_scope = cache_scope
     end
 
-    def store(embedding, response)
+    def store(embedding, response, token_info = {})
       key     = cache_key(embedding)
-      # Serialize embedding as raw 64-bit big-endian doubles to preserve full
-      # Float precision. MessagePack silently downcasts Ruby Float to 32-bit,
-      # which corrupts cosine similarity on deserialization.
       payload = MessagePack.pack({
-                                   "embedding" => embedding.pack("G*"), # binary string, lossless
-                                   "response" => response
+                                   "embedding" => embedding.pack("G*"),
+                                   "response" => response,
+                                   "token_info" => token_info
                                  })
       @redis.set(key, payload, ex: @ttl)
     rescue ::Redis::BaseError => e
@@ -28,28 +27,32 @@ module LlmOptimizer
     end
 
     def lookup(embedding)
-      keys = @redis.keys("#{KEY_NAMESPACE}*")
+      prefix = KEY_NAMESPACE
+      prefix += "#{@cache_scope}:" if @cache_scope
+      keys = @redis.keys("#{prefix}*")
+
+      keys.reject! { |k| k.count(":") > 2 } unless @cache_scope
+
       return nil if keys.empty?
 
       best_score    = -Float::INFINITY
-      best_response = nil
+      best_entry    = nil
 
       keys.each do |key|
         raw = @redis.get(key)
         next unless raw
 
         entry = MessagePack.unpack(raw)
-        # Unpack the binary string back to 64-bit doubles
         stored_embedding = entry["embedding"].unpack("G*")
         score = cosine_similarity(embedding, stored_embedding)
 
         if score > best_score
           best_score    = score
-          best_response = entry["response"]
+          best_entry    = entry
         end
       end
 
-      best_score >= @threshold ? best_response : nil
+      [best_entry["response"], best_entry["token_info"] || {}] if best_score >= @threshold
     rescue ::Redis::BaseError => e
       warn "[llm_optimizer] SemanticCache lookup failed: #{e.message}"
       nil
@@ -70,7 +73,9 @@ module LlmOptimizer
       # Use "G*" (64-bit big-endian double) to match Ruby's native Float precision.
       # "f*" (32-bit) truncates precision and produces inconsistent hashes for the
       # same embedding across serialize/deserialize round trips.
-      KEY_NAMESPACE + Digest::SHA256.hexdigest(embedding.pack("G*"))
+      prefix = KEY_NAMESPACE
+      prefix += "#{@cache_scope}:" if @cache_scope
+      prefix + Digest::SHA256.hexdigest(embedding.pack("G*"))
     end
   end
 end

data/lib/llm_optimizer/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module LlmOptimizer
-  VERSION = "0.1.5"
+  VERSION = "0.1.6"
 end

data/lib/llm_optimizer.rb

@@ -103,17 +103,18 @@ module LlmOptimizer
 
     messages, store = load_conversation(conversation_id, options, call_config)
     messages        = apply_history_manager(messages, call_config)
-    response        = raw_llm_call(prompt, messages: messages, model: model, config: call_config)
-    messages        = persist_conversation(store, conversation_id, messages, prompt, response)
-    store_in_cache(embedding, response, call_config)
+    response, token_info = raw_llm_call(prompt, messages: messages, model: model, config: call_config)
+    messages = persist_conversation(store, conversation_id, messages, prompt, response)
+    store_in_cache(embedding, response, call_config, token_info)
 
     latency_ms = elapsed_ms(start)
     emit_log(call_config.logger, call_config,
              cache_status: :miss, model_tier: model_tier,
              original_tokens: original_tokens, compressed_tokens: compressed_tokens,
              latency_ms: latency_ms, prompt: original_prompt, response: response)
+
     build_result(response, model, model_tier, :miss, original_tokens, compressed_tokens,
-                 latency_ms, messages)
+                 latency_ms, messages, token_info)
   rescue EmbeddingError => e
     configuration.logger.warn("[llm_optimizer] EmbeddingError (outer rescue): #{e.message}")
     fallback_result(original_prompt, original_tokens, options, start)

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: llm_optimizer
 version: !ruby/object:Gem::Version
-  version: 0.1.5
+  version: 0.1.6
 platform: ruby
 authors:
 - arun kumar