agentbill-sdk 5.0.1 → 7.17.0

data/lib/agentbill/tracing.rb

@@ -0,0 +1,343 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+require 'net/http'
+require 'json'
+require 'uri'
+
+require_relative 'exceptions'
+require_relative 'distributed'
+require_relative 'signals'
+
+module AgentBill
+  # Context manager for tracking AI operations with Cost Guard protection.
+  #
+  # Implements the flow from the sequence diagram:
+  # 1. Pre-validate with ai-cost-guard-router
+  # 2. If blocked, raise BudgetExceededError
+  # 3. If allowed, execute AI calls
+  # 4. After completion, send spans to otel-collector
+  class TracingContext
+    attr_reader :trace_id, :span_id, :parent_span_id, :request_id, :start_time, :validation_result
+
+    BASE_URL = 'https://api.agentbill.io'
+    VERSION = '7.16.1'
+
+    def initialize(options = {})
+      @api_key = options[:api_key] || ENV['AGENTBILL_API_KEY']
+      @customer_id = options[:customer_id]
+      @daily_budget = options[:daily_budget]
+      @monthly_budget = options[:monthly_budget]
+      @agent_id = options[:agent_id]
+      @base_url = options[:base_url] || BASE_URL
+      @debug = options[:debug] || false
+      @model = options[:model] || 'gpt-4'
+      @estimated_tokens = options[:estimated_tokens] || 1000
+
+      @trace_id = nil
+      @span_id = nil
+      @parent_span_id = nil
+      @request_id = nil
+      @start_time = nil
+      @validation_result = nil
+      @spans = []
+    end
+
+    def enter
+      generate_ids
+      @start_time = Time.now
+
+      # Set global config for signal() function
+      Signals.set_config(
+        api_key: @api_key,
+        customer_id: @customer_id,
+        agent_id: @agent_id,
+        base_url: @base_url,
+        debug: @debug
+      )
+
+      # CRITICAL: Pre-validate with Cost Guard
+      @validation_result = validate_budget
+
+      # CRITICAL: Raise exception if not allowed
+      unless @validation_result['allowed']
+        reason = @validation_result['reason'] || 'Request blocked by Cost Guard'
+        reason_lower = reason.downcase
+
+        if reason_lower.include?('budget')
+          raise BudgetExceededError.new(reason, @validation_result)
+        elsif reason_lower.include?('rate')
+          raise RateLimitExceededError.new(reason, @validation_result)
+        else
+          raise PolicyViolationError.new(reason, @validation_result)
+        end
+      end
+
+      self
+    end
+
+    def exit(error = nil)
+      end_time = Time.now
+
+      # Add root span for the entire context
+      add_span(
+        'agentbill_tracing',
+        {
+          'customer_id' => @customer_id,
+          'agent_id' => @agent_id,
+          'request_id' => @request_id,
+          'error' => error&.message
+        },
+        (@start_time.to_f * 1_000_000_000).to_i,
+        (end_time.to_f * 1_000_000_000).to_i,
+        error ? 1 : 0
+      )
+
+      # Export spans to otel-collector
+      export_spans
+
+      # Clear trace context
+      Distributed.clear_trace_context
+    end
+
+    def add_span(name, attributes, start_time_ns, end_time_ns, status = 0)
+      span = {
+        'trace_id' => @trace_id,
+        'span_id' => SecureRandom.hex(8),
+        'parent_span_id' => @span_id,
+        'name' => name,
+        'attributes' => attributes,
+        'start_time_unix_nano' => start_time_ns,
+        'end_time_unix_nano' => end_time_ns,
+        'status' => { 'code' => status }
+      }
+      @spans << span
+    end
+
+    def cache_response(model, messages, response_content, input_tokens = 0, output_tokens = 0)
+      prompt_hash = hash_prompt(messages)
+      url = URI("#{@base_url}/functions/v1/cache-ai-response")
+
+      payload = {
+        api_key: @api_key,
+        prompt_hash: prompt_hash,
+        response_content: response_content,
+        model: model,
+        prompt_content: messages.is_a?(String) ? messages : JSON.generate(messages),
+        tokens_used: input_tokens + output_tokens,
+        cacheable: true,
+        ttl_hours: 24,
+        customer_id: @customer_id,
+        request_id: @request_id
+      }
+
+      begin
+        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)
+        log("✓ Cached response: #{response.code}") if @debug
+      rescue StandardError => e
+        log("⚠️ Cache error (non-blocking): #{e.message}") if @debug
+      end
+    end
+
+    private
+
+    def log(message)
+      puts "[AgentBill] #{message}" if @debug
+    end
+
+    def generate_ids
+      existing = Distributed.get_trace_context
+
+      if existing && existing[:trace_id]
+        @trace_id = existing[:trace_id]
+        @parent_span_id = existing[:span_id]
+      else
+        @trace_id = SecureRandom.hex(16)
+        @parent_span_id = nil
+      end
+
+      @span_id = SecureRandom.hex(8)
+      Distributed.set_trace_context(@trace_id, @span_id)
+    end
+
+    def validate_budget
+      url = URI("#{@base_url}/functions/v1/ai-cost-guard-router")
+
+      payload = {
+        api_key: @api_key,
+        customer_id: @customer_id,
+        model: @model,
+        messages: [],
+        estimated_tokens: @estimated_tokens
+      }
+
+      payload[:daily_budget_override] = @daily_budget if @daily_budget
+      payload[:monthly_budget_override] = @monthly_budget if @monthly_budget
+      payload[:agent_id] = @agent_id if @agent_id
+
+      begin
+        http = Net::HTTP.new(url.host, url.port)
+        http.use_ssl = true
+        http.read_timeout = 10
+
+        request = Net::HTTP::Post.new(url)
+        request['Content-Type'] = 'application/json'
+        request.body = JSON.generate(payload)
+
+        response = http.request(request)
+
+        if response.code.to_i >= 500
+          log('⚠️ Router server error (failing open)') if @debug
+          return { 'allowed' => true, 'reason' => 'Router server error (failed open)' }
+        end
+
+        if response.code.to_i >= 400
+          log("❌ Router rejected: #{response.body}") if @debug
+          return { 'allowed' => false, 'reason' => response.body }
+        end
+
+        result = JSON.parse(response.body)
+        @request_id = result['request_id']
+
+        if @debug
+          if result['allowed']
+            log('✓ Budget validation passed')
+          else
+            log("❌ Budget validation failed: #{result['reason']}")
+          end
+        end
+
+        result
+      rescue StandardError => e
+        log("⚠️ Router network error: #{e.message} (failing open)") if @debug
+        { 'allowed' => true, 'reason' => 'Router network error (failed open)' }
+      end
+    end
+
+    def export_spans
+      return if @spans.empty?
+
+      url = URI("#{@base_url}/functions/v1/otel-collector")
+
+      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
+      resource_attributes << { key: 'agent.id', value: { stringValue: @agent_id } } if @agent_id
+
+      otlp_spans = @spans.map do |span|
+        otlp_span = {
+          traceId: span['trace_id'],
+          spanId: span['span_id'],
+          name: span['name'],
+          kind: 1,
+          startTimeUnixNano: span['start_time_unix_nano'].to_s,
+          endTimeUnixNano: span['end_time_unix_nano'].to_s,
+          attributes: (span['attributes'] || {}).map do |k, v|
+            value = case v
+                    when String then { stringValue: v }
+                    when Integer then { intValue: v }
+                    when Float then { doubleValue: v }
+                    else { stringValue: v.to_s }
+                    end
+            { key: k, value: value }
+          end,
+          status: span['status']
+        }
+        otlp_span[:parentSpanId] = span['parent_span_id'] if span['parent_span_id']
+        otlp_span
+      end
+
+      payload = {
+        resourceSpans: [{
+          resource: { attributes: resource_attributes },
+          scopeSpans: [{
+            scope: { name: 'agentbill', version: VERSION },
+            spans: otlp_spans
+          }]
+        }]
+      }
+
+      begin
+        http = Net::HTTP.new(url.host, url.port)
+        http.use_ssl = true
+        http.read_timeout = 10
+
+        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)
+
+        if @debug
+          if response.code.to_i == 200
+            log("✓ Exported #{@spans.length} spans to otel-collector")
+          else
+            log("⚠️ Span export failed: #{response.code}")
+          end
+        end
+      rescue StandardError => e
+        log("⚠️ Span export error: #{e.message}") if @debug
+      end
+    end
+
+    def hash_prompt(messages)
+      content = messages.is_a?(String) ? messages : JSON.generate(messages)
+      Digest::SHA256.hexdigest(content)
+    end
+  end
+
+  class << self
+    # Context manager for tracking AI operations with Cost Guard protection.
+    #
+    # @example
+    #   AgentBill.tracing(customer_id: 'cust-123', api_key: 'ab_xxx') do |ctx|
+    #     response = openai.chat.completions.create(...)
+    #     AgentBill.signal('chat_completed', revenue: 0.50)
+    #   end
+    #
+    # @raise [BudgetExceededError] If daily/monthly budget would be exceeded
+    # @raise [RateLimitExceededError] If rate limits would be exceeded
+    # @raise [PolicyViolationError] If other policy constraints are violated
+    def tracing(customer_id:, api_key: nil, daily_budget: nil, monthly_budget: nil,
+                agent_id: nil, base_url: nil, debug: false, model: nil, estimated_tokens: 1000)
+      effective_api_key = api_key || ENV['AGENTBILL_API_KEY']
+      raise ArgumentError, 'api_key is required. Provide it as argument or set AGENTBILL_API_KEY env var.' unless effective_api_key
+
+      ctx = TracingContext.new(
+        api_key: effective_api_key,
+        customer_id: customer_id,
+        daily_budget: daily_budget,
+        monthly_budget: monthly_budget,
+        agent_id: agent_id,
+        base_url: base_url,
+        debug: debug,
+        model: model,
+        estimated_tokens: estimated_tokens
+      )
+
+      begin
+        ctx.enter
+        yield ctx
+      rescue BudgetExceededError, RateLimitExceededError, PolicyViolationError
+        raise
+      rescue StandardError => e
+        ctx.exit(e)
+        raise
+      else
+        ctx.exit
+      end
+    end
+  end
+end

data/lib/agentbill/version.rb

@@ -1,3 +1,3 @@
 module AgentBill
-  VERSION = "5.0.1"
+  VERSION = "7.17.0"
 end