agentbill-sdk 5.0.1 → 7.17.0

data/lib/agentbill/wrappers.rb

@@ -0,0 +1,384 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+require 'net/http'
+require 'json'
+require 'uri'
+require 'digest'
+require 'time'
+
+require_relative 'exceptions'
+require_relative 'distributed'
+require_relative 'validation'
+
+module AgentBill
+  # Base class for AgentBill wrapper classes.
+  class BaseWrapper
+    BASE_URL = 'https://api.agentbill.io/functions/v1'
+    VERSION = '7.16.1'
+
+    attr_reader :provider_name
+
+    def initialize(api_key:, customer_id: nil, daily_budget: nil, monthly_budget: nil,
+                   base_url: nil, debug: false, **kwargs)
+      Validation.validate_api_key(api_key)
+      Validation.validate_customer_id(customer_id) if customer_id
+
+      @api_key = api_key
+      @customer_id = customer_id
+      @daily_budget = daily_budget
+      @monthly_budget = monthly_budget
+      @base_url = base_url || BASE_URL
+      @debug = debug
+      @provider_kwargs = kwargs
+      @provider_name = 'unknown'
+    end
+
+    protected
+
+    def log(message)
+      puts "[AgentBill] #{message}" if @debug
+    end
+
+    def validate_budget_sync(model, messages = nil)
+      payload = {
+        api_key: @api_key,
+        customer_id: @customer_id,
+        model: model,
+        daily_budget: @daily_budget,
+        monthly_budget: @monthly_budget
+      }
+
+      payload[:messages] = messages if messages
+
+      begin
+        url = URI("#{@base_url}/ai-cost-guard-router")
+        http = Net::HTTP.new(url.host, url.port)
+        http.use_ssl = true
+        http.read_timeout = 5
+
+        request = Net::HTTP::Post.new(url)
+        request['Content-Type'] = 'application/json'
+        request.body = JSON.generate(payload)
+
+        response = http.request(request)
+        result = JSON.parse(response.body)
+        log("Cost Guard validation: #{result}")
+        result
+      rescue StandardError => e
+        log("Cost Guard validation failed (allowing): #{e.message}")
+        { 'allowed' => true, 'reason' => 'validation_failed_open' }
+      end
+    end
+
+    def track_usage_sync(model:, input_tokens:, output_tokens:, latency_ms:, prompt_hash:,
+                          event_name: nil, metadata: nil)
+      ctx = Distributed.get_trace_context
+      if ctx
+        trace_id = ctx[:trace_id]
+        parent_span_id = ctx[:span_id]
+        span_id = SecureRandom.hex(8)
+        Distributed.set_trace_context(trace_id, span_id, parent_span_id)
+        log("Using trace context: trace_id=#{trace_id[0..7]}..., span_id=#{span_id[0..7]}...")
+      else
+        trace_id = SecureRandom.hex(16)
+        span_id = SecureRandom.hex(8)
+        parent_span_id = nil
+        Distributed.set_trace_context(trace_id, span_id)
+        log("Generated new trace context: trace_id=#{trace_id[0..7]}...")
+      end
+
+      now = Time.now.to_f
+      start_time_ns = ((now - latency_ms / 1000.0) * 1_000_000_000).to_i.to_s
+      end_time_ns = (now * 1_000_000_000).to_i.to_s
+
+      span_attributes = [
+        { key: 'gen_ai.request.model', value: { stringValue: model } },
+        { key: 'gen_ai.system', value: { stringValue: @provider_name } },
+        { key: 'gen_ai.usage.prompt_tokens', value: { intValue: input_tokens } },
+        { key: 'gen_ai.usage.completion_tokens', value: { intValue: output_tokens } },
+        { key: 'gen_ai.usage.total_tokens', value: { intValue: input_tokens + output_tokens } },
+        { key: 'agentbill.latency_ms', value: { doubleValue: latency_ms } },
+        { key: 'agentbill.prompt_hash', value: { stringValue: prompt_hash } },
+        { key: 'agentbill.event_name', value: { stringValue: event_name || 'ai_call' } }
+      ]
+
+      span_attributes << { key: 'agentbill.metadata', value: { stringValue: JSON.generate(metadata) } } if metadata
+
+      resource_attributes = [
+        { key: 'service.name', value: { stringValue: 'agentbill-ruby-sdk' } },
+        { key: 'service.version', value: { stringValue: VERSION } }
+      ]
+      resource_attributes << { key: 'customer.id', value: { stringValue: @customer_id } } if @customer_id
+
+      otel_span = {
+        traceId: trace_id,
+        spanId: span_id,
+        name: 'agentbill.trace.signal',
+        kind: 1,
+        startTimeUnixNano: start_time_ns,
+        endTimeUnixNano: end_time_ns,
+        attributes: span_attributes,
+        status: { code: 0 }
+      }
+      otel_span[:parentSpanId] = parent_span_id if parent_span_id
+
+      payload = {
+        resourceSpans: [{
+          resource: { attributes: resource_attributes },
+          scopeSpans: [{
+            scope: { name: 'agentbill', version: VERSION },
+            spans: [otel_span]
+          }]
+        }]
+      }
+
+      begin
+        url = URI("#{@base_url}/otel-collector")
+        http = Net::HTTP.new(url.host, url.port)
+        http.use_ssl = true
+        http.read_timeout = 5
+
+        request = Net::HTTP::Post.new(url)
+        request['Content-Type'] = 'application/json'
+        request['X-API-Key'] = @api_key
+        request.body = JSON.generate(payload)
+
+        response = http.request(request)
+        log("OTEL span sent: #{response.code}, trace_id=#{trace_id[0..7]}...")
+      rescue StandardError => e
+        log("Failed to send OTEL span: #{e.message}")
+      end
+    end
+
+    def cache_response_sync(model:, prompt_hash:, prompt_content:, response_content:,
+                             input_tokens:, output_tokens:, cost: 0.0, request_id: nil,
+                             cacheable: true, ttl_hours: 24)
+      prompt_str = prompt_content.is_a?(String) ? prompt_content : JSON.generate(prompt_content)
+
+      # v7.15.2: Send prompt_tokens and completion_tokens separately for accurate cost calculation
+      payload = {
+        api_key: @api_key,
+        prompt_hash: prompt_hash,
+        response_content: response_content,
+        model: model,
+        prompt_content: prompt_str,
+        tokens_used: input_tokens + output_tokens,
+        prompt_tokens: input_tokens,
+        completion_tokens: output_tokens,
+        cost: cost,
+        cacheable: cacheable,
+        ttl_hours: ttl_hours
+      }
+
+      payload[:customer_id] = @customer_id if @customer_id
+      payload[:request_id] = request_id if request_id
+
+      begin
+        url = URI("#{@base_url}/cache-ai-response")
+        http = Net::HTTP.new(url.host, url.port)
+        http.use_ssl = true
+        http.read_timeout = 5
+
+        request = Net::HTTP::Post.new(url)
+        request['Content-Type'] = 'application/json'
+        request.body = JSON.generate(payload)
+
+        response = http.request(request)
+        result = JSON.parse(response.body)
+        log("Cache response: cached=#{result['cached']}, cache_id=#{result['cache_id'] || 'none'}")
+      rescue StandardError => e
+        log("Failed to cache response (non-blocking): #{e.message}")
+      end
+    end
+
+    def hash_prompt(messages)
+      content = messages.is_a?(String) ? messages : JSON.generate(messages)
+      Digest::SHA256.hexdigest(content)
+    end
+  end
+
+  # OpenAI wrapper with automatic tracking
+  class AgentBillOpenAI < BaseWrapper
+    def initialize(**kwargs)
+      super(**kwargs)
+      @provider_name = 'openai'
+
+      # Lazy load OpenAI
+      require 'openai'
+      openai_options = @provider_kwargs.reject { |k, _| %i[api_key customer_id daily_budget monthly_budget base_url debug].include?(k) }
+      @client = OpenAI::Client.new(**openai_options)
+    end
+
+    def chat
+      @chat_namespace ||= ChatNamespace.new(self)
+    end
+
+    class ChatNamespace
+      def initialize(wrapper)
+        @wrapper = wrapper
+      end
+
+      def completions
+        @completions_namespace ||= CompletionsNamespace.new(@wrapper)
+      end
+    end
+
+    class CompletionsNamespace
+      def initialize(wrapper)
+        @wrapper = wrapper
+      end
+
+      def create(model:, messages:, **kwargs)
+        start_time = Time.now
+
+        # Validate budget first (CRITICAL: Always call router)
+        validation = @wrapper.send(:validate_budget_sync, model, messages)
+        
+        unless validation['allowed']
+          reason = validation['reason'] || 'Budget exceeded'
+          reason_lower = reason.downcase
+          
+          if reason_lower.include?('budget')
+            raise BudgetExceededError.new(reason, validation)
+          elsif reason_lower.include?('rate')
+            raise RateLimitExceededError.new(reason, validation)
+          else
+            raise PolicyViolationError.new(reason, validation)
+          end
+        end
+
+        # Check for cache hit
+        if validation['cache_hit'] && validation['cached_response']
+          @wrapper.send(:log, '✓ Cache hit - returning cached response')
+          return validation['cached_response']
+        end
+
+        # Make the actual API call
+        response = @wrapper.instance_variable_get(:@client).chat(
+          parameters: { model: model, messages: messages, **kwargs }
+        )
+
+        latency_ms = (Time.now - start_time) * 1000
+        input_tokens = response.dig('usage', 'prompt_tokens') || 0
+        output_tokens = response.dig('usage', 'completion_tokens') || 0
+        prompt_hash = @wrapper.send(:hash_prompt, messages)
+        
+        # v7.16.0: Track OpenAI native prompt prefix caching (different from AgentBill semantic cache)
+        # This is OpenAI's ~50% discount on repeated prompt prefixes, NOT our full semantic cache
+        cached_input_tokens = response.dig('usage', 'prompt_tokens_details', 'cached_tokens') || 0
+        reasoning_output_tokens = response.dig('usage', 'completion_tokens_details', 'reasoning_tokens') || 0
+        
+        metadata = {}
+        metadata[:openai_cached_input_tokens] = cached_input_tokens if cached_input_tokens > 0
+        metadata[:openai_reasoning_output_tokens] = reasoning_output_tokens if reasoning_output_tokens > 0
+
+        # Track usage via OTEL
+        @wrapper.send(:track_usage_sync,
+                      model: model,
+                      input_tokens: input_tokens,
+                      output_tokens: output_tokens,
+                      latency_ms: latency_ms,
+                      prompt_hash: prompt_hash,
+                      metadata: metadata.empty? ? nil : metadata)
+
+        # Cache the response
+        response_content = response.dig('choices', 0, 'message', 'content') || ''
+        @wrapper.send(:cache_response_sync,
+                      model: model,
+                      prompt_hash: prompt_hash,
+                      prompt_content: messages,
+                      response_content: response_content,
+                      input_tokens: input_tokens,
+                      output_tokens: output_tokens,
+                      request_id: validation['request_id'])
+
+        response
+      end
+    end
+  end
+
+  # Anthropic wrapper with automatic tracking
+  class AgentBillAnthropic < BaseWrapper
+    def initialize(**kwargs)
+      super(**kwargs)
+      @provider_name = 'anthropic'
+
+      # Lazy load Anthropic
+      require 'anthropic'
+      anthropic_options = @provider_kwargs.reject { |k, _| %i[api_key customer_id daily_budget monthly_budget base_url debug].include?(k) }
+      @client = Anthropic::Client.new(**anthropic_options)
+    end
+
+    def messages
+      @messages_namespace ||= MessagesNamespace.new(self)
+    end
+
+    class MessagesNamespace
+      def initialize(wrapper)
+        @wrapper = wrapper
+      end
+
+      def create(model:, max_tokens:, messages:, **kwargs)
+        start_time = Time.now
+
+        # Validate budget first (CRITICAL: Always call router)
+        validation = @wrapper.send(:validate_budget_sync, model, messages)
+        
+        unless validation['allowed']
+          reason = validation['reason'] || 'Budget exceeded'
+          reason_lower = reason.downcase
+          
+          if reason_lower.include?('budget')
+            raise BudgetExceededError.new(reason, validation)
+          elsif reason_lower.include?('rate')
+            raise RateLimitExceededError.new(reason, validation)
+          else
+            raise PolicyViolationError.new(reason, validation)
+          end
+        end
+
+        # Check for cache hit
+        if validation['cache_hit'] && validation['cached_response']
+          @wrapper.send(:log, '✓ Cache hit - returning cached response')
+          return validation['cached_response']
+        end
+
+        # Make the actual API call
+        response = @wrapper.instance_variable_get(:@client).messages(
+          model: model,
+          max_tokens: max_tokens,
+          messages: messages,
+          **kwargs
+        )
+
+        latency_ms = (Time.now - start_time) * 1000
+        input_tokens = response.dig('usage', 'input_tokens') || 0
+        output_tokens = response.dig('usage', 'output_tokens') || 0
+        prompt_hash = @wrapper.send(:hash_prompt, messages)
+
+        # Track usage via OTEL
+        @wrapper.send(:track_usage_sync,
+                      model: model,
+                      input_tokens: input_tokens,
+                      output_tokens: output_tokens,
+                      latency_ms: latency_ms,
+                      prompt_hash: prompt_hash)
+
+        # Cache the response
+        content_block = response.dig('content', 0)
+        response_content = content_block.is_a?(Hash) ? (content_block['text'] || '') : ''
+        @wrapper.send(:cache_response_sync,
+                      model: model,
+                      prompt_hash: prompt_hash,
+                      prompt_content: messages,
+                      response_content: response_content,
+                      input_tokens: input_tokens,
+                      output_tokens: output_tokens,
+                      request_id: validation['request_id'])
+
+        response
+      end
+    end
+  end
+end