Refine prompt, allow multiple assessments and add table to edit doc page

Author: embarnard

app/assets/stylesheets/application.scss

@@ -33,6 +33,7 @@
 @import "components/client-filters";
 @import "components/doc-help";
 @import "components/hub-form";
+@import "components/hub-doc-assessment";
 @import "components/illustrations";
 @import "components/progress-indicator";
 @import "components/marketing-comparison";

app/assets/stylesheets/components/_hub-doc-assessment.scss

@@ -0,0 +1,31 @@
+#doc-assessments-table {
+  margin: 10rem 5rem 5rem 5rem;
+
+
+  table {
+    border-collapse: collapse;
+    font-size: 12px;
+  }
+
+  tr.error {
+    background-color: #FEF2E7;
+  }
+
+  tr.fail {
+    background-color: #FEE7E7;
+  }
+
+  tr.pass {
+    background-color: #E7FEEC;
+  }
+
+  td, th {
+    border: 1px solid black;
+    padding: 10px 20px;
+  }
+
+  th {
+    font-weight: bold;
+    text-transform: capitalize;
+  }
+}

app/controllers/hub/documents_controller.rb

@@ -66,6 +66,14 @@ def destroy
       redirect_back(fallback_location: hub_client_documents_path)
     end
 
+    def rerun_screener
+      return head :forbidden if acts_like_production? || !current_user.admin?
+
+      DocScreenerJob.perform_now(@document.id)
+
+      redirect_back(fallback_location: edit_hub_client_document_path(client_id: @document.client.id, id: @document), notice: "Re-ran document screening.")
+    end
+
     private
 
     def load_document_type_options

app/jobs/doc_screener_job.rb

@@ -1,5 +1,7 @@
 class DocScreenerJob < ApplicationJob
   def perform(document_id)
+    return if Flipper.enabled?(:disable_ai_doc_screener)
+
     document = Document.find(document_id)
     return unless document.upload.attached?
 
@@ -9,7 +11,13 @@ def perform(document_id)
       input_blob_id: document.upload.blob_id
     )
 
-    return if assessment.status == "complete"
+    if assessment.status == "complete"
+      assessment = DocAssessment.create!(
+        document_id: document.id,
+        prompt_version: BedrockDocScreener::PROMPT_VERSION,
+        input_blob_id: document.upload.blob_id
+      )
+    end
 
     assessment.update!(
       status: "processing",

app/models/doc_assessment.rb

@@ -16,8 +16,7 @@
 #
 # Indexes
 #
-#  index_doc_assessments_on_doc_and_version_and_blob  (document_id,prompt_version,input_blob_id) UNIQUE
-#  index_doc_assessments_on_document_id               (document_id)
+#  index_doc_assessments_on_document_id  (document_id)
 #
 # Foreign Keys
 #

app/models/document.rb

@@ -44,6 +44,8 @@ class Document < ApplicationRecord
   belongs_to :contact_record, polymorphic: true, optional: true
   belongs_to :tax_return, optional: true
   belongs_to :uploaded_by, polymorphic: true, optional: true
+  has_many :assessments, class_name: "DocAssessment", dependent: :destroy
+  has_one :latest_assessment, -> { order(created_at: :desc) }, class_name: "DocAssessment"
 
   validates_presence_of :client
   validates_presence_of :upload
@@ -88,6 +90,7 @@ class Document < ApplicationRecord
   end
   after_save_commit { SearchIndexer.refresh_filterable_properties([client_id]) }
   after_destroy_commit { SearchIndexer.refresh_filterable_properties([client_id]) }
+  after_update_commit :rerun_screener_if_document_type_changed, if: :saved_change_to_document_type?
 
   # has_one_attached needs to be called after defining any callbacks that access attachments, like
   # the HEIC conversion; see https://github.com/rails/rails/issues/37304
@@ -150,6 +153,12 @@ def uploaded_by_name_label
 
   private
 
+  def rerun_screener_if_document_type_changed
+    return unless upload.attached?
+
+    DocScreenerJob.perform_later(id)
+  end
+
   def tax_return_belongs_to_client
     errors.add(:tax_return_id, I18n.t("forms.errors.tax_return_belongs_to_client")) unless tax_return.blank? || tax_return.client == client
   end

app/services/bedrock_doc_screener.rb

@@ -1,53 +1,80 @@
 require 'aws-sdk-bedrockruntime'
 require 'json'
 require 'base64'
+require "mini_magick"
+require "tempfile"
 
 module BedrockDocScreener
   MODEL_ID = 'us.anthropic.claude-haiku-4-5-20251001-v1:0'.freeze
   REGION = 'us-east-1'.freeze
   SUPPORTED_MEDIA_TYPES = %w[
     image/png
     image/jpeg
+    application/pdf
   ].freeze
-  # since the payload is `type: "image"` it won't work for application/pdf
-  # can convert the pdf to image or use another bedrock flow for docs
 
-  # Update the prompt version if you are updating the prompt
   PROMPT_VERSION = "v1".freeze
+
   def self.prompt_for(document_type:)
     <<~PROMPT
-      You are validating an uploaded client document.
-
-      Document type: #{document_type}
+      Clients are uploading documents and you need to verify the validity of the document using these rules:
+         1) If the photo is a poor quality image (poorly lit, blurry, cropped & missing information, pixelated screen etc.) 
+            so much so that it renders the document illegible, 
+            then set reason="unreadable" and set verdict="fail".
+         2) If the document does not fit any of the doc types in the available-doc-types list,
+            then set reason="no_doc_type_match" and set verdict="fail".
+         3) If it does not appear to match the stated doc type (in this case #{document_type}) 
+            but does match another type in the available-doc-types list,
+            then set reason="wrong_document_type", verdict="fail" 
+            and include the doc-types that might be match in the explanation field by their label name.
+         4) If the document is expired, then set reason="expired" and verdict="fail".
+         5) If the document is fake (for example if it is labeled as a 'sample'),
+            then set reason="fake" and verdict="fail".
+         6) If there is another reason that the document is not valid, 
+            then set reason="other" and verdict="fail".
+         7) If the document seems to be a valid document, readable and the selected document type matches,
+            then set reason="" and verdict="pass".
+         8) "confidence" must be between 0.0 and 1.0.
+         9) Do not include any keys other than "verdict", "reason", "explanation" and "confidence"
+         10) verdict should only be "pass" or "fail"
+      
+      available-doc-types: #{available_doc_types}
+
+      Selected document type: #{document_type}
 
       Return ONLY valid JSON with this exact schema:
-
       {
-        "verdict": "pass" | "fail" | "needs_review",
-        "reasons": [string, brief explanation],
+        "verdict": "pass" | "fail",
+        "reason": "unreadable" | "no_doc_type_match" | "wrong_document_type" | expired" | "fake" | "other",
+        "explanation": [Brief 1-2 sentence explanation of reason. Explain why valid/invalid.],
         "confidence": number between 0.0-1.0,
       }
-
-      Rules:
-      - If the document is unreadable, set verdict="needs_review" and include reason "unreadable".
-      - If it does not appear to match the stated document type, verdict="fail" and include reason "wrong_document_type".
-      - If it appears valid and readable, verdict="pass".
-      - "confidence" must be between 0.0 and 1.0.
-      - Do not include any keys other than verdict, reasons and confidence
     PROMPT
   end
 
+  def self.available_doc_types
+    # matches @doc_type_options in document controller
+    available_doc_types = [DocumentTypes::Identity, DocumentTypes::SsnItin] + (DocumentTypes::ALL_TYPES - DocumentTypes::IDENTITY_TYPES - DocumentTypes::SECONDARY_IDENTITY_TYPES)
+    available_doc_types.map { |d| {key: d.key, label: d.label} }
+  end
+
   def self.screen_document!(document:)
     raise "Document has no upload attached" unless document.upload.attached?
 
     media_type = document.upload.content_type
     raise "Unsupported media type: #{media_type}" unless SUPPORTED_MEDIA_TYPES.include?(media_type)
 
-    base64_data = Base64.strict_encode64(document.upload.download)
+    input = if media_type == "application/pdf"
+              pdf_to_png_base64(document.upload)
+            else
+              [{
+                 media_type: media_type,
+                 base64_data: Base64.strict_encode64(document.upload.download)
+               }]
+            end
 
     body_hash = construct_bedrock_payload(
-      base64_data: base64_data,
-      media_type: media_type,
+      images: input,
       user_prompt: prompt_for(document_type: document.document_type)
     )
 
@@ -61,15 +88,17 @@ def self.screen_document!(document:)
     [result_json, raw_response_json]
   end
 
-  def self.construct_bedrock_payload(base64_data:, media_type:, user_prompt:)
+  def self.construct_bedrock_payload(images:, user_prompt:)
     {
-      anthropic_version: 'bedrock-2023-05-31',
+      anthropic_version: "bedrock-2023-05-31",
       max_tokens: 250,
       messages: [{
-                   role: 'user',
+                   role: "user",
                    content: [
-                     { type: 'image', source: { type: 'base64', media_type: media_type, data: base64_data } },
-                     { type: 'text', text: user_prompt }
+                     *images.map do |img|
+                       { type: "image", source: { type: "base64", media_type: img[:media_type], data: img[:base64_data] } }
+                     end,
+                     { type: "text", text: user_prompt }
                    ]
                  }]
     }
@@ -89,17 +118,58 @@ def self.extract_text_from_response(response)
     Array(response['content'])
       .select { |content| content['type'] == 'text' }
       .map { |content| content['text'] }
-      .join('\n')
+      .join("\n")
       .strip
   end
 
   def self.parse_strict_json!(text)
-    cleaned = text.to_s.strip
+    s = text.to_s
 
-    cleaned = cleaned.sub(/\A```(?:json)?\s*/i, "").sub(/\s*```\z/, "").strip
+    blocks = s.scan(/```(?:json)?\s*(\{.*?\})\s*```/m)
+    if blocks.any?
+      json_str = blocks.last.first
+      return JSON.parse(json_str)
+    end
 
+    cleaned = s.strip
+    cleaned = cleaned.sub(/\A```(?:json)?\s*/i, "").sub(/\s*```\z/, "").strip
     JSON.parse(cleaned)
   rescue JSON::ParserError => e
-    raise "Model did not return valid JSON. Error=#{e.message}. Output=#{text.inspect}"
+    raise "Bedrock did not return valid JSON. \n Error: #{e.message} \n Output: #{text.inspect}"
+  end
+
+
+  def self.pdf_to_png_base64(upload)
+    images = []
+
+    Tempfile.create(["upload", ".pdf"]) do |pdf|
+      pdf.binmode
+      pdf.write(upload.download)
+      pdf.flush
+
+      MiniMagick::Image.open(pdf.path).pages.each_with_index do |page, index|
+        Tempfile.create(["pdf_page_#{index}", ".png"]) do |png|
+          MiniMagick::Tool::Convert.new do |convert|
+            convert.density(200)
+            convert.quality(90)
+            convert << "#{pdf.path}[#{index}]"
+            convert << png.path
+          end
+
+          data = File.binread(png.path)
+
+          images << {
+            media_type: "image/png",
+            base64_data: Base64.strict_encode64(data)
+          }
+        end
+      end
+    end
+
+    raise "pdf produced no pages" if images.empty?
+    images
+  rescue MiniMagick::Error, MiniMagick::Invalid => e
+    raise "failed to convert pdf pages to images (perhaps minimagick or ghostscript issue). #{e.class}: #{e.message}"
   end
+
 end

app/views/hub/documents/edit.html.erb

@@ -9,9 +9,9 @@
         <%= render 'form', url: [:hub, @client, @document], method: "patch", file_upload_enabled: false %>
       </div>
       <div id="image-and-button-wrapper">
-        <div class="item" id="image-container" >
+        <div class="item" id="image-container">
           <% if @document.is_pdf? %>
-            <embed src="<%= transient_storage_url(@document.upload.blob) %>" width="800px" height="100%" type="application/pdf" />
+            <embed src="<%= transient_storage_url(@document.upload.blob) %>" width="800px" height="100%" type="application/pdf"/>
           <% else %>
             <%= image_tag transient_storage_url(@document.upload.blob), id: 'image', class: 'rotatable-image', data: { rotation: 0 } %>
           <% end %>
@@ -23,4 +23,39 @@
     </div>
 
   </div>
-<% end %>
+
+  <% if !acts_like_production? && current_user.admin? %>
+    <div id="doc-assessments-table">
+      <%= button_to "Re-run screener",
+                    rerun_screener_hub_client_document_path(client_id: @document.client.id, id: @document.id),
+                    method: :post %>
+      <br>
+      <% if @document.assessments.present? %>
+        <table class="index-table">
+          <tr>
+            <th>job complete?</th>
+            <th>error?</th>
+            <th>updated at</th>
+            <th>verdict</th>
+            <th>confidence</th>
+            <th>reason</th>
+            <th>explanation</th>
+          </tr>
+          <% @document.assessments.order('created_at DESC').each do |assessment| %>
+            <% result = assessment&.result_json %>
+            <% error = assessment&.error %>
+          <tr class="<%= error.present? ? "error" : result["verdict"] %>">
+            <td><%= assessment.status %></td>
+            <td><%= "Error: #{JSON.pretty_generate(error)}" unless error.blank? %></td>
+            <td><%= long_formatted_datetime(assessment.updated_at) %></td>
+            <td><%= result["verdict"] %></td>
+            <td><%= result["confidence"] %></td>
+            <td><u><%= result["reason"] %></u></td>
+            <td><%= result["explanation"] %></td>
+          </tr>
+      <% end %>
+      </table>
+  <% end %>
+  </div>
+<% end %>
+<% end %>

config/routes.rb

@@ -286,6 +286,7 @@ def scoped_navigation_routes(context, navigation)
           resources :documents do
             get "/archived", to: "documents#archived", on: :collection, as: :archived
             get "/confirm", to: "documents#confirm", on: :member, as: :confirm
+            post :rerun_screener, on: :member
           end
           resources :notes, only: [:create, :index]
           resources :messages, only: [:index]

db/migrate/20251217230133_create_doc_assessments.rb

@@ -16,8 +16,5 @@ def change
 
       t.timestamps
     end
-
-    # prevents dupes for the same document/file/prompt-version
-    add_index :doc_assessments, %i[document_id prompt_version input_blob_id], unique: true, name: "index_doc_assessments_on_doc_and_version_and_blob"
   end
 end