diff --git a/Gemfile b/Gemfile index b4752c745..78fe3016f 100644 --- a/Gemfile +++ b/Gemfile @@ -31,6 +31,14 @@ gem 'haikunator' gem 'liquid' # Parse Markdown to HTML gem 'commonmarker' +# Convert PDF pages to images for content extraction +gem 'pdftoimage' +# Google Vision API for text extraction from images +gem 'google-cloud-vision' +# Tesseract OCR fallback when Vision API not available +gem 'rtesseract' +# Image manipulation for PDF page batching +gem 'mini_magick' # Validate Data against JSON Schema gem 'json_schemer' # used in swagger build diff --git a/enterprise/app/jobs/captain/documents/pdf_processing_job.rb b/enterprise/app/jobs/captain/documents/pdf_processing_job.rb new file mode 100644 index 000000000..2bc300ce5 --- /dev/null +++ b/enterprise/app/jobs/captain/documents/pdf_processing_job.rb @@ -0,0 +1,30 @@ +class Captain::Documents::PdfProcessingJob < ApplicationJob + queue_as :low + + def perform(document) + return unless document.pdf_document? + + Rails.logger.info "Starting PDF processing for document #{document.id}" + + begin + service = Captain::Tools::PdfProcessingService.new(document) + extracted_content = service.process + + Rails.logger.info "Successfully processed PDF document #{document.id}, extracted #{extracted_content.length} characters" + + rescue Captain::Tools::PdfProcessingService::PdfProcessingError => e + Rails.logger.error "PDF processing failed for document #{document.id}: #{e.message}" + + # Mark document as failed (you may want to add a 'failed' status to the enum) + document.update!(status: :in_progress) # Keep as in_progress for potential retry + + # Optionally, you could add a retry mechanism or notification + raise e + + rescue StandardError => e + Rails.logger.error "Unexpected error processing PDF document #{document.id}: #{e.message}" + document.update!(status: :in_progress) + raise e + end + end +end \ No newline at end of file diff --git a/enterprise/app/models/captain/document.rb b/enterprise/app/models/captain/document.rb index d2a02f5b5..6ed8d6e8a 100644 --- a/enterprise/app/models/captain/document.rb +++ b/enterprise/app/models/captain/document.rb @@ -37,7 +37,13 @@ class Captain::Document < ApplicationRecord available: 1 } + enum document_type: { + url: 0, + pdf: 1 + } + before_create :ensure_within_plan_limit + before_create :detect_document_type after_create_commit :enqueue_crawl_job after_create_commit :update_document_usage after_destroy :update_document_usage @@ -52,7 +58,33 @@ class Captain::Document < ApplicationRecord def enqueue_crawl_job return if status != 'in_progress' - Captain::Documents::CrawlJob.perform_later(self) + if pdf_document? + Captain::Documents::PdfProcessingJob.perform_later(self) + else + Captain::Documents::CrawlJob.perform_later(self) + end + end + + def url_document? + document_type == 'url' + end + + def pdf_document? + document_type == 'pdf' + end + + def detect_document_type + return unless external_link.present? + + # Check if URL points to a PDF file + uri = URI.parse(external_link) + if uri.path.downcase.end_with?('.pdf') || external_link.downcase.include?('.pdf') + self.document_type = :pdf + else + self.document_type = :url + end + rescue URI::InvalidURIError + self.document_type = :url end def enqueue_response_builder_job diff --git a/enterprise/app/services/captain/tools/pdf_page_batcher_service.rb b/enterprise/app/services/captain/tools/pdf_page_batcher_service.rb new file mode 100644 index 000000000..7e0374b20 --- /dev/null +++ b/enterprise/app/services/captain/tools/pdf_page_batcher_service.rb @@ -0,0 +1,57 @@ +require 'mini_magick' + +class Captain::Tools::PdfPageBatcherService + PAGES_PER_BATCH = 3 + + def initialize(temp_files_tracker) + @temp_files = temp_files_tracker + end + + def batch_pages(image_files) + return image_files if image_files.length <= PAGES_PER_BATCH + + batched_images = [] + + image_files.each_slice(PAGES_PER_BATCH) do |batch| + if batch.length == 1 + # Single page, use as-is + batched_images << batch.first + else + # Multiple pages, combine them + combined_image = combine_images(batch) + batched_images << combined_image + end + end + + batched_images + end + + private + + def combine_images(image_paths) + # Create a vertical montage of the images + output_path = create_temp_file_path('.png') + + MiniMagick::Tool::Montage.new do |montage| + image_paths.each { |path| montage << path } + montage.tile('1x') # Arrange in single column (vertically) + montage.geometry('+0+10') # Add 10px spacing between images + montage.background('white') + montage << output_path + end + + @temp_files << output_path + output_path + rescue StandardError => e + Rails.logger.error "Failed to combine images: #{e.message}" + # Fallback to first image if combining fails + image_paths.first + end + + def create_temp_file_path(extension) + temp_file = Tempfile.new(['combined_pages', extension]) + temp_file.close + @temp_files << temp_file + temp_file.path + end +end \ No newline at end of file diff --git a/enterprise/app/services/captain/tools/pdf_processing_service.rb b/enterprise/app/services/captain/tools/pdf_processing_service.rb new file mode 100644 index 000000000..fbe36bfb2 --- /dev/null +++ b/enterprise/app/services/captain/tools/pdf_processing_service.rb @@ -0,0 +1,75 @@ +class Captain::Tools::PdfProcessingService + def initialize(document) + @document = document + @temp_files = [] + end + + def process + return unless @document.pdf_document? + + content = extract_text_content + + # Follow the same pattern as SimplePageCrawlParserJob + @document.update!( + content: content[0..14_999], # Same limit as crawler + status: :available + ) + + content + rescue StandardError => e + Rails.logger.error "PDF processing failed for document #{@document.id}: #{e.message}" + @document.update!(content: "Error processing PDF: #{e.message}", status: :available) + raise e + ensure + cleanup_temp_files + end + + private + + def extract_text_content + require 'down' + + pdf_file = Down.download(@document.external_link) + @temp_files << pdf_file + + output_dir = Dir.mktmpdir + @temp_files << output_dir + + # Convert PDF to images + system("pdftoimage", "-png", "-r", "150", pdf_file.path, "#{output_dir}/page") + + image_files = Dir.glob("#{output_dir}/page-*.png").sort + return fallback_content if image_files.empty? + + # Batch pages for efficient processing + batcher = Captain::Tools::PdfPageBatcherService.new(@temp_files) + batched_images = batcher.batch_pages(image_files) + + # Extract text using Google Vision + vision_extractor = Captain::Tools::VisionTextExtractorService.new + extracted_text = vision_extractor.extract_text_from_multiple_images(batched_images) + + extracted_text.present? ? extracted_text : fallback_content + end + + def fallback_content + "PDF Document: #{@document.name || 'Untitled'}\nSource: #{@document.external_link}" + end + + def cleanup_temp_files + @temp_files.each do |file_or_dir| + case file_or_dir + when String # directory path + FileUtils.rm_rf(file_or_dir) + when Tempfile + file_or_dir.close + file_or_dir.unlink + else # Down::ChunkedIO or similar + file_or_dir.close if file_or_dir.respond_to?(:close) + file_or_dir.unlink if file_or_dir.respond_to?(:unlink) + end + rescue StandardError => e + Rails.logger.warn "Failed to cleanup temp file: #{e.message}" + end + end +end \ No newline at end of file diff --git a/enterprise/app/services/captain/tools/vision_text_extractor_service.rb b/enterprise/app/services/captain/tools/vision_text_extractor_service.rb new file mode 100644 index 000000000..547503986 --- /dev/null +++ b/enterprise/app/services/captain/tools/vision_text_extractor_service.rb @@ -0,0 +1,68 @@ +class Captain::Tools::VisionTextExtractorService + class TextExtractionError < StandardError; end + + def initialize + @use_vision = vision_available? + @vision = Google::Cloud::Vision.image_annotator if @use_vision + end + + def extract_text_from_image(image_path) + if @use_vision + extract_with_vision(image_path) + else + extract_with_tesseract(image_path) + end + end + + def extract_text_from_multiple_images(image_paths) + Rails.logger.info "Using #{@use_vision ? 'Google Vision' : 'Tesseract'} for text extraction" + + extracted_texts = [] + + image_paths.each_with_index do |image_path, index| + Rails.logger.info "Extracting text from image #{index + 1}/#{image_paths.length}" + + text = extract_text_from_image(image_path) + extracted_texts << text if text.present? + end + + extracted_texts.join("\n\n") + end + + private + + def vision_available? + require 'google/cloud/vision' + Google::Cloud::Vision.image_annotator + true + rescue StandardError => e + Rails.logger.info "Google Vision not available, falling back to Tesseract: #{e.message}" + false + end + + def extract_with_vision(image_path) + image = Google::Cloud::Vision::Image.new(image_path) + response = @vision.text_detection(image: image) + + if response.error + raise TextExtractionError, "Vision API error: #{response.error.message}" + end + + text_annotation = response.full_text_annotation + return "" unless text_annotation + + text_annotation.text + rescue Google::Cloud::Error => e + Rails.logger.error "Google Vision API error: #{e.message}" + raise TextExtractionError, "Failed to extract text: #{e.message}" + end + + def extract_with_tesseract(image_path) + require 'rtesseract' + + RTesseract.new(image_path).to_s.strip + rescue StandardError => e + Rails.logger.error "Tesseract extraction error: #{e.message}" + raise TextExtractionError, "Failed to extract text with Tesseract: #{e.message}" + end +end \ No newline at end of file diff --git a/enterprise/app/views/api/v1/models/captain/_document.json.jbuilder b/enterprise/app/views/api/v1/models/captain/_document.json.jbuilder index 83724b9cd..cce21fe6d 100644 --- a/enterprise/app/views/api/v1/models/captain/_document.json.jbuilder +++ b/enterprise/app/views/api/v1/models/captain/_document.json.jbuilder @@ -4,6 +4,7 @@ json.assistant do end json.content resource.content json.created_at resource.created_at.to_i +json.document_type resource.document_type json.external_link resource.external_link json.id resource.id json.name resource.name diff --git a/enterprise/db/migrate/20250623_add_document_type_and_file_to_captain_documents.rb b/enterprise/db/migrate/20250623_add_document_type_and_file_to_captain_documents.rb new file mode 100644 index 000000000..989596674 --- /dev/null +++ b/enterprise/db/migrate/20250623_add_document_type_and_file_to_captain_documents.rb @@ -0,0 +1,6 @@ +class AddDocumentTypeToCaptainDocuments < ActiveRecord::Migration[7.0] + def change + add_column :captain_documents, :document_type, :integer, default: 0, null: false + add_index :captain_documents, :document_type + end +end \ No newline at end of file diff --git a/spec/enterprise/models/captain/document_spec.rb b/spec/enterprise/models/captain/document_spec.rb new file mode 100644 index 000000000..fc5e52130 --- /dev/null +++ b/spec/enterprise/models/captain/document_spec.rb @@ -0,0 +1,69 @@ +require 'rails_helper' + +RSpec.describe Captain::Document, type: :model do + let(:account) { create(:account) } + let(:assistant) { create(:captain_assistant, account: account) } + + describe 'document_type detection' do + it 'detects PDF documents from URL' do + document = assistant.documents.build( + name: 'Test PDF', + external_link: 'https://example.com/document.pdf' + ) + + document.save! + + expect(document.document_type).to eq('pdf') + expect(document.pdf_document?).to be true + expect(document.url_document?).to be false + end + + it 'detects URL documents for non-PDF links' do + document = assistant.documents.build( + name: 'Test Page', + external_link: 'https://example.com/page.html' + ) + + document.save! + + expect(document.document_type).to eq('url') + expect(document.url_document?).to be true + expect(document.pdf_document?).to be false + end + + it 'defaults to URL for invalid URIs' do + document = assistant.documents.build( + name: 'Invalid URL', + external_link: 'not-a-valid-url' + ) + + document.save! + + expect(document.document_type).to eq('url') + end + end + + describe 'job enqueueing' do + it 'enqueues PDF processing job for PDF documents' do + document = assistant.documents.build( + name: 'Test PDF', + external_link: 'https://example.com/document.pdf' + ) + + expect(Captain::Documents::PdfProcessingJob).to receive(:perform_later).with(document) + + document.save! + end + + it 'enqueues crawl job for URL documents' do + document = assistant.documents.build( + name: 'Test Page', + external_link: 'https://example.com/page.html' + ) + + expect(Captain::Documents::CrawlJob).to receive(:perform_later).with(document) + + document.save! + end + end +end \ No newline at end of file diff --git a/spec/enterprise/services/captain/tools/pdf_processing_service_spec.rb b/spec/enterprise/services/captain/tools/pdf_processing_service_spec.rb new file mode 100644 index 000000000..6e6c4ebbd --- /dev/null +++ b/spec/enterprise/services/captain/tools/pdf_processing_service_spec.rb @@ -0,0 +1,76 @@ +require 'rails_helper' + +RSpec.describe Captain::Tools::PdfProcessingService do + let(:account) { create(:account) } + let(:assistant) { create(:captain_assistant, account: account) } + let(:document) do + assistant.documents.create!( + name: 'Test PDF', + external_link: 'https://example.com/test.pdf', + document_type: :pdf + ) + end + + subject { described_class.new(document) } + + describe '#process' do + context 'when document is not a PDF' do + before { allow(document).to receive(:pdf_document?).and_return(false) } + + it 'returns early without processing' do + expect(subject.process).to be_nil + end + end + + context 'when document is a PDF' do + let(:mock_file) { double('file', path: '/tmp/test.pdf', close: nil, unlink: nil) } + let(:mock_batcher) { double('batcher') } + let(:mock_vision) { double('vision') } + + before do + allow(Down).to receive(:download).and_return(mock_file) + allow(Dir).to receive(:mktmpdir).and_return('/tmp/pdf_images') + allow(Dir).to receive(:glob).and_return(['/tmp/image1.png', '/tmp/image2.png']) + allow(FileUtils).to receive(:rm_rf) + + allow(Captain::Tools::PdfPageBatcherService).to receive(:new).and_return(mock_batcher) + allow(mock_batcher).to receive(:batch_pages).and_return(['/tmp/batched1.png']) + + allow(Captain::Tools::VisionTextExtractorService).to receive(:new).and_return(mock_vision) + allow(mock_vision).to receive(:extract_text_from_multiple_images).and_return('Extracted text content') + end + + it 'processes the PDF and updates document' do + expect { subject.process }.to change { document.reload.status }.to('available') + end + + it 'sets extracted content on the document' do + subject.process + expect(document.reload.content).to eq('Extracted text content') + end + + it 'uses batching service for page processing' do + expect(mock_batcher).to receive(:batch_pages).with(['/tmp/image1.png', '/tmp/image2.png']) + subject.process + end + + it 'uses Vision API for text extraction' do + expect(mock_vision).to receive(:extract_text_from_multiple_images).with(['/tmp/batched1.png']) + subject.process + end + end + + context 'when processing fails' do + before do + allow(Down).to receive(:download).and_raise(StandardError.new('Download failed')) + end + + it 'logs error and updates document with error message' do + expect(Rails.logger).to receive(:error).with(/PDF processing failed/) + + expect { subject.process }.to raise_error(StandardError) + expect(document.reload.content).to include('Error processing PDF') + end + end + end +end \ No newline at end of file diff --git a/spec/enterprise/services/captain/tools/vision_text_extractor_service_spec.rb b/spec/enterprise/services/captain/tools/vision_text_extractor_service_spec.rb new file mode 100644 index 000000000..f224968da --- /dev/null +++ b/spec/enterprise/services/captain/tools/vision_text_extractor_service_spec.rb @@ -0,0 +1,98 @@ +require 'rails_helper' + +RSpec.describe Captain::Tools::VisionTextExtractorService do + let(:image_path) { '/tmp/test_image.png' } + + describe 'when Google Vision is available' do + subject { described_class.new } + + let(:mock_vision) { double('vision_client') } + let(:mock_response) { double('response', error: nil) } + let(:mock_text_annotation) { double('text_annotation', text: 'Vision extracted text') } + + before do + allow(Google::Cloud::Vision).to receive(:image_annotator).and_return(mock_vision) + allow(Google::Cloud::Vision::Image).to receive(:new).with(image_path).and_return(double('image')) + allow(mock_vision).to receive(:text_detection).and_return(mock_response) + allow(mock_response).to receive(:full_text_annotation).and_return(mock_text_annotation) + end + + describe '#extract_text_from_image' do + it 'uses Vision API for text extraction' do + result = subject.extract_text_from_image(image_path) + expect(result).to eq('Vision extracted text') + end + + context 'when Vision API returns error' do + before do + allow(mock_response).to receive(:error).and_return(double('error', message: 'API Error')) + end + + it 'raises TextExtractionError' do + expect { subject.extract_text_from_image(image_path) }.to raise_error( + Captain::Tools::VisionTextExtractorService::TextExtractionError, + /Vision API error/ + ) + end + end + end + end + + describe 'when Google Vision is not available' do + subject { described_class.new } + + let(:mock_tesseract) { double('tesseract', to_s: 'Tesseract extracted text') } + + before do + allow(Google::Cloud::Vision).to receive(:image_annotator).and_raise(StandardError.new('No credentials')) + allow(RTesseract).to receive(:new).with(image_path).and_return(mock_tesseract) + end + + describe '#extract_text_from_image' do + it 'falls back to Tesseract for text extraction' do + result = subject.extract_text_from_image(image_path) + expect(result).to eq('Tesseract extracted text') + end + + context 'when Tesseract fails' do + before do + allow(RTesseract).to receive(:new).and_raise(StandardError.new('Tesseract error')) + end + + it 'raises TextExtractionError' do + expect { subject.extract_text_from_image(image_path) }.to raise_error( + Captain::Tools::VisionTextExtractorService::TextExtractionError, + /Failed to extract text with Tesseract/ + ) + end + end + end + end + + describe '#extract_text_from_multiple_images' do + let(:image_paths) { ['/tmp/image1.png', '/tmp/image2.png'] } + + before do + allow(Google::Cloud::Vision).to receive(:image_annotator).and_return(double('vision')) + allow(subject).to receive(:extract_text_from_image).with('/tmp/image1.png').and_return('Text from page 1') + allow(subject).to receive(:extract_text_from_image).with('/tmp/image2.png').and_return('Text from page 2') + end + + it 'extracts text from multiple images and joins them' do + result = subject.extract_text_from_multiple_images(image_paths) + expect(result).to eq("Text from page 1\n\nText from page 2") + end + + it 'skips empty results' do + allow(subject).to receive(:extract_text_from_image).with('/tmp/image2.png').and_return('') + + result = subject.extract_text_from_multiple_images(image_paths) + expect(result).to eq('Text from page 1') + end + + it 'logs which extraction method is being used' do + expect(Rails.logger).to receive(:info).with(/Using (Google Vision|Tesseract) for text extraction/) + subject.extract_text_from_multiple_images(image_paths) + end + end +end \ No newline at end of file