prompt_warden 0.1.0 → 0.1.1

data/lib/prompt_warden/uploader.rb

@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+require 'faraday'
+require 'faraday/retry'
+require 'singleton'
+require 'tmpdir'
+require 'fileutils'
+
+module PromptWarden
+  class Uploader
+    include Singleton
+
+    def initialize
+      cfg    = PromptWarden.configuration
+      @token = cfg.project_token
+      @logger = cfg.logger
+      api_url = cfg.respond_to?(:api_url) ? cfg.api_url : 'https://httpbin.org/post'
+
+      @client = Faraday.new(url: api_url) do |f|
+        f.request :retry,
+                  max: cfg.max_retries,
+                  interval: 0.05,
+                  backoff_factor: 2,
+                  methods: %i[get post],
+                  retry_statuses: [500, 502, 503]
+        f.adapter :net_http
+      end
+
+      @queue = Queue.new
+      start_worker
+    end
+
+    # Enqueue compressed batch for async upload
+    def enqueue(payload)
+      @queue << payload
+    end
+
+    # Retry all failed uploads from disk
+    def self.retry_failed!
+      files = Dir.glob(File.join(Dir.tmpdir, 'pw_failed_*.json.gz'))
+      logger = PromptWarden.configuration.logger
+
+      if files.empty?
+        logger.info "No failed upload files found to retry"
+        return
+      end
+
+      logger.info "Found #{files.size} failed upload file(s) to retry"
+
+      files.each do |file|
+        begin
+          instance.enqueue(File.binread(file))
+          File.delete(file)
+          logger.info "Successfully retried and deleted: #{File.basename(file)}"
+        rescue => e
+          # If upload fails again, leave the file for next retry
+          logger.error "Retry failed for #{File.basename(file)}: #{e.class}: #{e.message}"
+        end
+      end
+    end
+
+    private
+
+    def start_worker
+      Thread.new do
+        loop do
+          payload = @queue.pop
+          begin
+            resp = @client.post do |req|
+              req.headers['Authorization']     = "Bearer #{@token}"
+              req.headers['Content-Encoding']  = 'gzip'
+              req.body = payload
+            end
+            @logger.info  "PromptWarden upload: #{resp.status}"
+          rescue StandardError => e
+            @logger.error "PromptWarden upload failed: #{e.message}"
+            # Disk fallback: write failed batch to tmpdir
+            begin
+              dir = Dir.tmpdir
+              FileUtils.mkdir_p(dir) unless Dir.exist?(dir)
+              path = File.join(dir, 'pw_failed_') + Time.now.utc.strftime('%Y%m%d%H%M%S%L') + '.json.gz'
+              File.binwrite(path, payload)
+              @logger.info "PromptWarden wrote failed batch to #{path}"
+            rescue StandardError => file_err
+              @logger.error "PromptWarden failed to write fallback file: #{file_err.class}: #{file_err.message}"
+            end
+            # TODO: persist to disk for retry on process restart
+          end
+        end
+      end.tap { |t| t.name = 'PromptWarden::UploaderWorker' }
+    end
+  end
+end

data/lib/prompt_warden/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module PromptWarden
-  VERSION = "0.1.0"
+  VERSION = '0.1.1'
 end

data/lib/prompt_warden.rb

@@ -1,8 +1,37 @@
 # frozen_string_literal: true
 
-require_relative "prompt_warden/version"
+require_relative 'prompt_warden/version'
+require_relative 'prompt_warden/configuration'
+require_relative 'prompt_warden/railtie' if defined?(Rails)
+require_relative 'prompt_warden/event'
+require_relative 'prompt_warden/buffer'
+require_relative 'prompt_warden/uploader'
+require_relative 'prompt_warden/policy'
+require_relative 'prompt_warden/adapter'
+require_relative 'prompt_warden/cli'
+require_relative 'prompt_warden/cost_calculator'
 
 module PromptWarden
-  class Error < StandardError; end
-  # Your code goes here...
+  class << self
+    def configuration    = (@configuration ||= Configuration.new)
+    def configure        = yield(configuration).tap { configuration.validate! }
+    def reset!           = (@configuration = nil)
+
+    def record(event_attrs)
+      buffer.push(Event.new(**event_attrs))
+    end
+
+    def calculate_cost(prompt:, model:, response_tokens: nil)
+      CostCalculator.calculate_cost(prompt: prompt, model: model, response_tokens: response_tokens)
+    end
+
+    private
+
+    def buffer
+      @buffer ||= Buffer.new
+    end
+  end
 end
+
+# Auto‑load adapters after core is ready
+PromptWarden::Adapter.auto_load_all!

data/prompt_warden.gemspec

@@ -1,42 +1,50 @@
-# prompt_warden.gemspec
-require_relative "lib/prompt_warden/version"
+# frozen_string_literal: true
+
+require_relative 'lib/prompt_warden/version'
 
 Gem::Specification.new do |spec|
   # ---- Core identity --------------------------------------------------
-  spec.name          = "prompt_warden"
-  spec.version       = PromptWarden::VERSION        # keep in one place
-  spec.authors       = ["Tyler Hammett"]
-  spec.email         = ["hello@promptwarden.io"]
-  spec.summary       = "Rails-native guard-rails & cost logging for LLM calls."
+  spec.name          = 'prompt_warden'
+  spec.version       = PromptWarden::VERSION # keep in one place
+  spec.authors       = ['Tyler Hammett']
+  spec.email         = ['hello@promptwarden.io']
+  spec.summary       = 'Record, audit, and guard AI prompt usage with automatic SDK instrumentation.'
   spec.description   = <<~DESC
-    PromptWarden watches every OpenAI / Anthropic request,
-    blocks risky prompts, and uploads JSON batches to the PromptWarden SaaS.
+    PromptWarden provides automatic instrumentation for OpenAI, Anthropic, and Langchain SDKs.
+    Features include policy enforcement, cost calculation, real-time monitoring, and alert recording.
+    Includes CLI tool for live event streaming and comprehensive test coverage.
   DESC
-  spec.license       = "MIT"
-  spec.homepage      = "https://promptwarden.io"
+  spec.license       = 'MIT'
+  spec.homepage      = 'https://promptwarden.io'
 
   # ---- Compatibility ---------------------------------------------------
-  spec.required_ruby_version = ">= 3.1"
+  spec.required_ruby_version = '>= 3.1'
 
   # ---- Files & executables --------------------------------------------
-  spec.files         = Dir.chdir(File.expand_path(__dir__)) do
-                         `git ls-files -z`.split("\x0")
-                       end
-  spec.bindir        = "exe"          # if you later add CLI binstubs
-  spec.executables   = []             # none yet
-  spec.require_paths = ["lib"]
+  spec.files = Dir.chdir(File.expand_path(__dir__)) do
+    `git ls-files -z`.split("\x0")
+  end
+  spec.bindir        = 'bin'
+  spec.executables   = ['pw_tail']
+  spec.require_paths = ['lib']
 
   # ---- Runtime dependencies -------------------------------------------
   # HTTP client for uploads
-  spec.add_runtime_dependency "faraday",      "~> 2.9"
+  spec.add_runtime_dependency 'faraday', '~> 2.9'
+  spec.add_runtime_dependency 'faraday-retry', '~> 2.2'
+
+  # YAML for policy configuration
+  spec.add_runtime_dependency 'yaml'
 
   # ActiveSupport for Notifications
-  spec.add_runtime_dependency "activesupport", ">= 6.1"
+  spec.add_runtime_dependency 'activesupport', '>= 6.1'
 
   # ---- Development / test dependencies -------------------------------
-  spec.add_development_dependency "rspec",          "~> 3.12"
-  spec.add_development_dependency "rubocop",        "~> 1.60"
-  spec.add_development_dependency "rubocop-rspec",  "~> 2.26"
-  spec.add_development_dependency "bundler",        ">= 2.3.0"
-  spec.add_development_dependency "rake",           "~> 13.2"
+  spec.add_development_dependency 'bundler',        '>= 2.3.0'
+  spec.add_development_dependency 'rake',           '~> 13.2'
+  spec.add_development_dependency 'rspec',          '~> 3.12'
+  spec.add_development_dependency 'rubocop',        '~> 1.60'
+  spec.add_development_dependency 'rubocop-rspec',  '~> 2.26'
+  spec.add_development_dependency 'timecop',        '~> 0.9'
+  spec.add_development_dependency 'webmock',        '~> 3.20'
 end

data/spec/adapter_auto_detect_spec.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+# --- Helper to stub an SDK constant --------------------------
+module OpenAI
+  class Client
+    def initialize(api_key: nil, **opts)
+      @api_key = api_key
+    end
+
+    def chat(*); { "choices" => [{ "message" => { "content" => "hi" } }] }; end
+  end
+end
+
+def stub_openai_constant
+  # No-op: the constant is now always defined
+  PromptWarden::Adapter.auto_load_all!
+  if defined?(PromptWarden::Instrumentation::OpenAI)
+    OpenAI::Client.prepend(PromptWarden::Instrumentation::OpenAI)
+  end
+end
+
+RSpec.describe 'PromptWarden adapter auto‑detection' do
+  # ------------------------------------------------------------------
+  it 'loads adapter when gem is in Gem.loaded_specs' do
+    Gem.loaded_specs['openai'] = Gem::Specification.new('openai') # simulate Gemfile
+
+    require 'prompt_warden'                                       # auto‑loader runs
+    stub_openai_constant                                          # constant after load
+
+    PromptWarden.configure { |c| c.project_token = 'tok' }
+    expect(OpenAI::Client.ancestors)
+      .to include(PromptWarden::Instrumentation::OpenAI)
+  ensure
+    Gem.loaded_specs.delete('openai')
+  end
+
+  # ------------------------------------------------------------------
+  it 'loads adapter when constant exists (no gem spec)' do
+    stub_openai_constant                                          # constant first
+    require 'prompt_warden'                                       # auto‑loader sees it
+
+    PromptWarden.configure { |c| c.project_token = 'tok' }
+    client = OpenAI::Client.new
+    expect(client).to respond_to(:chat)                           # wrapper prepended
+  end
+
+  # ------------------------------------------------------------------
+  it 'executes user‑registered adapter via register_adapter' do
+    require 'prompt_warden'
+    called = false
+
+    PromptWarden.configure do |c|
+      c.project_token = 'tok'
+      c.register_adapter(:fake_sdk) { called = true }
+    end
+
+    # Fake gem appears later in runtime
+    Gem.loaded_specs['fake_sdk'] = Gem::Specification.new('fake_sdk')
+    PromptWarden.configuration.run_pending_adapters
+
+    expect(called).to be true
+  ensure
+    Gem.loaded_specs.delete('fake_sdk')
+  end
+end

data/spec/anthropic_adapter_spec.rb

@@ -0,0 +1,137 @@
+# frozen_string_literal: true
+
+require_relative '../lib/prompt_warden/instrumentation/anthropic'
+
+RSpec.describe 'PromptWarden Anthropic adapter' do
+  before do
+    require 'prompt_warden'
+    PromptWarden.configure { |c| c.project_token = 'tok' }
+    allow(PromptWarden).to receive(:record)
+    allow(PromptWarden::Policy.instance).to receive(:check_alerts).and_return([])
+    allow(PromptWarden).to receive(:calculate_cost).and_return(0.005)
+  end
+
+  it 'records prompt usage via instrumentation' do
+    client = Object.new
+    client.extend(PromptWarden::Instrumentation::Anthropic)
+    # Call the instrumentation method directly
+    client.send(:_pw_wrap, :complete, {
+      prompt: 'hi',
+      model: 'claude-3-sonnet-20240229',
+      max_tokens: 1000
+    }) do
+      { 'completion' => 'anthropic reply' }
+    end
+    expect(PromptWarden).to have_received(:record).with(hash_including(
+      prompt: 'hi',
+      model: 'claude-3-sonnet-20240229',
+      cost_usd: 0.005,
+      status: 'ok',
+      alerts: []
+    ))
+  end
+
+  it 'includes alerts when policy alerts are detected' do
+    alerts = [{ type: 'regex', rule: '/confidential/i', level: 'warn' }]
+    allow(PromptWarden::Policy.instance).to receive(:check_alerts).and_return(alerts)
+    client = Object.new
+    client.extend(PromptWarden::Instrumentation::Anthropic)
+    # Call the instrumentation method directly
+    client.send(:_pw_wrap, :complete, {
+      prompt: 'This is confidential information',
+      model: 'claude-3-sonnet-20240229'
+    }) do
+      { 'completion' => 'anthropic reply' }
+    end
+    expect(PromptWarden).to have_received(:record).with(hash_including(
+      prompt: 'This is confidential information',
+      model: 'claude-3-sonnet-20240229',
+      status: 'ok',
+      alerts: alerts
+    ))
+  end
+
+  it 'records failed status on error with alerts' do
+    alerts = [{ type: 'regex', rule: '/ETA/i', level: 'warn' }]
+    allow(PromptWarden::Policy.instance).to receive(:check_alerts).and_return(alerts)
+    client = Object.new
+    client.extend(PromptWarden::Instrumentation::Anthropic)
+    expect {
+      client.send(:_pw_wrap, :complete, {
+        prompt: 'What is the ETA?',
+        model: 'claude-3-sonnet-20240229'
+      }) do
+        raise 'test error'
+      end
+    }.to raise_error('test error')
+    expect(PromptWarden).to have_received(:record).with(hash_including(
+      prompt: 'What is the ETA?',
+      model: 'claude-3-sonnet-20240229',
+      status: 'failed',
+      alerts: alerts
+    ))
+  end
+
+  it 'handles messages format for Anthropic API' do
+    client = Object.new
+    client.extend(PromptWarden::Instrumentation::Anthropic)
+
+    client.send(:_pw_wrap, :messages, {
+      messages: [
+        { role: 'user', content: 'Hello' },
+        { role: 'assistant', content: 'Hi there' },
+        { role: 'user', content: 'How are you?' }
+      ],
+      model: 'claude-3-sonnet-20240229'
+    }) do
+      {
+        'content' => [{ 'type' => 'text', 'text' => 'anthropic reply' }],
+        'usage' => { 'output_tokens' => 30 }
+      }
+    end
+
+    expect(PromptWarden).to have_received(:record).with(hash_including(
+      prompt: "Hello\nHi there\nHow are you?",
+      model: 'claude-3-sonnet-20240229',
+      status: 'ok',
+      alerts: []
+    ))
+  end
+
+  it 'uses enhanced cost calculation for estimates and actual costs' do
+    client = Object.new
+    client.extend(PromptWarden::Instrumentation::Anthropic)
+
+    # Mock cost calculation calls
+    allow(PromptWarden).to receive(:calculate_cost).with(
+      prompt: 'hi',
+      model: 'claude-3-sonnet-20240229'
+    ).and_return(0.003)  # Estimate
+
+    allow(PromptWarden).to receive(:calculate_cost).with(
+      prompt: 'hi',
+      model: 'claude-3-sonnet-20240229',
+      response_tokens: 30
+    ).and_return(0.005)  # Actual cost
+
+    client.send(:_pw_wrap, :complete, {
+      prompt: 'hi',
+      model: 'claude-3-sonnet-20240229'
+    }) do
+      {
+        'completion' => 'anthropic reply',
+        'usage' => { 'output_tokens' => 30 }
+      }
+    end
+
+    expect(PromptWarden).to have_received(:calculate_cost).with(
+      prompt: 'hi',
+      model: 'claude-3-sonnet-20240229'
+    )
+    expect(PromptWarden).to have_received(:calculate_cost).with(
+      prompt: 'hi',
+      model: 'claude-3-sonnet-20240229',
+      response_tokens: 30
+    )
+  end
+end

data/spec/buffer_spec.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require 'webmock/rspec'
+require 'timecop'
+
+RSpec.describe PromptWarden::Buffer do
+  let(:config) do
+    PromptWarden::Configuration.new.tap do |c|
+      c.project_token  = 'test'
+      c.flush_interval = 10 # long, we’ll call flush! manually
+      c.batch_bytes    = 1_000
+    end
+  end
+
+  before { PromptWarden.reset! } # ensure fresh singleton
+
+  # spec/buffer_spec.rb fragment
+  it 'enqueues events and flushes when batch_bytes exceeded' do
+    uploader = instance_double(PromptWarden::Uploader, enqueue: true)
+    allow(PromptWarden::Uploader).to receive(:instance).and_return(uploader)
+
+    config.batch_bytes = 1_000 # raise threshold
+    buffer = described_class.new(config)
+
+    3.times { buffer.push(PromptWarden::Event.new(prompt: 'hi').to_h) }
+
+    buffer.flush! # ensure final flush for deterministic spec
+
+    expect(uploader).to have_received(:enqueue).once
+  end
+
+  it 'timer flushes automatically after interval' do
+    uploader = instance_double(PromptWarden::Uploader, enqueue: true)
+    allow(PromptWarden::Uploader).to receive(:instance).and_return(uploader)
+
+    # use very small interval so we can wait for real time
+    config.flush_interval = 0.05
+    buffer = described_class.new(config)
+    buffer.push(PromptWarden::Event.new(prompt: 'auto').to_h)
+
+    sleep 0.1 # two ticks of interval
+    expect(uploader).to have_received(:enqueue).at_least(:once)
+  end
+end