prompt_warden 0.1.0 → 0.1.1

data/spec/openai_adapter_spec.rb

@@ -0,0 +1,153 @@
+# frozen_string_literal: true
+
+require_relative '../lib/prompt_warden/instrumentation/openai'
+
+RSpec.describe 'PromptWarden OpenAI 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 for chat method' do
+    client = Object.new
+    client.extend(PromptWarden::Instrumentation::OpenAI)
+    # Call the instrumentation method directly
+    client.send(:_pw_wrap, :chat, {
+      messages: [{ role: 'user', content: 'hi' }],
+      model: 'gpt-4o',
+      max_tokens: 1000
+    }) do
+      { 'choices' => [{ 'message' => { 'content' => 'openai reply' } }] }
+    end
+    expect(PromptWarden).to have_received(:record).with(hash_including(
+      prompt: 'hi',
+      model: 'gpt-4o',
+      cost_usd: 0.005,
+      status: 'ok',
+      alerts: []
+    ))
+  end
+
+  it 'records prompt usage via instrumentation for completions method' do
+    client = Object.new
+    client.extend(PromptWarden::Instrumentation::OpenAI)
+    # Call the instrumentation method directly
+    client.send(:_pw_wrap, :completions, {
+      prompt: 'hello',
+      model: 'gpt-3.5-turbo',
+      max_tokens: 500
+    }) do
+      { 'choices' => [{ 'message' => { 'content' => 'openai completion' } }] }
+    end
+    expect(PromptWarden).to have_received(:record).with(hash_including(
+      prompt: 'hello',
+      model: 'gpt-3.5-turbo',
+      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::OpenAI)
+    # Call the instrumentation method directly
+    client.send(:_pw_wrap, :chat, {
+      messages: [{ role: 'user', content: 'This is confidential information' }],
+      model: 'gpt-4o'
+    }) do
+      { 'choices' => [{ 'message' => { 'content' => 'openai reply' } }] }
+    end
+    expect(PromptWarden).to have_received(:record).with(hash_including(
+      prompt: 'This is confidential information',
+      model: 'gpt-4o',
+      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::OpenAI)
+    expect {
+      client.send(:_pw_wrap, :chat, {
+        messages: [{ role: 'user', content: 'What is the ETA?' }],
+        model: 'gpt-4o'
+      }) 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: 'gpt-4o',
+      status: 'failed',
+      alerts: alerts
+    ))
+  end
+
+  it 'extracts prompt from messages array' do
+    client = Object.new
+    client.extend(PromptWarden::Instrumentation::OpenAI)
+    # Call the instrumentation method directly
+    client.send(:_pw_wrap, :chat, {
+      messages: [
+        { role: 'user', content: 'Hello' },
+        { role: 'assistant', content: 'Hi there' },
+        { role: 'user', content: 'How are you?' }
+      ],
+      model: 'gpt-4o'
+    }) do
+      { 'choices' => [{ 'message' => { 'content' => 'openai reply' } }] }
+    end
+    expect(PromptWarden).to have_received(:record).with(hash_including(
+      prompt: "Hello\nHi there\nHow are you?",
+      model: 'gpt-4o',
+      status: 'ok',
+      alerts: []
+    ))
+  end
+
+  it 'uses enhanced cost calculation for estimates and actual costs' do
+    client = Object.new
+    client.extend(PromptWarden::Instrumentation::OpenAI)
+
+    # Mock cost calculation calls
+    allow(PromptWarden).to receive(:calculate_cost).with(
+      prompt: 'hi',
+      model: 'gpt-4o'
+    ).and_return(0.003)  # Estimate
+
+    allow(PromptWarden).to receive(:calculate_cost).with(
+      prompt: 'hi',
+      model: 'gpt-4o',
+      response_tokens: 50
+    ).and_return(0.005)  # Actual cost
+
+    client.send(:_pw_wrap, :chat, {
+      messages: [{ role: 'user', content: 'hi' }],
+      model: 'gpt-4o'
+    }) do
+      {
+        'choices' => [{ 'message' => { 'content' => 'openai reply' } }],
+        'usage' => { 'completion_tokens' => 50 }
+      }
+    end
+
+    expect(PromptWarden).to have_received(:calculate_cost).with(
+      prompt: 'hi',
+      model: 'gpt-4o'
+    )
+    expect(PromptWarden).to have_received(:calculate_cost).with(
+      prompt: 'hi',
+      model: 'gpt-4o',
+      response_tokens: 50
+    )
+  end
+end

data/spec/policy_spec.rb

@@ -0,0 +1,170 @@
+# frozen_string_literal: true
+
+require 'tmpdir'
+require 'yaml'
+
+RSpec.describe PromptWarden::Policy do
+  let(:policy_path) { File.join(Dir.mktmpdir, 'policy.yml') }
+
+  def write_rules(rules)
+    File.write(policy_path, rules.to_yaml)
+    ENV['PROMPT_WARDEN_POLICY'] = policy_path
+    PromptWarden::Policy.instance.reload!
+  end
+
+  after { ENV.delete('PROMPT_WARDEN_POLICY') }
+
+  # -----------------------------------------------------------
+  it 'blocks when cost exceeds max_cost_usd' do
+    write_rules(max_cost_usd: 0.01)
+
+    policy = described_class.instance
+    expect do
+      policy.check!(prompt: 'hi', cost_estimate: 0.02)
+    end.to raise_error(PromptWarden::PolicyError, /exceeds/)
+  end
+
+  # -----------------------------------------------------------
+  it 'blocks when prompt matches reject regex' do
+    write_rules(reject_if_regex: ['/password/i'])
+
+    policy = described_class.instance
+    expect do
+      policy.check!(prompt: 'My password is 123', cost_estimate: 0)
+    end.to raise_error(PromptWarden::PolicyError, /reject regex/)
+  end
+
+  # -----------------------------------------------------------
+  it 'returns :ok for compliant prompt' do
+    write_rules(max_cost_usd: 1, reject_if_regex: [])
+
+    policy = described_class.instance
+    result = policy.check!(prompt: 'hello', cost_estimate: 0.5)
+    expect(result).to eq :ok
+  end
+
+  # -----------------------------------------------------------
+  describe '#check_alerts' do
+    it 'returns empty array when no alert rules configured' do
+      write_rules({})
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'any text', cost_estimate: 0.1)
+      expect(alerts).to eq []
+    end
+
+    it 'returns empty array when prompt does not match warn regex' do
+      write_rules(warn_if_regex: ['/password/i'])
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'hello world', cost_estimate: 0.1)
+      expect(alerts).to eq []
+    end
+
+    it 'returns alert when prompt matches warn regex' do
+      write_rules(warn_if_regex: ['/password/i'])
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'My password is 123', cost_estimate: 0.1)
+      expect(alerts).to eq [{ type: 'regex', rule: '/password/i', level: 'warn' }]
+    end
+
+    it 'returns multiple alerts when prompt matches multiple warn regexes' do
+      write_rules(warn_if_regex: ['/password/i', '/ssn/i'])
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'My password is 123 and SSN is 456', cost_estimate: 0.1)
+      expect(alerts).to contain_exactly(
+        { type: 'regex', rule: '/password/i', level: 'warn' },
+        { type: 'regex', rule: '/ssn/i', level: 'warn' }
+      )
+    end
+
+    it 'returns cost alert when cost exceeds limit' do
+      write_rules(max_cost_usd: 0.05)
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'hello', cost_estimate: 0.1)
+      expect(alerts).to eq [{ type: 'cost', limit: 0.05, level: 'block' }]
+    end
+
+    it 'returns both regex and cost alerts when both conditions are met' do
+      write_rules(max_cost_usd: 0.05, warn_if_regex: ['/password/i'])
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'My password is 123', cost_estimate: 0.1)
+      expect(alerts).to contain_exactly(
+        { type: 'regex', rule: '/password/i', level: 'warn' },
+        { type: 'cost', limit: 0.05, level: 'block' }
+      )
+    end
+
+    it 'handles case insensitive regex flags' do
+      write_rules(warn_if_regex: ['/ETA/i'])
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'The eta is 5 minutes', cost_estimate: 0.1)
+      expect(alerts).to eq [{ type: 'regex', rule: '/ETA/i', level: 'warn' }]
+    end
+
+    it 'handles case sensitive regex when no flags' do
+      write_rules(warn_if_regex: ['/ETA'])
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'The eta is 5 minutes', cost_estimate: 0.1)
+      expect(alerts).to eq []
+    end
+
+    it 'handles regex with multiple flags' do
+      write_rules(warn_if_regex: ['/eta/im'])
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: "ETA\nis 5 minutes", cost_estimate: 0.1)
+      expect(alerts).to eq [{ type: 'regex', rule: '/eta/im', level: 'warn' }]
+    end
+
+    it 'handles plain string patterns' do
+      write_rules(warn_if_regex: ['password'])
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'My password is 123', cost_estimate: 0.1)
+      expect(alerts).to eq [{ type: 'regex', rule: 'password', level: 'warn' }]
+    end
+
+    it 'handles empty warn_if_regex array' do
+      write_rules(warn_if_regex: [])
+
+      policy = described_class.instance
+      alerts = policy.check_alerts(prompt: 'any text', cost_estimate: 0.1)
+      expect(alerts).to eq []
+    end
+  end
+
+  # -----------------------------------------------------------
+  describe 'integration with check!' do
+    it 'allows check! to pass even when alerts are present' do
+      write_rules(
+        max_cost_usd: 1,
+        reject_if_regex: [],
+        warn_if_regex: ['/password/i']
+      )
+
+      policy = described_class.instance
+      result = policy.check!(prompt: 'My password is 123', cost_estimate: 0.5)
+      expect(result).to eq :ok
+    end
+
+    it 'still blocks on reject regex even when alerts are present' do
+      write_rules(
+        max_cost_usd: 1,
+        reject_if_regex: ['/ssn/i'],
+        warn_if_regex: ['/password/i']
+      )
+
+      policy = described_class.instance
+      expect do
+        policy.check!(prompt: 'My password is 123 and SSN is 456', cost_estimate: 0.5)
+      end.to raise_error(PromptWarden::PolicyError, /reject regex/)
+    end
+  end
+end

data/spec/prompt_warden_spec.rb

@@ -3,11 +3,11 @@
 RSpec.describe PromptWarden do
   subject(:mod) { described_class }
 
-  it "has a semantic version number" do
+  it 'has a semantic version number' do
     expect(PromptWarden::VERSION).to match(/\A\d+\.\d+\.\d+\z/)
   end
 
-  it "is defined as a module" do
+  it 'is defined as a module' do
     expect(mod).to be_a(Module)
   end
 end

data/spec/spec_helper.rb

@@ -1,15 +1,14 @@
 # frozen_string_literal: true
 
-require "prompt_warden"
+require_relative '../lib/prompt_warden'
+require 'webmock/rspec'
+WebMock.disable_net_connect!(allow_localhost: true)
 
 RSpec.configure do |config|
-  # Enable flags like --only-failures and --next-failure
-  config.example_status_persistence_file_path = ".rspec_status"
-
-  # Disable RSpec exposing methods globally on `Module` and `main`
+  config.example_status_persistence_file_path = '.rspec_status'
   config.disable_monkey_patching!
+  config.expect_with(:rspec) { |c| c.syntax = :expect }
 
-  config.expect_with :rspec do |c|
-    c.syntax = :expect
-  end
+  # fresh configuration between examples
+  config.before { PromptWarden.reset! if defined?(PromptWarden) }
 end

data/spec/uploader_spec.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+require 'webmock/rspec'
+
+RSpec.describe PromptWarden::Uploader do
+  let(:api_url) { 'https://example.com/ingest' }
+  let(:config) do
+    PromptWarden::Configuration.new.tap do |c|
+      c.project_token = 'tok_123'
+      c.api_url       = api_url
+      c.max_retries   = 2
+    end
+  end
+
+  before do
+    PromptWarden.reset!
+    PromptWarden.configure do |c|
+      c.project_token = config.project_token
+      c.api_url       = config.api_url
+      c.max_retries   = config.max_retries
+    end
+  end
+
+  after do
+    # Clean up any failed files in tmpdir
+    Dir.glob(File.join(Dir.tmpdir, 'pw_failed_*.json.gz')).each { |f| File.delete(f) }
+  end
+
+  it 'POSTs gzip payload with auth header' do
+    stub = stub_request(:post, api_url)
+           .with(headers: { 'Authorization' => 'Bearer tok_123',
+                            'Content-Encoding' => 'gzip' })
+           .to_return(status: 200)
+
+    uploader = described_class.instance
+    uploader.enqueue('compressed bytes')
+
+    # give worker thread a tick
+    sleep 0.05
+
+    expect(stub).to have_been_requested
+  end
+
+  it 'retries once on 500 then succeeds' do
+    stub_request(:post, api_url)
+      .to_return({ status: 500 }, { status: 200 })
+
+    uploader = described_class.instance
+    uploader.enqueue('data')
+
+    # wait long enough for retry thread (interval=0.05, backoff 2x)
+    sleep 0.2
+
+    expect(a_request(:post, api_url)).to have_been_made.times(2)
+  end
+
+  it 'writes failed batch to disk on network error' do
+    # Simulate network error
+    stub_request(:post, api_url).to_raise(Faraday::ConnectionFailed.new('fail'))
+    uploader = described_class.instance
+    uploader.enqueue('diskfail')
+    sleep 0.05
+    failed_files = Dir.glob(File.join(Dir.tmpdir, 'pw_failed_*.json.gz'))
+    expect(failed_files.size).to be >= 1
+    expect(File.binread(failed_files.first)).to eq 'diskfail'
+  end
+
+  it 'pw:retry_failed task uploads and deletes failed files' do
+    # Write a failed file manually
+    path = File.join(Dir.tmpdir, 'pw_failed_test.json.gz')
+    File.binwrite(path, 'retryme')
+    stub = stub_request(:post, api_url).to_return(status: 200)
+    # Run the rake task logic directly
+    PromptWarden::Uploader.retry_failed!
+    sleep 0.05
+    expect(stub).to have_been_requested
+    expect(File.exist?(path)).to be false
+  end
+end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: prompt_warden
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.1.1
 platform: ruby
 authors:
 - Tyler Hammett
 autorequire:
-bindir: exe
+bindir: bin
 cert_chain: []
-date: 2025-07-22 00:00:00.000000000 Z
+date: 2025-07-24 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: faraday
@@ -24,6 +24,34 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: '2.9'
+- !ruby/object:Gem::Dependency
+  name: faraday-retry
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.2'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.2'
+- !ruby/object:Gem::Dependency
+  name: yaml
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 - !ruby/object:Gem::Dependency
   name: activesupport
   requirement: !ruby/object:Gem::Requirement
@@ -38,6 +66,34 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '6.1'
+- !ruby/object:Gem::Dependency
+  name: bundler
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 2.3.0
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 2.3.0
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.2'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.2'
 - !ruby/object:Gem::Dependency
   name: rspec
   requirement: !ruby/object:Gem::Requirement
@@ -81,39 +137,41 @@ dependencies:
       - !ruby/object:Gem::Version
         version: '2.26'
 - !ruby/object:Gem::Dependency
-  name: bundler
+  name: timecop
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: 2.3.0
+        version: '0.9'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - ">="
+    - - "~>"
       - !ruby/object:Gem::Version
-        version: 2.3.0
+        version: '0.9'
 - !ruby/object:Gem::Dependency
-  name: rake
+  name: webmock
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '13.2'
+        version: '3.20'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '13.2'
+        version: '3.20'
 description: |
-  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.
 email:
 - hello@promptwarden.io
-executables: []
+executables:
+- pw_tail
 extensions: []
 extra_rdoc_files: []
 files:
@@ -128,13 +186,38 @@ files:
 - README.md
 - Rakefile
 - bin/console
+- bin/pw_tail
 - bin/setup
+- examples/policy.yml
 - lib/prompt_warden.rb
+- lib/prompt_warden/adapter.rb
+- lib/prompt_warden/buffer.rb
+- lib/prompt_warden/cli.rb
+- lib/prompt_warden/configuration.rb
+- lib/prompt_warden/cost_calculator.rb
+- lib/prompt_warden/event.rb
+- lib/prompt_warden/instrumentation/anthropic.rb
+- lib/prompt_warden/instrumentation/langchain.rb
+- lib/prompt_warden/instrumentation/openai.rb
+- lib/prompt_warden/policy.rb
+- lib/prompt_warden/railtie.rb
+- lib/prompt_warden/uploader.rb
 - lib/prompt_warden/version.rb
 - prompt_warden.gemspec
 - sig/prompt_warden.rbs
+- spec/adapter_auto_detect_spec.rb
+- spec/anthropic_adapter_spec.rb
+- spec/buffer_spec.rb
+- spec/cli_spec.rb
+- spec/configuration_spec.rb
+- spec/cost_calculator_spec.rb
+- spec/event_spec.rb
+- spec/langchain_adapter_spec.rb
+- spec/openai_adapter_spec.rb
+- spec/policy_spec.rb
 - spec/prompt_warden_spec.rb
 - spec/spec_helper.rb
+- spec/uploader_spec.rb
 homepage: https://promptwarden.io
 licenses:
 - MIT
@@ -157,5 +240,5 @@ requirements: []
 rubygems_version: 3.4.19
 signing_key:
 specification_version: 4
-summary: Rails-native guard-rails & cost logging for LLM calls.
+summary: Record, audit, and guard AI prompt usage with automatic SDK instrumentation.
 test_files: []