completion-kit 0.28.42 → 0.28.43

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9976b877524a1b4f64e96c78c3d5f158f1157c749902d27dcf37efe18c9bdb13
-  data.tar.gz: bc2fd622095d7556b40fb2de3dcd3836ea47b77c4a784e9fd1e6c1fd0d5dc381
+  metadata.gz: 3cd88afe200e8524d2cad09e6214612f645d9edb42ca0300c8c72c3dd2edf77c
+  data.tar.gz: 8873756fb01f18fa8405dc486e2f8d4bfbe4a11b37c6f1288ec80012e0a6bd18
 SHA512:
-  metadata.gz: 21b4460aea3386e6ecb936ef2a3303401929a9d26cba52e3fe9200efca47128625a4617f07f86539d43d9f700e3bba24d86f69a877921e128cf7ede75a5e1c7d
-  data.tar.gz: 611a2a701c2db81edb16cee9ca7fa1d69e62b945ad07787cd2a979f33b762979abf94265150627c883a5bccd2c56b656b449b212eb1d0e3bde373c66cd5884c7
+  metadata.gz: 5b6f6806c9d3522221ecc2cb9249ed87082961712d1f8ed75860998e5774bc28693ea5dc0957d2a7ce371fb23f97578b87ca1b7ed181a8fe8d81e989682061c6
+  data.tar.gz: 57c412ed0623239f5f88a731ca46a70404ad4f4211bc7ff06ab612057522b23a6edd13ac6ae9dfde578e1ba22dc5eeda0ce5ae3b9ec3488391ba844475ff860d

data/app/services/completion_kit/azure_foundry_client.rb

@@ -4,7 +4,6 @@ module CompletionKit
     def generate_completion(prompt, options = {})
       @temperature_dropped = false
       return "Error: Azure provider is not fully configured" unless configured?
-      return "Error: API endpoint resolves to a private address" unless ProviderEndpoint.safe?(api_endpoint)
 
       model = options[:model]
       max_tokens = options[:max_tokens] || 1000
@@ -47,13 +46,14 @@ module CompletionKit
       raise
     rescue Faraday::Error
       raise
+    rescue ProviderEndpoint::UnsafeEndpoint => e
+      "Error: API endpoint #{e.message}"
     rescue => e
       "Error: #{e.message}"
     end
 
     def available_models
       return [] unless configured?
-      return [] unless ProviderEndpoint.safe?(api_endpoint)
 
       path = foundry_project? ? "#{azure_base_url}/deployments?api-version=v1" : models_path
       response = build_connection(azure_base_url).get(path) do |req|
@@ -88,6 +88,10 @@ module CompletionKit
 
     private
 
+    def attach_adapter(builder, url)
+      ProviderEndpoint.pin(builder, url)
+    end
+
     def api_key
       @config[:api_key]
     end

data/app/services/completion_kit/llm_client.rb

@@ -102,8 +102,12 @@ module CompletionKit
         f.options.timeout = timeout if timeout
         f.options.open_timeout = open_timeout if open_timeout
         f.request :retry, max: 2, interval: 0.5
-        f.adapter Faraday.default_adapter
+        attach_adapter(f, url)
       end
     end
+
+    def attach_adapter(builder, _url)
+      builder.adapter Faraday.default_adapter
+    end
   end
 end

data/app/services/completion_kit/mcp_tools/datasets.rb

@@ -90,10 +90,7 @@ module CompletionKit
       end
 
       def self.create_from_url(args)
-        issues = ProviderEndpoint.validate(args["url"])
-        return error_result("URL is not allowed (#{issues.join(", ")}).") if issues.any?
-
-        response = csv_connection.get(args["url"])
+        response = csv_connection(args["url"]).get(args["url"])
         return error_result("Could not download CSV (HTTP #{response.status}).") unless response.success?
 
         body = response.body.to_s
@@ -106,15 +103,17 @@ module CompletionKit
         else
           error_result(dataset.errors.full_messages.join(", "))
         end
+      rescue ProviderEndpoint::UnsafeEndpoint => e
+        error_result("URL is not allowed: it #{e.message}.")
       rescue Faraday::Error => e
         error_result("Could not download CSV: #{e.message}")
       end
 
-      def self.csv_connection
+      def self.csv_connection(url)
         Faraday.new do |f|
           f.options.timeout = 30
           f.options.open_timeout = 5
-          f.adapter Faraday.default_adapter
+          ProviderEndpoint.pin(f, url)
         end
       end
     end

data/app/services/completion_kit/model_discovery_service.rb

@@ -40,14 +40,20 @@ module CompletionKit
       end
     end
 
-    def fetch_connection(base_url)
+    def fetch_connection(base_url, pinned: false)
       Faraday.new(url: base_url) do |f|
         f.options.timeout = 15
         f.options.open_timeout = 5
-        f.adapter Faraday.default_adapter
+        pinned ? pin_endpoint(f, base_url) : f.adapter(Faraday.default_adapter)
       end
     end
 
+    def pin_endpoint(builder, url)
+      ProviderEndpoint.pin(builder, url)
+    rescue ProviderEndpoint::UnsafeEndpoint => e
+      raise DiscoveryError, "The model endpoint #{e.message}."
+    end
+
     def raise_fetch_error!(response)
       label = case response.status
               when 401, 403 then "Invalid API key for #{@provider}"
@@ -122,7 +128,7 @@ module CompletionKit
     def fetch_ollama_models
       raise DiscoveryError, "A model endpoint URL is required." if @api_endpoint.blank?
       base_url = ollama_root_url
-      response = fetch_connection(base_url).get("/v1/models") do |req|
+      response = fetch_connection(base_url, pinned: true).get("/v1/models") do |req|
         req.headers["Authorization"] = "Bearer #{@api_key}" if @api_key.present?
       end
       raise DiscoveryError, custom_endpoint_error_message(response) unless response.success?
@@ -154,7 +160,7 @@ module CompletionKit
     end
 
     def custom_endpoint_host
-      ProviderEndpoint.parse(@api_endpoint)&.host || @api_endpoint.to_s
+      ProviderEndpoint.parse(@api_endpoint).host
     end
 
     def azure_custom_host?
@@ -170,7 +176,7 @@ module CompletionKit
       raise DiscoveryError, "An Azure endpoint URL is required." if @api_endpoint.blank?
 
       path = azure_foundry_project? ? "#{azure_base_url}/deployments?api-version=v1" : azure_models_path
-      response = fetch_connection(azure_base_url).get(path) do |req|
+      response = fetch_connection(azure_base_url, pinned: true).get(path) do |req|
         req.headers["api-key"] = @api_key
       end
       raise DiscoveryError, azure_error_message(response) unless response.success?
@@ -185,7 +191,7 @@ module CompletionKit
     end
 
     def fetch_azure_catalog_count
-      response = fetch_connection(azure_base_url).get("/openai/v1/models") do |req|
+      response = fetch_connection(azure_base_url, pinned: true).get("/openai/v1/models") do |req|
         req.headers["api-key"] = @api_key
       end
       return nil unless response.success?
@@ -464,7 +470,7 @@ module CompletionKit
         f.options.timeout = 60
         f.options.open_timeout = 5
         f.request :retry, max: 1, interval: 0.5
-        f.adapter Faraday.default_adapter
+        pin_endpoint(f, ollama_root_url)
       end
       conn.post do |req|
         req.url "/v1/chat/completions"
@@ -479,7 +485,7 @@ module CompletionKit
         f.options.timeout = 60
         f.options.open_timeout = 5
         f.request :retry, max: 1, interval: 0.5
-        f.adapter Faraday.default_adapter
+        pin_endpoint(f, azure_base_url)
       end
       response = azure_probe_post(conn, model_id, input, max_tokens, max_completion: false)
       if response.status == 400 && azure_max_tokens_unsupported?(response.body)

data/app/services/completion_kit/ollama_client.rb

@@ -4,7 +4,6 @@ module CompletionKit
     def generate_completion(prompt, options = {})
       @temperature_dropped = false
       return "Error: API endpoint not configured" unless configured?
-      return "Error: API endpoint resolves to a private address" unless ProviderEndpoint.safe?(api_endpoint)
 
       model = options[:model]
       max_tokens = options[:max_tokens] || 1000
@@ -36,13 +35,14 @@ module CompletionKit
       raise
     rescue Faraday::Error
       raise
+    rescue ProviderEndpoint::UnsafeEndpoint => e
+      "Error: API endpoint #{e.message}"
     rescue => e
       "Error: #{e.message}"
     end
 
     def available_models
       return [] unless configured?
-      return [] unless ProviderEndpoint.safe?(api_endpoint)
 
       response = build_connection(api_endpoint).get("/v1/models") do |req|
         req.headers["Authorization"] = "Bearer #{api_key}" if api_key.present?
@@ -68,6 +68,10 @@ module CompletionKit
 
     private
 
+    def attach_adapter(builder, url)
+      ProviderEndpoint.pin(builder, url)
+    end
+
     def api_key
       @config[:api_key] || ENV["OLLAMA_API_KEY"]
     end

data/app/services/completion_kit/provider_endpoint.rb

@@ -1,24 +1,41 @@
 require "ipaddr"
-require "resolv"
+require "socket"
 
 module CompletionKit
   module ProviderEndpoint
-    ZERO_NET = IPAddr.new("0.0.0.0/8").freeze
+    class UnsafeEndpoint < StandardError; end
+
+    BLOCKED_NETS = [IPAddr.new("0.0.0.0/8"), IPAddr.new("100.64.0.0/10")].freeze
+    ISSUE_MESSAGES = {
+      invalid_url: "is not a valid http or https URL",
+      unresolvable: "could not be resolved",
+      unsafe_host: "resolves to a private address"
+    }.freeze
 
     module_function
 
     def validate(url)
       uri = parse(url)
       return [:invalid_url] unless uri
-      addrs = addresses(uri.host)
-      return [:unresolvable] if addrs.empty?
-      return [:unsafe_host] if addrs.any? { |ip| unsafe?(ip) }
-      []
+      issues_for(addresses(uri.host))
     end
 
     def safe?(url)
-      errors = validate(url)
-      errors.empty? || errors == [:unresolvable]
+      validate(url).empty?
+    end
+
+    def pinned_address(url)
+      uri = parse(url)
+      raise UnsafeEndpoint, ISSUE_MESSAGES[:invalid_url] unless uri
+      addrs = addresses(uri.host)
+      issue = issues_for(addrs).first
+      raise UnsafeEndpoint, ISSUE_MESSAGES[issue] if issue
+      (addrs.find(&:ipv4?) || addrs.first).to_s
+    end
+
+    def pin(builder, url)
+      address = pinned_address(url)
+      builder.adapter(:net_http) { |http| http.ipaddr = address }
     end
 
     def parse(value)
@@ -28,18 +45,31 @@ module CompletionKit
       nil
     end
 
+    def issues_for(addrs)
+      return [:unresolvable] if addrs.empty?
+      return [:unsafe_host] if addrs.any? { |ip| unsafe?(ip) }
+      []
+    end
+
     def addresses(host)
-      bare = host.delete_prefix("[").delete_suffix("]")
+      bare = host.delete_prefix("[").delete_suffix("]").delete_suffix(".")
       [IPAddr.new(bare)]
     rescue IPAddr::InvalidAddressError
-      Resolv.getaddresses(host).map { |addr| IPAddr.new(addr) }
+      resolve(bare).map { |addr| IPAddr.new(addr) }
+    end
+
+    def resolve(host)
+      Addrinfo.getaddrinfo(host, nil, nil, :STREAM).map(&:ip_address).uniq
+    rescue SocketError
+      []
     end
 
     def unsafe?(ip)
+      ip = ip.native
       return true if ip.private?
       return true if ip.link_local?
       return true if ip.to_i.zero?
-      return true if ip.ipv4? && ZERO_NET.include?(ip)
+      return true if ip.ipv4? && BLOCKED_NETS.any? { |net| net.include?(ip) }
       return true if ip.loopback? && !CompletionKit.config.allow_loopback_endpoints
       false
     end

data/lib/completion_kit/version.rb

@@ -1,3 +1,3 @@
 module CompletionKit
-  VERSION = "0.28.42"
+  VERSION = "0.28.43"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: completion-kit
 version: !ruby/object:Gem::Version
-  version: 0.28.42
+  version: 0.28.43
 platform: ruby
 authors:
 - Damien Bastin