agentbill-sdk 5.0.1 → 7.17.0

data/lib/agentbill/distributed.rb

@@ -0,0 +1,109 @@
+# AgentBill Distributed Tracing for Ruby
+#
+# Functions for cross-service trace propagation and correlation.
+# This is the SINGLE SOURCE OF TRUTH for trace context in the Ruby SDK.
+
+module AgentBill
+  module Distributed
+    # Thread-local storage for trace context
+    @trace_context = {}
+    @tracing_token = nil
+
+    class << self
+      # Get the current trace context (trace_id, span_id)
+      # This is the SINGLE SOURCE OF TRUTH used by the tracer.
+      def get_trace_context
+        Thread.current[:agentbill_trace_context]
+      end
+
+      # Set the trace context manually.
+      # Called by the tracer after starting a span.
+      def set_trace_context(trace_id, span_id)
+        Thread.current[:agentbill_trace_context] = {
+          trace_id: trace_id,
+          span_id: span_id
+        }
+      end
+
+      # Clear the current trace context
+      def clear_trace_context
+        Thread.current[:agentbill_trace_context] = nil
+        Thread.current[:agentbill_tracing_token] = nil
+      end
+
+      # Get the current tracing token
+      def get_tracing_token
+        Thread.current[:agentbill_tracing_token]
+      end
+
+      # Generate a new tracing token for distributed trace propagation
+      def generate_tracing_token
+        ctx = get_trace_context
+        random_suffix = SecureRandom.hex(4)
+
+        if ctx
+          token = "agentbill-v1-#{ctx[:trace_id]}-#{ctx[:span_id]}-#{random_suffix}"
+        else
+          trace_id = SecureRandom.hex(16)
+          span_id = SecureRandom.hex(8)
+          set_trace_context(trace_id, span_id)
+          token = "agentbill-v1-#{trace_id}-#{span_id}-#{random_suffix}"
+        end
+
+        Thread.current[:agentbill_tracing_token] = token
+        token
+      end
+
+      # Set the tracing token received from an upstream service
+      def set_tracing_token(token)
+        return false unless token&.start_with?('agentbill-v1-')
+
+        parts = token.split('-')
+        return false if parts.length < 5
+
+        trace_id = parts[2]
+        span_id = parts[3]
+
+        set_trace_context(trace_id, span_id)
+        Thread.current[:agentbill_tracing_token] = token
+        true
+      rescue
+        false
+      end
+
+      # Get HTTP headers for trace propagation
+      def propagate_trace_headers
+        headers = {}
+
+        token = get_tracing_token
+        headers['X-AgentBill-Trace'] = token if token
+
+        ctx = get_trace_context
+        if ctx
+          headers['traceparent'] = "00-#{ctx[:trace_id]}-#{ctx[:span_id]}-01"
+        end
+
+        headers
+      end
+
+      # Extract trace context from incoming HTTP headers
+      def extract_trace_from_headers(headers)
+        # Try AgentBill token first
+        token = headers['X-AgentBill-Trace'] || headers['x-agentbill-trace']
+        return true if token && set_tracing_token(token)
+
+        # Try W3C traceparent
+        traceparent = headers['traceparent'] || headers['Traceparent']
+        if traceparent
+          parts = traceparent.split('-')
+          if parts.length >= 3
+            set_trace_context(parts[1], parts[2])
+            return true
+          end
+        end
+
+        false
+      end
+    end
+  end
+end

data/lib/agentbill/exceptions.rb

@@ -0,0 +1,84 @@
+# frozen_string_literal: true
+
+# AgentBill Exceptions
+#
+# Custom exception classes for Cost Guard protection and error handling.
+
+module AgentBill
+  # Base class for all AgentBill errors
+  class AgentBillError < StandardError
+    attr_reader :details
+
+    def initialize(message, details = {})
+      @details = details
+      super(message)
+    end
+  end
+
+  # Raised when a request would exceed the configured budget limits.
+  #
+  # This is thrown by the context manager or wrapper when ai-cost-guard-router
+  # returns allowed=false due to budget constraints.
+  #
+  # @example
+  #   begin
+  #     AgentBill.tracing(customer_id: 'cust-123', daily_budget: 10.00) do |ctx|
+  #       response = openai.chat.completions.create(...)
+  #     end
+  #   rescue AgentBill::BudgetExceededError => e
+  #     puts "Budget exceeded: #{e.reason}"
+  #     # Handle gracefully - queue for later, show user message, etc.
+  #   end
+  class BudgetExceededError < AgentBillError
+    attr_reader :reason, :code
+
+    def initialize(reason, details = {})
+      @reason = reason
+      @code = 'BUDGET_EXCEEDED'
+      super("Budget exceeded: #{reason}", details)
+    end
+  end
+
+  # Raised when a request would exceed rate limits.
+  class RateLimitExceededError < AgentBillError
+    attr_reader :reason, :code
+
+    def initialize(reason, details = {})
+      @reason = reason
+      @code = 'RATE_LIMIT_EXCEEDED'
+      super("Rate limit exceeded: #{reason}", details)
+    end
+  end
+
+  # Raised when a request violates a Cost Guard policy.
+  class PolicyViolationError < AgentBillError
+    attr_reader :reason, :code
+
+    def initialize(reason, details = {})
+      @reason = reason
+      @code = 'POLICY_VIOLATION'
+      super("Policy violation: #{reason}", details)
+    end
+  end
+
+  # Raised when input validation fails.
+  class ValidationError < AgentBillError
+    attr_reader :field, :code
+
+    def initialize(message, field = nil)
+      @field = field
+      @code = 'VALIDATION_ERROR'
+      super("Validation error: #{message}")
+    end
+  end
+
+  # Raised when tracing operations fail.
+  class TracingError < AgentBillError
+    attr_reader :code
+
+    def initialize(message, details = {})
+      @code = 'TRACING_ERROR'
+      super("Tracing error: #{message}", details)
+    end
+  end
+end

data/lib/agentbill/ollama_wrapper.rb

@@ -0,0 +1,153 @@
+# Ollama Wrapper for AgentBill Ruby SDK (Local, $0 cost)
+require_relative 'pricing'
+
+module AgentBill
+  class OllamaWrapper
+    def initialize(client, config, tracer)
+      @client = client
+      @config = config
+      @tracer = tracer
+    end
+
+    def wrap
+      wrap_chat
+      wrap_generate
+      @client
+    end
+
+    private
+
+    def wrap_chat
+      original_chat = @client.method(:chat)
+      config = @config
+      tracer = @tracer
+
+      @client.define_singleton_method(:chat) do |params|
+        model = params[:model] || 'llama2'
+
+        start_time = Time.now
+        span = tracer.start_span('ollama.chat', {
+          # ✅ OTEL GenAI compliant attributes
+          'gen_ai.system' => 'ollama',
+          'gen_ai.request.model' => model,
+          'gen_ai.operation.name' => 'chat',
+          # ⚠️ Backward compatibility
+          'model' => model,
+          'provider' => 'ollama'
+        })
+
+        begin
+          response = original_chat.call(params)
+          latency = ((Time.now - start_time) * 1000).round
+
+          input_tokens = response[:prompt_eval_count] || 0
+          output_tokens = response[:eval_count] || 0
+
+          track_usage(model, 'ollama', input_tokens, output_tokens, latency, 0, config)
+
+          span.set_attributes({
+            # ✅ OTEL GenAI compliant attributes
+            'gen_ai.usage.input_tokens' => input_tokens,
+            'gen_ai.usage.output_tokens' => output_tokens,
+            # ⚠️ Backward compatibility
+            'response.prompt_eval_count' => input_tokens,
+            'response.eval_count' => output_tokens,
+            'latency_ms' => latency,
+            'cost' => 0
+          })
+          span.set_status(0)
+
+          puts "[AgentBill] ✓ Ollama chat tracked (local, $0.00)" if config[:debug]
+          response
+        rescue => e
+          span.set_status(1, e.message)
+          raise
+        ensure
+          span.finish
+        end
+      end
+    end
+
+    def wrap_generate
+      original_generate = @client.method(:generate)
+      config = @config
+      tracer = @tracer
+
+      @client.define_singleton_method(:generate) do |params|
+        model = params[:model] || 'llama2'
+
+        start_time = Time.now
+        span = tracer.start_span('ollama.generate', {
+          # ✅ OTEL GenAI compliant attributes
+          'gen_ai.system' => 'ollama',
+          'gen_ai.request.model' => model,
+          'gen_ai.operation.name' => 'text_completion',
+          # ⚠️ Backward compatibility
+          'model' => model,
+          'provider' => 'ollama'
+        })
+
+        begin
+          response = original_generate.call(params)
+          latency = ((Time.now - start_time) * 1000).round
+
+          input_tokens = response[:prompt_eval_count] || 0
+          output_tokens = response[:eval_count] || 0
+
+          track_usage(model, 'ollama', input_tokens, output_tokens, latency, 0, config)
+
+          span.set_attributes({
+            # ✅ OTEL GenAI compliant attributes
+            'gen_ai.usage.input_tokens' => input_tokens,
+            'gen_ai.usage.output_tokens' => output_tokens,
+            # ⚠️ Backward compatibility
+            'response.prompt_eval_count' => input_tokens,
+            'response.eval_count' => output_tokens,
+            'latency_ms' => latency,
+            'cost' => 0
+          })
+          span.set_status(0)
+
+          puts "[AgentBill] ✓ Ollama generate tracked (local, $0.00)" if config[:debug]
+          response
+        rescue => e
+          span.set_status(1, e.message)
+          raise
+        ensure
+          span.finish
+        end
+      end
+    end
+
+    def track_usage(model, provider, input_tokens, output_tokens, latency, cost, config)
+      uri = URI("#{config[:base_url] || 'https://api.agentbill.io'}/functions/v1/track-ai-usage")
+
+      payload = {
+        api_key: config[:api_key],
+        customer_id: config[:customer_id],
+        agent_id: config[:agent_id],
+        event_name: 'ai_request',
+        model: model,
+        provider: provider,
+        prompt_tokens: input_tokens,
+        completion_tokens: output_tokens,
+        latency_ms: latency,
+        cost: cost
+      }
+
+      begin
+        http = Net::HTTP.new(uri.host, uri.port)
+        http.use_ssl = true
+        http.read_timeout = 10
+
+        request = Net::HTTP::Post.new(uri.path)
+        request['Content-Type'] = 'application/json'
+        request.body = payload.to_json
+
+        http.request(request)
+      rescue => e
+        puts "[AgentBill] Tracking failed: #{e.message}" if config[:debug]
+      end
+    end
+  end
+end

data/lib/agentbill/orders.rb

@@ -0,0 +1,283 @@
+# AgentBill Orders Resource - CRUD operations with line items
+# 
+# IMPORTANT: At least one agent must be associated with an order via
+# agent_ids or agent_external_ids. Orders cannot be created without agents.
+#
+# Example:
+#   ab = AgentBill::Client.init(api_key: 'agb_...')
+#   
+#   # List orders
+#   result = ab.orders.list(limit: 10)
+#   
+#   # Create order (agent_ids or agent_external_ids REQUIRED)
+#   order = ab.orders.create(
+#     account_id: 'acct-123',
+#     name: 'Enterprise Q1',
+#     agent_ids: ['agent-uuid-1']  # REQUIRED
+#   )
+#   
+#   # Add line item
+#   ab.orders.add_line_item(
+#     order_id: order['id'],
+#     description: 'AI Usage',
+#     quantity: 1000,
+#     unit_price: 0.05,
+#     item_type: 'usage'
+#   )
+#   
+#   # Get order by ID
+#   order = ab.orders.get('order-123')
+#   
+#   # Update order
+#   order = ab.orders.update('order-123', status: 'active')
+#   
+#   # Delete order
+#   ab.orders.delete('order-123')
+
+require 'net/http'
+require 'json'
+
+module AgentBill
+  # Valid item_type values for add_line_item:
+  # - usage: Variable consumption billing (API calls, tokens, requests)
+  # - platform_fee: Platform access and service charges
+  # - setup_fee: One-time onboarding and implementation costs
+  # - base_fee: Fixed recurring monthly/annual charges
+  # - overage: Usage exceeding included plan amounts
+  # - custom: Flexible billing for unique scenarios
+  # - flat_rate: Fixed-price service packages or retainers
+  # - product: Physical or digital product sales
+  # - service: Professional services and consulting
+  # - subscription: Recurring subscription billing
+  VALID_ITEM_TYPES = %w[usage platform_fee setup_fee base_fee overage custom product service subscription flat_rate].freeze
+
+  class OrdersResource
+    def initialize(config)
+      @config = config
+      @base_url = config[:base_url] || 'https://api.agentbill.io'
+      @api_key = config[:api_key]
+      @debug = config[:debug] || false
+    end
+
+    # List orders with pagination and filtering
+    #
+    # @param account_id [String, nil] Filter by account ID
+    # @param status [String, nil] Filter by status (draft, active, completed, cancelled)
+    # @param limit [Integer] Number of orders to return (default: 50)
+    # @param offset [Integer] Offset for pagination (default: 0)
+    # @param search [String, nil] Search term to filter by name or external_id
+    # @return [Hash] Hash with 'data' (list of orders) and 'pagination' info
+    def list(account_id: nil, status: nil, limit: 50, offset: 0, search: nil)
+      uri = URI("#{@base_url}/functions/v1/api-orders")
+      params = { limit: limit.to_s, offset: offset.to_s }
+      params[:search] = search if search
+      params[:account_id] = account_id if account_id
+      params[:status] = status if status
+      uri.query = URI.encode_www_form(params)
+
+      response = make_request(:get, uri)
+      handle_response(response, 'list orders')
+    end
+
+    # Create a new order
+    #
+    # IMPORTANT: At least one agent must be specified via agent_ids or agent_external_ids.
+    # Orders cannot be created without agent association for proper billing attribution.
+    #
+    # @param account_id [String] Account (customer) UUID (required unless account_external_id provided)
+    # @param name [String] Order name (required)
+    # @param agent_ids [Array<String>, nil] List of agent UUIDs (required if agent_external_ids not provided)
+    # @param agent_external_ids [Array<String>, nil] List of agent external IDs (required if agent_ids not provided)
+    # @param external_id [String, nil] Your external ID for this order
+    # @param account_external_id [String, nil] External ID of the account
+    # @param billing_contact_id [String, nil] Contact UUID for billing
+    # @param billing_contact_external_id [String, nil] External ID of billing contact
+    # @param start_date [String, nil] Order start date (ISO format)
+    # @param end_date [String, nil] Order end date (ISO format)
+    # @param status [String, nil] Order status (draft, active, completed, cancelled)
+    # @param currency [String, nil] Currency code (default: USD)
+    # @param metadata [Hash, nil] Additional metadata
+    # @return [Hash] Created order object
+    # @raise [ArgumentError] If neither agent_ids nor agent_external_ids provided
+    def create(account_id:, name:, agent_ids: nil, agent_external_ids: nil, external_id: nil,
+               account_external_id: nil, billing_contact_id: nil, billing_contact_external_id: nil,
+               start_date: nil, end_date: nil, status: nil, currency: nil, metadata: nil)
+      # Validate agent requirement
+      if (agent_ids.nil? || agent_ids.empty?) && (agent_external_ids.nil? || agent_external_ids.empty?)
+        raise ArgumentError, 'At least one agent must be specified via agent_ids or agent_external_ids. ' \
+                             'Orders cannot be created without agent association for billing attribution.'
+      end
+
+      uri = URI("#{@base_url}/functions/v1/api-orders")
+      
+      payload = { account_id: account_id, name: name }
+      payload[:agent_ids] = agent_ids if agent_ids&.any?
+      payload[:agent_external_ids] = agent_external_ids if agent_external_ids&.any?
+      payload[:external_id] = external_id if external_id
+      payload[:account_external_id] = account_external_id if account_external_id
+      payload[:billing_contact_id] = billing_contact_id if billing_contact_id
+      payload[:billing_contact_external_id] = billing_contact_external_id if billing_contact_external_id
+      payload[:start_date] = start_date if start_date
+      payload[:end_date] = end_date if end_date
+      payload[:status] = status if status
+      payload[:currency] = currency if currency
+      payload[:metadata] = metadata if metadata
+
+      response = make_request(:post, uri, payload)
+      handle_response(response, 'create order')
+    end
+
+    # Get an order by ID or external_id (includes line items)
+    #
+    # @param order_id [String] Order UUID or external_id
+    # @param by_external_id [Boolean] If true, treat order_id as external_id
+    # @return [Hash] Order object with order_line_items
+    def get(order_id, by_external_id: false)
+      uri = URI("#{@base_url}/functions/v1/api-orders")
+      params = by_external_id ? { external_id: order_id } : { id: order_id }
+      uri.query = URI.encode_www_form(params)
+
+      response = make_request(:get, uri)
+      handle_response(response, 'get order')
+    end
+
+    # Update an order
+    #
+    # @param order_id [String] Order UUID
+    # @param name [String, nil] New name
+    # @param billing_contact_id [String, nil] New billing contact ID
+    # @param start_date [String, nil] New start date
+    # @param end_date [String, nil] New end date
+    # @param status [String, nil] New status
+    # @param currency [String, nil] New currency
+    # @param metadata [Hash, nil] New metadata (replaces existing)
+    # @return [Hash] Updated order object
+    def update(order_id, name: nil, billing_contact_id: nil, start_date: nil, 
+               end_date: nil, status: nil, currency: nil, metadata: nil)
+      uri = URI("#{@base_url}/functions/v1/api-orders")
+      
+      payload = { id: order_id }
+      payload[:name] = name if name
+      payload[:billing_contact_id] = billing_contact_id if billing_contact_id
+      payload[:start_date] = start_date if start_date
+      payload[:end_date] = end_date if end_date
+      payload[:status] = status if status
+      payload[:currency] = currency if currency
+      payload[:metadata] = metadata if metadata
+
+      response = make_request(:patch, uri, payload)
+      handle_response(response, 'update order')
+    end
+
+    # Delete an order and its line items
+    #
+    # @param order_id [String] Order UUID to delete
+    def delete(order_id)
+      uri = URI("#{@base_url}/functions/v1/api-orders")
+      payload = { id: order_id }
+
+      response = make_request(:delete, uri, payload)
+      handle_response(response, 'delete order')
+    end
+
+    # Add a line item to an order
+    #
+    # @param order_id [String] Order UUID (required)
+    # @param description [String] Line item description (required)
+    # @param quantity [Numeric] Quantity (required)
+    # @param unit_price [Numeric] Price per unit (required)
+    # @param product_id [String, nil] Product reference ID
+    # @param product_name [String, nil] Product name
+    # @param product_sku [String, nil] Product SKU
+    # @param usage_metric [String, nil] Usage metric name
+    # @param usage_quantity [Numeric, nil] Usage quantity
+    # @param discount_type [String, nil] "percentage" or "fixed"
+    # @param discount_value [Numeric, nil] Discount amount/percentage
+    # @param tax_rate [Numeric, nil] Tax rate percentage
+    # @param notes [String, nil] Additional notes
+    # @param sort_order [Integer, nil] Display order
+    # @param agent_id [String, nil] Associated agent ID
+    # @param item_type [String, nil] Type: usage, platform_fee, setup_fee, base_fee, overage, custom, product, service, subscription
+    # @param metadata [Hash, nil] Additional metadata
+    # @return [Hash] Created line item object
+    def add_line_item(order_id:, description:, quantity:, unit_price:, product_id: nil, 
+                      product_name: nil, product_sku: nil, usage_metric: nil, usage_quantity: nil,
+                      discount_type: nil, discount_value: nil, tax_rate: nil, notes: nil,
+                      sort_order: nil, agent_id: nil, item_type: nil, metadata: nil)
+      uri = URI("#{@base_url}/functions/v1/api-orders?action=line-items")
+      
+      payload = {
+        order_id: order_id,
+        description: description,
+        quantity: quantity,
+        unit_price: unit_price
+      }
+      payload[:product_id] = product_id if product_id
+      payload[:product_name] = product_name if product_name
+      payload[:product_sku] = product_sku if product_sku
+      payload[:usage_metric] = usage_metric if usage_metric
+      payload[:usage_quantity] = usage_quantity if usage_quantity
+      payload[:discount_type] = discount_type if discount_type
+      payload[:discount_value] = discount_value if discount_value
+      payload[:tax_rate] = tax_rate if tax_rate
+      payload[:notes] = notes if notes
+      payload[:sort_order] = sort_order if sort_order
+      payload[:agent_id] = agent_id if agent_id
+      payload[:item_type] = item_type if item_type
+      payload[:metadata] = metadata if metadata
+
+      response = make_request(:post, uri, payload)
+      handle_response(response, 'add line item')
+    end
+
+    # Remove a line item from an order
+    #
+    # @param line_item_id [String] Line item UUID to delete
+    def remove_line_item(line_item_id)
+      uri = URI("#{@base_url}/functions/v1/api-orders?action=line-items")
+      payload = { id: line_item_id }
+
+      response = make_request(:delete, uri, payload)
+      handle_response(response, 'remove line item')
+    end
+
+    private
+
+    def make_request(method, uri, payload = nil)
+      http = Net::HTTP.new(uri.host, uri.port)
+      http.use_ssl = true
+      http.read_timeout = 30
+
+      case method
+      when :get
+        request = Net::HTTP::Get.new(uri)
+      when :post
+        request = Net::HTTP::Post.new(uri.request_uri)
+        request.body = payload.to_json if payload
+      when :patch
+        request = Net::HTTP::Patch.new(uri.path)
+        request.body = payload.to_json if payload
+      when :delete
+        request = Net::HTTP::Delete.new(uri.request_uri)
+        request.body = payload.to_json if payload
+      end
+
+      request['Content-Type'] = 'application/json'
+      request['X-API-Key'] = @api_key
+
+      http.request(request)
+    end
+
+    def handle_response(response, operation)
+      body = response.body ? JSON.parse(response.body) : {}
+      
+      unless %w[200 201 204].include?(response.code)
+        error_msg = body['error'] || "Failed to #{operation}"
+        raise error_msg
+      end
+
+      puts "[AgentBill] #{operation} successful" if @debug
+      body
+    end
+  end
+end