From e7538132277ab34b8d8d0240c71b89aa85e95ffb Mon Sep 17 00:00:00 2001 From: Anthony Tseng Date: Tue, 22 Apr 2025 18:10:01 -0700 Subject: [PATCH] [AI Chat] Full page screenshots feature (#28714) --- browser/ai_chat/BUILD.gn | 3 + browser/ai_chat/DEPS | 4 + .../ai_chat/full_screenshotter_unittest.cc | 505 ++++++++++++++++++ browser/ai_chat/upload_file_helper.cc | 68 +-- .../ai_chat/upload_file_helper_unittest.cc | 27 +- components/ai_chat/content/browser/BUILD.gn | 5 + components/ai_chat/content/browser/DEPS | 2 + .../content/browser/ai_chat_tab_helper.cc | 30 ++ .../content/browser/ai_chat_tab_helper.h | 11 + .../content/browser/full_screenshotter.cc | 254 +++++++++ .../content/browser/full_screenshotter.h | 105 ++++ components/ai_chat/core/browser/DEPS | 1 + .../ai_chat/core/browser/ai_chat_database.cc | 65 ++- .../core/browser/ai_chat_database_unittest.cc | 59 +- .../core/browser/ai_chat_service_unittest.cc | 5 +- .../core/browser/conversation_handler.cc | 46 +- .../core/browser/conversation_handler.h | 7 +- .../browser/conversation_handler_unittest.cc | 235 +++++--- .../browser/engine/conversation_api_client.cc | 3 +- .../browser/engine/conversation_api_client.h | 1 + .../engine_consumer_conversation_api.cc | 26 +- ...gine_consumer_conversation_api_unittest.cc | 43 +- .../browser/engine/engine_consumer_oai.cc | 51 +- .../engine/engine_consumer_oai_unittest.cc | 62 ++- components/ai_chat/core/browser/test_utils.cc | 34 +- components/ai_chat/core/browser/test_utils.h | 2 +- components/ai_chat/core/browser/utils.cc | 45 ++ components/ai_chat/core/browser/utils.h | 4 + .../ai_chat/core/browser/utils_unittest.cc | 25 +- .../ai_chat/core/common/mojom/ai_chat.mojom | 31 +- components/ai_chat/core/common/test_utils.cc | 26 +- components/ai_chat/core/common/test_utils.h | 5 +- .../common/conversation_history_utils.ts | 51 ++ .../attachment_button_menu/index.tsx | 18 +- .../page/components/input_box/index.tsx | 4 + .../components/input_box/style.module.scss | 3 + .../components/uploaded_img_item/index.tsx | 4 +- .../page/state/conversation_context.tsx | 61 ++- .../page/stories/components_panel.tsx | 121 ++++- .../components/conversation_entries/index.tsx | 21 +- .../conversation_entries/style.module.scss | 6 + .../untrusted_conversation_frame_api.ts | 17 +- .../page/components/BeginGeneration.tsx | 1 + .../AIChat/Components/AIChatView.swift | 2 +- .../Messages/AIChatResponseMessageView.swift | 2 +- ios/browser/api/ai_chat/conversation_client.h | 2 +- .../api/ai_chat/conversation_client.mm | 3 +- .../aichat_database_dump_version_1.sql | 1 + .../aichat_database_dump_version_2.sql | 1 + .../aichat_database_dump_version_3.sql | 24 + 50 files changed, 1798 insertions(+), 334 deletions(-) create mode 100644 browser/ai_chat/full_screenshotter_unittest.cc create mode 100644 components/ai_chat/content/browser/full_screenshotter.cc create mode 100644 components/ai_chat/content/browser/full_screenshotter.h create mode 100644 components/ai_chat/resources/common/conversation_history_utils.ts create mode 100644 test/data/ai_chat/aichat_database_dump_version_3.sql diff --git a/browser/ai_chat/BUILD.gn b/browser/ai_chat/BUILD.gn index c2222056d49..f9f59d80e64 100644 --- a/browser/ai_chat/BUILD.gn +++ b/browser/ai_chat/BUILD.gn @@ -59,6 +59,7 @@ source_set("unit_tests") { sources = [ "ai_chat_throttle_unittest.cc", "brave_open_ai_chat_permission_context_unittest.cc", + "full_screenshotter_unittest.cc", "upload_file_helper_unittest.cc", ] @@ -72,6 +73,8 @@ source_set("unit_tests") { "//brave/components/constants", "//chrome/common", "//chrome/test:test_support", + "//components/paint_preview/common/mojom", + "//components/services/paint_preview_compositor/public/mojom", "//content/public/browser", "//content/test:test_support", "//testing/gtest", diff --git a/browser/ai_chat/DEPS b/browser/ai_chat/DEPS index 079d4323f9e..e86d5b970ba 100644 --- a/browser/ai_chat/DEPS +++ b/browser/ai_chat/DEPS @@ -5,5 +5,9 @@ include_rules = [ specific_include_rules = { "brave_open_ai_chat_permission_context_unittest\.cc": [ "+brave/components/permissions/contexts/brave_open_ai_chat_permission_context.h", + ], + "full_screenshotter_unittest\.cc" : [ + "+components/paint_preview/common/mojom", + "+content/test", ] } diff --git a/browser/ai_chat/full_screenshotter_unittest.cc b/browser/ai_chat/full_screenshotter_unittest.cc new file mode 100644 index 00000000000..cc1379374d8 --- /dev/null +++ b/browser/ai_chat/full_screenshotter_unittest.cc @@ -0,0 +1,505 @@ +// Copyright (c) 2025 The Brave Authors. All rights reserved. +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this file, +// You can obtain one at https://mozilla.org/MPL/2.0/. + +#include "brave/components/ai_chat/content/browser/full_screenshotter.h" + +#include + +#include "base/test/test_future.h" +#include "chrome/test/base/chrome_render_view_host_test_harness.h" +#include "components/paint_preview/common/mojom/paint_preview_recorder.mojom.h" +#include "components/services/paint_preview_compositor/public/mojom/paint_preview_compositor.mojom.h" +#include "content/public/test/test_renderer_host.h" +#include "content/public/test/web_contents_tester.h" +#include "content/test/test_render_view_host.h" +#include "content/test/test_web_contents.h" +#include "mojo/public/cpp/bindings/associated_receiver.h" +#include "testing/gtest/include/gtest/gtest.h" +#include "third_party/blink/public/common/associated_interfaces/associated_interface_provider.h" +#include "ui/gfx/geometry/rect.h" +#include "ui/gfx/image/image_unittest_util.h" + +namespace ai_chat { + +namespace { + +class TestView : public content::TestRenderWidgetHostView { + public: + explicit TestView(content::RenderWidgetHost* widget) + : TestRenderWidgetHostView(widget) {} + + gfx::Rect GetViewBounds() override { return view_bounds_; } + void SetViewBounds(const gfx::Rect& bounds) { view_bounds_ = bounds; } + + private: + gfx::Rect view_bounds_; +}; + +class MockPaintPreviewRecorder + : public paint_preview::mojom::PaintPreviewRecorder { + public: + MockPaintPreviewRecorder() = default; + ~MockPaintPreviewRecorder() override = default; + + MockPaintPreviewRecorder(const MockPaintPreviewRecorder&) = delete; + MockPaintPreviewRecorder& operator=(const MockPaintPreviewRecorder&) = delete; + + void CapturePaintPreview( + paint_preview::mojom::PaintPreviewCaptureParamsPtr params, + paint_preview::mojom::PaintPreviewRecorder::CapturePaintPreviewCallback + callback) override { + std::move(callback).Run(status_, std::move(response_)); + } + + void SetResponse( + paint_preview::mojom::PaintPreviewStatus status, + paint_preview::mojom::PaintPreviewCaptureResponsePtr&& response) { + status_ = status; + response_ = std::move(response); + } + + void BindRequest(mojo::ScopedInterfaceEndpointHandle handle) { + binding_.reset(); + binding_.Bind( + mojo::PendingAssociatedReceiver< + paint_preview::mojom::PaintPreviewRecorder>(std::move(handle))); + } + + private: + paint_preview::mojom::PaintPreviewStatus status_; + paint_preview::mojom::PaintPreviewCaptureResponsePtr response_; + mojo::AssociatedReceiver binding_{ + this}; +}; + +class MockPaintPreviewCompositorClient + : public paint_preview::PaintPreviewCompositorClient { + public: + explicit MockPaintPreviewCompositorClient( + scoped_refptr task_runner) + : response_status_(paint_preview::mojom::PaintPreviewCompositor:: + BeginCompositeStatus::kSuccess), + bitmap_status_(paint_preview::mojom::PaintPreviewCompositor:: + BitmapStatus::kSuccess), + is_empty_bitmap_(false), + token_(base::UnguessableToken::Create()), + task_runner_(task_runner) {} + ~MockPaintPreviewCompositorClient() override = default; + + MockPaintPreviewCompositorClient(const MockPaintPreviewCompositorClient&) = + delete; + MockPaintPreviewCompositorClient& operator=( + const MockPaintPreviewCompositorClient&) = delete; + + const std::optional& Token() const override { + return token_; + } + + void SetDisconnectHandler(base::OnceClosure closure) override { + disconnect_handler_ = std::move(closure); + } + + void BeginSeparatedFrameComposite( + paint_preview::mojom::PaintPreviewBeginCompositeRequestPtr request, + paint_preview::mojom::PaintPreviewCompositor:: + BeginSeparatedFrameCompositeCallback callback) override { + NOTREACHED(); + } + + void BitmapForSeparatedFrame( + const base::UnguessableToken& frame_guid, + const gfx::Rect& clip_rect, + float scale_factor, + paint_preview::mojom::PaintPreviewCompositor:: + BitmapForSeparatedFrameCallback callback, + bool run_task_on_default_task_runner = true) override { + NOTREACHED(); + } + + void BeginMainFrameComposite( + paint_preview::mojom::PaintPreviewBeginCompositeRequestPtr request, + paint_preview::mojom::PaintPreviewCompositor:: + BeginMainFrameCompositeCallback callback) override { + auto response = + paint_preview::mojom::PaintPreviewBeginCompositeResponse::New(); + response->root_frame_guid = root_frame_guid_; + if (!frames_.empty()) { + response->frames = std::move(frames_); + } + task_runner_->PostTask(FROM_HERE, + base::BindOnce(std::move(callback), response_status_, + std::move(response))); + } + + void BitmapForMainFrame( + const gfx::Rect& clip_rect, + float scale_factor, + paint_preview::mojom::PaintPreviewCompositor::BitmapForMainFrameCallback + callback, + bool run_task_on_default_task_runner = true) override { + task_runner_->PostDelayedTask( + FROM_HERE, + base::BindOnce(std::move(callback), bitmap_status_, + is_empty_bitmap_ + ? SkBitmap() + : gfx::test::CreateBitmap(clip_rect.width(), + clip_rect.height())), + base::Seconds(1)); + } + + void SetRootFrameUrl(const GURL& url) override { + // no-op. + } + + void SetBeginMainFrameResponseStatus( + paint_preview::mojom::PaintPreviewCompositor::BeginCompositeStatus + status) { + response_status_ = status; + } + + void SetBitmapStatus( + paint_preview::mojom::PaintPreviewCompositor::BitmapStatus status) { + bitmap_status_ = status; + } + + void SetIsEmptyBitmap(bool is_empty) { is_empty_bitmap_ = is_empty; } + + void Disconnect() { + if (disconnect_handler_) { + std::move(disconnect_handler_).Run(); + } + } + + void SetCompositeResponse( + base::flat_map> frames, + const base::UnguessableToken& root_guid) { + frames_ = std::move(frames); + root_frame_guid_ = root_guid; + } + + private: + paint_preview::mojom::PaintPreviewCompositor::BeginCompositeStatus + response_status_; + paint_preview::mojom::PaintPreviewCompositor::BitmapStatus bitmap_status_; + bool is_empty_bitmap_; + std::optional token_; + base::OnceClosure disconnect_handler_; + scoped_refptr task_runner_; + base::flat_map> + frames_; + base::UnguessableToken root_frame_guid_; +}; + +class MockPaintPreviewCompositorService + : public paint_preview::PaintPreviewCompositorService { + public: + explicit MockPaintPreviewCompositorService( + scoped_refptr task_runner) + : task_runner_(task_runner) {} + ~MockPaintPreviewCompositorService() override = default; + + MockPaintPreviewCompositorService(const MockPaintPreviewCompositorService&) = + delete; + MockPaintPreviewCompositorService& operator=( + const MockPaintPreviewCompositorService&) = delete; + + std::unique_ptr + CreateCompositor(base::OnceClosure connected_closure) override { + task_runner_->PostTask(FROM_HERE, std::move(connected_closure)); + return std::unique_ptr( + new MockPaintPreviewCompositorClient(task_runner_), + base::OnTaskRunnerDeleter(task_runner_)); + } + + void OnMemoryPressure(base::MemoryPressureListener::MemoryPressureLevel + memory_pressure_level) override { + // no-op. + } + + bool HasActiveClients() const override { NOTREACHED(); } + + void SetDisconnectHandler(base::OnceClosure disconnect_handler) override { + disconnect_handler_ = std::move(disconnect_handler); + } + + void Disconnect() { + if (disconnect_handler_) { + std::move(disconnect_handler_).Run(); + } + } + + private: + base::OnceClosure disconnect_handler_; + scoped_refptr task_runner_; +}; + +MockPaintPreviewCompositorClient* AsMockClient( + paint_preview::PaintPreviewCompositorClient* client) { + return static_cast(client); +} + +} // namespace + +class FullScreenshotterTest : public ChromeRenderViewHostTestHarness { + public: + FullScreenshotterTest() = default; + ~FullScreenshotterTest() override = default; + + FullScreenshotterTest(const FullScreenshotterTest&) = delete; + FullScreenshotterTest& operator=(const FullScreenshotterTest&) = delete; + + protected: + void SetUp() override { + ChromeRenderViewHostTestHarness::SetUp(); + NavigateAndCommit(GURL("https://brave.com/"), + ui::PageTransition::PAGE_TRANSITION_FIRST); + // Store the original RenderWidgetHost, allowing it to be injected back from + // the destructor. + original_rwhv_ = GetRenderWidgetHostImpl()->GetView(); + + // Set a new RenderWidgetHost that allows us control over its size. + rwhv_ = std::make_unique(GetRenderWidgetHostImpl()); + SetView(rwhv_.get()); + SetSize(gfx::Size(320, 240)); + full_screenshotter_ = std::make_unique(); + } + void TearDown() override { + SetView(original_rwhv_); + original_rwhv_ = nullptr; + full_screenshotter_.reset(); + ChromeRenderViewHostTestHarness::TearDown(); + } + + base::expected>, std::string> + CaptureScreenshots(content::WebContents* web_contents) { + base::test::TestFuture< + base::expected>, std::string>> + future; + full_screenshotter()->CaptureScreenshots(web_contents, + future.GetCallback()); + return future.Take(); + } + + FullScreenshotter* full_screenshotter() { return full_screenshotter_.get(); } + + content::RenderWidgetHostImpl* GetRenderWidgetHostImpl() const { + return content::RenderWidgetHostImpl::From( + web_contents()->GetRenderWidgetHostView()->GetRenderWidgetHost()); + } + + void SetView(content::RenderWidgetHostViewBase* rwhv) { + GetRenderWidgetHostImpl()->SetView(rwhv); + } + + void SetSize(const gfx::Size& size) { rwhv_->SetViewBounds(gfx::Rect(size)); } + + void OverrideInterface(MockPaintPreviewRecorder* recorder) { + blink::AssociatedInterfaceProvider* remote_interfaces = + web_contents()->GetPrimaryMainFrame()->GetRemoteAssociatedInterfaces(); + remote_interfaces->OverrideBinderForTesting( + paint_preview::mojom::PaintPreviewRecorder::Name_, + base::BindRepeating(&MockPaintPreviewRecorder::BindRequest, + base::Unretained(recorder))); + } + + std::unique_ptr + CreateCompositorService() { + auto task_runner = base::SingleThreadTaskRunner::GetCurrentDefault(); + return std::unique_ptr( + new MockPaintPreviewCompositorService(task_runner), + base::OnTaskRunnerDeleter(task_runner)); + } + + private: + std::unique_ptr full_screenshotter_; + std::unique_ptr rwhv_; + raw_ptr original_rwhv_ = nullptr; +}; + +TEST_F(FullScreenshotterTest, InvalidWebContentsAndView) { + auto result = CaptureScreenshots(nullptr); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), "The given web contents is no longer valid"); + + for (auto size : {gfx::Size(320, 0), gfx::Size(0, 240)}) { + SetSize(size); + auto result2 = CaptureScreenshots(web_contents()); + ASSERT_FALSE(result2.has_value()); + EXPECT_EQ(result2.error(), "No visible render widget host view available"); + } +} + +TEST_F(FullScreenshotterTest, CaptureFailedAllErrorStates) { + const paint_preview::mojom::PaintPreviewStatus kErrorStatuses[] = { + paint_preview::mojom::PaintPreviewStatus::kAlreadyCapturing, + paint_preview::mojom::PaintPreviewStatus::kCaptureFailed, + paint_preview::mojom::PaintPreviewStatus::kGuidCollision, + paint_preview::mojom::PaintPreviewStatus::kFileCreationError, + // Covers !paint_preview::CaptureResult.capture_success + paint_preview::mojom::PaintPreviewStatus::kPartialSuccess, + paint_preview::mojom::PaintPreviewStatus::kFailed, + }; + + for (auto status : kErrorStatuses) { + MockPaintPreviewRecorder recorder; + recorder.SetResponse( + status, paint_preview::mojom::PaintPreviewCaptureResponse::New()); + OverrideInterface(&recorder); + + auto result = CaptureScreenshots(web_contents()); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), + base::StringPrintf( + "Failed to capture a screenshot (CaptureStatus=%d)", + static_cast(paint_preview::PaintPreviewBaseService:: + CaptureStatus::kCaptureFailed))); + } + // We won't get CaptureStatus::kClientCreationFailed since we check + // WebContents before calling CapturePaintPreview and no + // CaptureStatus::kContentUnsupported because we don't prodvide policy. +} + +TEST_F(FullScreenshotterTest, BeginMainFrameCompositeFailed) { + auto compositor_service = CreateCompositorService(); + full_screenshotter()->InitCompositorServiceForTest( + std::move(compositor_service)); + for (auto status : {paint_preview::mojom::PaintPreviewCompositor:: + BeginCompositeStatus::kCompositingFailure, + paint_preview::mojom::PaintPreviewCompositor:: + BeginCompositeStatus::kDeserializingFailure, + paint_preview::mojom::PaintPreviewCompositor:: + BeginCompositeStatus::kSuccess}) { + AsMockClient(full_screenshotter()->GetCompositorClientForTest()) + ->SetBeginMainFrameResponseStatus(status); + + MockPaintPreviewRecorder recorder; + auto response = paint_preview::mojom::PaintPreviewCaptureResponse::New(); + response->skp.emplace(mojo_base::BigBuffer(true)); + recorder.SetResponse(paint_preview::mojom::PaintPreviewStatus::kOk, + std::move(response)); + OverrideInterface(&recorder); + + auto result = CaptureScreenshots(web_contents()); + ASSERT_FALSE(result.has_value()); + if (status == paint_preview::mojom::PaintPreviewCompositor:: + BeginCompositeStatus::kSuccess) { + EXPECT_EQ(result.error(), "Root frame data not found"); + } else { + EXPECT_EQ(result.error(), "BeginMainFrameComposite failed"); + } + } +} + +TEST_F(FullScreenshotterTest, CompositionSucceeded) { + struct { + gfx::Size viewport_size; + gfx::Size page_size; + size_t num_of_screenshots; + } test_cases[]{ + {gfx::Size(800, 600), gfx::Size(1024, 768), 2u}, + {gfx::Size(1024, 768), gfx::Size(800, 600), 1u}, + {gfx::Size(1024, 768), gfx::Size(1024, 1536), 2u}, + {gfx::Size(1024, 768), gfx::Size(1024, 3072), 4u}, + {gfx::Size(1024, 768), gfx::Size(2048, 768), 1u}, + {gfx::Size(1024, 768), gfx::Size(2048, 1536), 2u}, + {gfx::Size(2560, 1440), gfx::Size(1024, 768), 1u}, + {gfx::Size(2560, 1440), gfx::Size(2560, 7200), 5u}, + }; + for (const auto& test_case : test_cases) { + SCOPED_TRACE(testing::Message() + << "viewport size: " << test_case.viewport_size.ToString() + << "; page size: " << test_case.page_size.ToString() + << "; screenshots number: " << test_case.num_of_screenshots); + SetSize(test_case.viewport_size); + auto compositor_service = CreateCompositorService(); + full_screenshotter()->InitCompositorServiceForTest( + std::move(compositor_service)); + + auto* client = + AsMockClient(full_screenshotter()->GetCompositorClientForTest()); + client->SetBeginMainFrameResponseStatus( + paint_preview::mojom::PaintPreviewCompositor::BeginCompositeStatus:: + kSuccess); + + // Set up frames with root frame data + base::flat_map> + frames; + auto root_frame = paint_preview::mojom::FrameData::New(); + root_frame->scroll_extents = test_case.page_size; + auto token = base::UnguessableToken::Create(); + frames.insert({token, std::move(root_frame)}); + client->SetCompositeResponse(std::move(frames), token); + + // Setup successful capture response + MockPaintPreviewRecorder recorder; + auto response = paint_preview::mojom::PaintPreviewCaptureResponse::New(); + response->skp.emplace(mojo_base::BigBuffer(true)); + recorder.SetResponse(paint_preview::mojom::PaintPreviewStatus::kOk, + std::move(response)); + OverrideInterface(&recorder); + + auto result = CaptureScreenshots(web_contents()); + EXPECT_TRUE(result.has_value()); + EXPECT_EQ(result.value().size(), test_case.num_of_screenshots); + EXPECT_TRUE(std::ranges::any_of( + result.value(), [](const auto& entry) { return !entry.empty(); })); + } +} + +TEST_F(FullScreenshotterTest, BitmapForMainFrameFailed) { + SetSize(gfx::Size(1024, 768)); + auto compositor_service = CreateCompositorService(); + full_screenshotter()->InitCompositorServiceForTest( + std::move(compositor_service)); + for (auto status : + {paint_preview::mojom::PaintPreviewCompositor::BitmapStatus:: + kAllocFailed, + paint_preview::mojom::PaintPreviewCompositor::BitmapStatus:: + kMissingFrame, + paint_preview::mojom::PaintPreviewCompositor::BitmapStatus::kSuccess}) { + auto* client = + AsMockClient(full_screenshotter()->GetCompositorClientForTest()); + client->SetBeginMainFrameResponseStatus( + paint_preview::mojom::PaintPreviewCompositor::BeginCompositeStatus:: + kSuccess); + client->SetBitmapStatus(status); + if (status == + paint_preview::mojom::PaintPreviewCompositor::BitmapStatus::kSuccess) { + client->SetIsEmptyBitmap(true); + } + + // Set up frames with root frame data + base::flat_map> + frames; + auto root_frame = paint_preview::mojom::FrameData::New(); + root_frame->scroll_extents = gfx::Size(800, 600); + auto token = base::UnguessableToken::Create(); + frames.insert({token, std::move(root_frame)}); + client->SetCompositeResponse(std::move(frames), token); + + MockPaintPreviewRecorder recorder; + auto response = paint_preview::mojom::PaintPreviewCaptureResponse::New(); + response->skp.emplace(mojo_base::BigBuffer(true)); + recorder.SetResponse(paint_preview::mojom::PaintPreviewStatus::kOk, + std::move(response)); + OverrideInterface(&recorder); + + auto result = CaptureScreenshots(web_contents()); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error(), + base::StringPrintf("Failed to get bitmap (BitmapStatus=%d)", + static_cast(status))); + } +} + +} // namespace ai_chat diff --git a/browser/ai_chat/upload_file_helper.cc b/browser/ai_chat/upload_file_helper.cc index eb8bf627289..f5b00953c7e 100644 --- a/browser/ai_chat/upload_file_helper.cc +++ b/browser/ai_chat/upload_file_helper.cc @@ -14,14 +14,13 @@ #include "base/no_destructor.h" #include "base/sequence_checker.h" #include "base/task/thread_pool.h" +#include "brave/components/ai_chat/core/browser/utils.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h" #include "chrome/browser/profiles/profile.h" #include "content/public/browser/web_contents.h" #include "services/data_decoder/public/cpp/data_decoder.h" #include "services/data_decoder/public/cpp/decode_image.h" #include "third_party/skia/include/core/SkBitmap.h" -#include "third_party/skia/include/core/SkCanvas.h" -#include "third_party/skia/include/core/SkImage.h" #include "ui/gfx/codec/png_codec.h" #include "ui/shell_dialogs/selected_file_info.h" @@ -37,49 +36,6 @@ namespace ai_chat { namespace { using UploadImageCallback = mojom::AIChatUIHandler::UploadImageCallback; -SkBitmap ScaleBitmap(const SkBitmap& bitmap) { - constexpr int kTargetWidth = 1024; - constexpr int kTargetHeight = 768; - - // Don't need to scale if dimensions are already smaller than target - // dimensions - if (bitmap.width() <= kTargetWidth && bitmap.height() <= kTargetHeight) { - return bitmap; - } - - SkBitmap scaled_bitmap; - scaled_bitmap.allocN32Pixels(kTargetWidth, kTargetHeight); - - SkCanvas canvas(scaled_bitmap); - canvas.clear(SK_ColorTRANSPARENT); - - // Use high-quality scaling options - SkSamplingOptions sampling_options(SkFilterMode::kLinear, - SkMipmapMode::kLinear); - - // Maintain aspect ratio while fitting within target dimensions - float src_aspect = static_cast(bitmap.width()) / bitmap.height(); - float dst_aspect = static_cast(kTargetWidth) / kTargetHeight; - - SkRect dst_rect; - if (src_aspect > dst_aspect) { - // Source is wider - fit to width - float scaled_height = kTargetWidth / src_aspect; - float y_offset = (kTargetHeight - scaled_height) / 2; - dst_rect = SkRect::MakeXYWH(0, y_offset, kTargetWidth, scaled_height); - } else { - // Source is taller - fit to height - float scaled_width = kTargetHeight * src_aspect; - float x_offset = (kTargetWidth - scaled_width) / 2; - dst_rect = SkRect::MakeXYWH(x_offset, 0, scaled_width, kTargetHeight); - } - - // Draw scaled bitmap with high-quality sampling - canvas.drawImageRect(bitmap.asImage(), dst_rect, sampling_options); - - return scaled_bitmap; -} - // base::ReadFileToBytes doesn't handle content uri so we need to read from // base::File which covers content uri. std::optional> ReadFileToBytes( @@ -109,8 +65,8 @@ void OnImageDecoded( } auto encode_image = base::BindOnce( [](const SkBitmap& decoded_bitmap) { - return gfx::PNGCodec::EncodeBGRASkBitmap(ScaleBitmap(decoded_bitmap), - false); + return gfx::PNGCodec::EncodeBGRASkBitmap( + ScaleDownBitmap(decoded_bitmap), false); }, decoded_bitmap); base::ThreadPool::PostTaskAndReplyWithResult(FROM_HERE, {base::MayBlock()}, @@ -201,17 +157,18 @@ void UploadFileHelper::MultiFilesSelected( [](UploadImageCallback callback, std::vector>, std::string>> results) { - std::vector uploaded_images; + std::vector uploaded_files; for (const auto& [image_data, filename] : results) { if (image_data) { - uploaded_images.push_back(mojom::UploadedImage::New( - std::move(filename), image_data->size(), *image_data)); + uploaded_files.push_back(mojom::UploadedFile::New( + std::move(filename), image_data->size(), *image_data, + mojom::UploadedFileType::kImage)); } } std::move(callback).Run( - uploaded_images.empty() + uploaded_files.empty() ? std::nullopt - : std::make_optional(std::move(uploaded_images))); + : std::make_optional(std::move(uploaded_files))); }, std::move(upload_image_callback_))); @@ -288,9 +245,10 @@ void UploadFileHelper::OnImageEncoded( std::move(upload_image_callback_).Run(std::nullopt); return; } - std::vector images; - images.push_back( - mojom::UploadedImage::New(std::move(filename), output->size(), *output)); + std::vector images; + images.push_back(mojom::UploadedFile::New(std::move(filename), output->size(), + *output, + mojom::UploadedFileType::kImage)); std::move(upload_image_callback_).Run(std::make_optional(std::move(images))); } diff --git a/browser/ai_chat/upload_file_helper_unittest.cc b/browser/ai_chat/upload_file_helper_unittest.cc index 9d607720963..e3b780e1545 100644 --- a/browser/ai_chat/upload_file_helper_unittest.cc +++ b/browser/ai_chat/upload_file_helper_unittest.cc @@ -11,6 +11,7 @@ #include "base/memory/ref_counted_memory.h" #include "base/test/test_future.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-forward.h" +#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-shared.h" #include "chrome/browser/ui/chrome_select_file_policy.h" #include "chrome/common/pref_names.h" #include "chrome/test/base/scoped_testing_local_state.h" @@ -46,8 +47,8 @@ class UploadFileHelperTest : public content::RenderViewHostTestHarness { file_helper_ = std::make_unique(web_contents(), profile); } - std::optional> UploadImageSync() { - base::test::TestFuture>> + std::optional> UploadImageSync() { + base::test::TestFuture>> future; file_helper_->UploadImage( std::make_unique(web_contents()), @@ -135,9 +136,9 @@ TEST_F(UploadFileHelperTest, ImageRead) { ASSERT_TRUE(sample_result); ASSERT_EQ(1u, sample_result->size()); EXPECT_EQ((*sample_result)[0]->filename, "sample_png.png"); - EXPECT_EQ((*sample_result)[0]->filesize, - (*sample_result)[0]->image_data.size()); - auto encoded_bitmap = gfx::PNGCodec::Decode((*sample_result)[0]->image_data); + EXPECT_EQ((*sample_result)[0]->filesize, (*sample_result)[0]->data.size()); + EXPECT_EQ((*sample_result)[0]->type, mojom::UploadedFileType::kImage); + auto encoded_bitmap = gfx::PNGCodec::Decode((*sample_result)[0]->data); EXPECT_TRUE(gfx::test::AreBitmapsClose(sample_bitmap, encoded_bitmap, 1)); // Check dimensions are the same. EXPECT_EQ(sample_bitmap.width(), encoded_bitmap.width()); @@ -154,10 +155,10 @@ TEST_F(UploadFileHelperTest, ImageRead) { ASSERT_TRUE(large_result); ASSERT_EQ(1u, large_result->size()); EXPECT_EQ((*large_result)[0]->filename, "large_png.png"); - EXPECT_EQ((*large_result)[0]->filesize, - (*large_result)[0]->image_data.size()); + EXPECT_EQ((*large_result)[0]->filesize, (*large_result)[0]->data.size()); + EXPECT_EQ((*large_result)[0]->type, mojom::UploadedFileType::kImage); EXPECT_LE((*large_result)[0]->filesize, large_png_bytes->size()); - encoded_bitmap = gfx::PNGCodec::Decode((*large_result)[0]->image_data); + encoded_bitmap = gfx::PNGCodec::Decode((*large_result)[0]->data); EXPECT_EQ(1024, encoded_bitmap.width()); EXPECT_EQ(768, encoded_bitmap.height()); @@ -172,15 +173,17 @@ TEST_F(UploadFileHelperTest, ImageRead) { ASSERT_EQ(2u, result->size()); EXPECT_EQ((*result)[0]->filename, "sample_png.png"); - EXPECT_EQ((*result)[0]->filesize, (*result)[0]->image_data.size()); - auto encoded_bitmap1 = gfx::PNGCodec::Decode((*result)[0]->image_data); + EXPECT_EQ((*result)[0]->filesize, (*result)[0]->data.size()); + EXPECT_EQ((*result)[0]->type, mojom::UploadedFileType::kImage); + auto encoded_bitmap1 = gfx::PNGCodec::Decode((*result)[0]->data); EXPECT_TRUE(gfx::test::AreBitmapsClose(sample_bitmap, encoded_bitmap1, 1)); EXPECT_EQ(sample_bitmap.width(), encoded_bitmap1.width()); EXPECT_EQ(sample_bitmap.height(), encoded_bitmap1.height()); EXPECT_EQ((*result)[1]->filename, "large_png.png"); - EXPECT_EQ((*result)[1]->filesize, (*result)[1]->image_data.size()); - auto encoded_bitmap2 = gfx::PNGCodec::Decode((*result)[1]->image_data); + EXPECT_EQ((*result)[1]->filesize, (*result)[1]->data.size()); + EXPECT_EQ((*result)[1]->type, mojom::UploadedFileType::kImage); + auto encoded_bitmap2 = gfx::PNGCodec::Decode((*result)[1]->data); EXPECT_EQ(1024, encoded_bitmap2.width()); EXPECT_EQ(768, encoded_bitmap2.height()); } diff --git a/components/ai_chat/content/browser/BUILD.gn b/components/ai_chat/content/browser/BUILD.gn index 126c207cee5..8e0aa6f37cd 100644 --- a/components/ai_chat/content/browser/BUILD.gn +++ b/components/ai_chat/content/browser/BUILD.gn @@ -13,6 +13,8 @@ static_library("browser") { "ai_chat_tab_helper.h", "ai_chat_throttle.cc", "ai_chat_throttle.h", + "full_screenshotter.cc", + "full_screenshotter.h", "model_service_factory.cc", "model_service_factory.h", "page_content_fetcher.cc", @@ -35,6 +37,9 @@ static_library("browser") { "//components/favicon/content", "//components/favicon/core", "//components/keyed_service/content:content", + "//components/paint_preview/browser", + "//components/paint_preview/common", + "//components/paint_preview/public", "//components/prefs", "//components/strings:components_strings_grit", "//components/user_prefs", diff --git a/components/ai_chat/content/browser/DEPS b/components/ai_chat/content/browser/DEPS index 416379002ca..2d7cbb8127d 100644 --- a/components/ai_chat/content/browser/DEPS +++ b/components/ai_chat/content/browser/DEPS @@ -12,5 +12,7 @@ include_rules = [ "+services/service_manager/public", "+third_party/blink/public/common", "+third_party/blink/public/mojom/permissions", + "+third_party/skia/include", "+ui/base", + "+ui/gfx", ] diff --git a/components/ai_chat/content/browser/ai_chat_tab_helper.cc b/components/ai_chat/content/browser/ai_chat_tab_helper.cc index 2219424501b..cfd2cf56564 100644 --- a/components/ai_chat/content/browser/ai_chat_tab_helper.cc +++ b/components/ai_chat/content/browser/ai_chat_tab_helper.cc @@ -36,6 +36,7 @@ #include "brave/components/ai_chat/core/browser/associated_content_driver.h" #include "brave/components/ai_chat/core/browser/constants.h" #include "brave/components/ai_chat/core/browser/utils.h" +#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h" #include "brave/components/ai_chat/core/common/mojom/page_content_extractor.mojom.h" #include "components/favicon/content/content_favicon_driver.h" #include "components/strings/grit/components_strings.h" @@ -432,6 +433,35 @@ bool AIChatTabHelper::HasOpenAIChatPermission() const { return permission_status.status == content::PermissionStatus::GRANTED; } +void AIChatTabHelper::GetScreenshots( + mojom::ConversationHandler::GetScreenshotsCallback callback) { + full_screenshotter_ = std::make_unique(); + full_screenshotter_->CaptureScreenshots( + web_contents(), + base::BindOnce(&AIChatTabHelper::OnScreenshotsCaptured, + weak_ptr_factory_.GetWeakPtr(), std::move(callback))); +} + +void AIChatTabHelper::OnScreenshotsCaptured( + mojom::ConversationHandler::GetScreenshotsCallback callback, + base::expected>, std::string> result) { + if (result.has_value()) { + std::vector screenshots; + size_t screenshot_index = 0; + for (auto& screenshot : result.value()) { + size_t screenshot_size = screenshot.size(); + screenshots.push_back(mojom::UploadedFile::New( + base::StringPrintf("fullscreenshot_%i.png", screenshot_index++), + screenshot_size, std::move(screenshot), + mojom::UploadedFileType::kScreenshot)); + } + std::move(callback).Run(std::move(screenshots)); + } else { + VLOG(1) << result.error(); + std::move(callback).Run(std::nullopt); + } +} + WEB_CONTENTS_USER_DATA_KEY_IMPL(AIChatTabHelper); } // namespace ai_chat diff --git a/components/ai_chat/content/browser/ai_chat_tab_helper.h b/components/ai_chat/content/browser/ai_chat_tab_helper.h index 086ed126f81..398785f4f6d 100644 --- a/components/ai_chat/content/browser/ai_chat_tab_helper.h +++ b/components/ai_chat/content/browser/ai_chat_tab_helper.h @@ -11,11 +11,13 @@ #include #include #include +#include #include "base/functional/callback.h" #include "base/functional/callback_forward.h" #include "base/memory/raw_ptr.h" #include "base/memory/weak_ptr.h" +#include "brave/components/ai_chat/content/browser/full_screenshotter.h" #include "brave/components/ai_chat/core/browser/associated_content_driver.h" #include "brave/components/ai_chat/core/browser/conversation_handler.h" #include "brave/components/ai_chat/core/common/mojom/page_content_extractor.mojom.h" @@ -178,6 +180,13 @@ class AIChatTabHelper : public content::WebContentsObserver, bool HasOpenAIChatPermission() const override; + void GetScreenshots( + mojom::ConversationHandler::GetScreenshotsCallback callback) override; + + void OnScreenshotsCaptured( + mojom::ConversationHandler::GetScreenshotsCallback callback, + base::expected>, std::string>); + void OnFetchPageContentComplete(GetPageContentCallback callback, std::string content, bool is_video, @@ -221,6 +230,8 @@ class AIChatTabHelper : public content::WebContentsObserver, // A scoper only used for PDF viewing. std::unique_ptr scoped_accessibility_mode_; + std::unique_ptr full_screenshotter_; + mojo::AssociatedReceiver page_content_extractor_receiver_{this}; diff --git a/components/ai_chat/content/browser/full_screenshotter.cc b/components/ai_chat/content/browser/full_screenshotter.cc new file mode 100644 index 00000000000..61c24ef0078 --- /dev/null +++ b/components/ai_chat/content/browser/full_screenshotter.cc @@ -0,0 +1,254 @@ +/* Copyright (c) 2024 The Brave Authors. All rights reserved. + * This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this file, + * You can obtain one at https://mozilla.org/MPL/2.0/. */ + +#include "brave/components/ai_chat/content/browser/full_screenshotter.h" + +#include +#include + +#include "base/strings/stringprintf.h" +#include "base/task/sequenced_task_runner.h" +#include "base/task/task_traits.h" +#include "base/task/thread_pool.h" +#include "base/types/expected.h" +#include "brave/components/ai_chat/core/browser/utils.h" +#include "components/paint_preview/browser/compositor_utils.h" +#include "components/paint_preview/browser/paint_preview_base_service.h" +#include "components/paint_preview/common/recording_map.h" +#include "content/public/browser/render_widget_host_view.h" +#include "mojo/public/cpp/base/proto_wrapper.h" +#include "third_party/skia/include/core/SkBitmap.h" +#include "ui/gfx/codec/png_codec.h" +#include "ui/gfx/geometry/rect.h" + +namespace ai_chat { + +FullScreenshotter::FullScreenshotter() + : paint_preview::PaintPreviewBaseService( + /*file_mixin=*/nullptr, // in-memory captures + /*policy=*/nullptr, // all content is deemed amenable + /*is_off_the_record=*/true), + paint_preview_compositor_service_(nullptr, + base::OnTaskRunnerDeleter(nullptr)), + paint_preview_compositor_client_(nullptr, + base::OnTaskRunnerDeleter(nullptr)) {} + +FullScreenshotter::~FullScreenshotter() = default; + +FullScreenshotter::PendingScreenshots::PendingScreenshots() = default; +FullScreenshotter::PendingScreenshots::~PendingScreenshots() = default; + +void FullScreenshotter::CaptureScreenshots( + const raw_ptr web_contents, + CaptureScreenshotsCallback callback) { + if (!web_contents) { + std::move(callback).Run( + base::unexpected("The given web contents is no longer valid")); + return; + } + auto* view = web_contents->GetRenderWidgetHostView(); + if (!view || view->GetVisibleViewportSize().IsEmpty()) { + std::move(callback).Run( + base::unexpected("No visible render widget host view available")); + return; + } + viewport_bounds_ = view->GetVisibleViewportSize(); + + // Start capturing via Paint Preview. + CaptureParams capture_params; + capture_params.web_contents = web_contents; + capture_params.persistence = + paint_preview::RecordingPersistence::kMemoryBuffer; + CapturePaintPreview( + capture_params, + base::BindOnce(&FullScreenshotter::OnScreenshotCaptured, + weak_ptr_factory_.GetWeakPtr(), std::move(callback))); +} + +void FullScreenshotter::InitCompositorServiceForTest( + std::unique_ptr service) { + paint_preview_compositor_service_ = std::move(service); + paint_preview_compositor_client_ = + paint_preview_compositor_service_->CreateCompositor(base::DoNothing()); +} + +void FullScreenshotter::OnScreenshotCaptured( + CaptureScreenshotsCallback callback, + paint_preview::PaintPreviewBaseService::CaptureStatus status, + std::unique_ptr result) { + if (status != PaintPreviewBaseService::CaptureStatus::kOk || + !result->capture_success) { + std::move(callback).Run(base::unexpected( + base::StringPrintf("Failed to capture a screenshot (CaptureStatus=%d)", + static_cast(status)))); + return; + } + + if (!paint_preview_compositor_client_) { + if (!paint_preview_compositor_service_) { + paint_preview_compositor_service_ = paint_preview::StartCompositorService( + base::BindOnce(&FullScreenshotter::OnCompositorServiceDisconnected, + weak_ptr_factory_.GetWeakPtr())); + } + paint_preview_compositor_client_ = + paint_preview_compositor_service_->CreateCompositor( + base::BindOnce(&FullScreenshotter::SendCompositeRequest, + weak_ptr_factory_.GetWeakPtr(), std::move(callback), + PrepareCompositeRequest(std::move(result)))); + } else { + SendCompositeRequest(std::move(callback), + PrepareCompositeRequest(std::move(result))); + } +} + +paint_preview::mojom::PaintPreviewBeginCompositeRequestPtr +FullScreenshotter::PrepareCompositeRequest( + std::unique_ptr capture_result) { + paint_preview::mojom::PaintPreviewBeginCompositeRequestPtr + begin_composite_request = + paint_preview::mojom::PaintPreviewBeginCompositeRequest::New(); + std::pair + map_and_proto = paint_preview::RecordingMapFromCaptureResult( + std::move(*capture_result)); + begin_composite_request->recording_map = std::move(map_and_proto.first); + if (begin_composite_request->recording_map.empty()) { + VLOG(2) << "Captured an empty screenshot"; + return nullptr; + } + begin_composite_request->preview = + mojo_base::ProtoWrapper(std::move(map_and_proto.second)); + return begin_composite_request; +} + +void FullScreenshotter::SendCompositeRequest( + CaptureScreenshotsCallback callback, + paint_preview::mojom::PaintPreviewBeginCompositeRequestPtr + begin_composite_request) { + if (!begin_composite_request) { + std::move(callback).Run( + base::unexpected("Invalid begin_composite_request")); + return; + } + + CHECK(paint_preview_compositor_client_); + paint_preview_compositor_client_->BeginMainFrameComposite( + std::move(begin_composite_request), + base::BindOnce(&FullScreenshotter::OnCompositeFinished, + weak_ptr_factory_.GetWeakPtr(), std::move(callback))); +} + +void FullScreenshotter::OnCompositorServiceDisconnected() { + VLOG(2) << "Compositor service is disconnected"; + paint_preview_compositor_client_.reset(); + paint_preview_compositor_service_.reset(); +} + +void FullScreenshotter::OnCompositeFinished( + CaptureScreenshotsCallback callback, + paint_preview::mojom::PaintPreviewCompositor::BeginCompositeStatus status, + paint_preview::mojom::PaintPreviewBeginCompositeResponsePtr response) { + if (status != paint_preview::mojom::PaintPreviewCompositor:: + BeginCompositeStatus::kSuccess) { + std::move(callback).Run(base::unexpected("BeginMainFrameComposite failed")); + return; + } + + if (!response->frames.contains(response->root_frame_guid)) { + std::move(callback).Run(base::unexpected("Root frame data not found")); + return; + } + + const auto& frame_data = response->frames[response->root_frame_guid]; + const auto& content_size = frame_data->scroll_extents; + + auto pending = std::make_unique(); + + // Calculate number of full viewport screenshots needed + int total_height = content_size.height(); + int viewport_height = viewport_bounds_.height(); + int num_screenshots = (total_height + viewport_height - 1) / + viewport_height; // Ceiling division + + pending->completed_images.resize(num_screenshots); + + // Queue up screenshot rectangles + for (int i = 0; i < num_screenshots; ++i) { + int y = i * viewport_height; + int height = std::min(viewport_height, total_height - y); + pending->remaining_rects.emplace(0, y, content_size.width(), height); + } + + pending->callback = std::move(callback); + CaptureNextScreenshot(std::move(pending)); +} + +void FullScreenshotter::CaptureNextScreenshot( + std::unique_ptr pending) { + if (pending->remaining_rects.empty()) { + // All screenshots captured, return results + std::move(pending->callback) + .Run(base::ok(std::move(pending->completed_images))); + return; + } + + // Take the next rect to capture + gfx::Rect capture_rect = pending->remaining_rects.front(); + pending->remaining_rects.pop(); + + paint_preview_compositor_client_->BitmapForMainFrame( + capture_rect, 1, + base::BindOnce(&FullScreenshotter::OnBitmapReceived, + weak_ptr_factory_.GetWeakPtr(), std::move(pending), + pending->completed_images.size() - + pending->remaining_rects.size() - 1)); +} + +void FullScreenshotter::OnBitmapReceived( + std::unique_ptr pending, + size_t index, + paint_preview::mojom::PaintPreviewCompositor::BitmapStatus status, + const SkBitmap& bitmap) { + if (status != paint_preview::mojom::PaintPreviewCompositor::BitmapStatus:: + kSuccess || + bitmap.empty()) { + std::move(pending->callback) + .Run(base::unexpected( + base::StringPrintf("Failed to get bitmap (BitmapStatus=%d)", + static_cast(status)))); + return; + } + + base::ThreadPool::PostTaskAndReplyWithResult( + FROM_HERE, {base::MayBlock()}, + base::BindOnce(&FullScreenshotter::EncodeBitmap, ScaleDownBitmap(bitmap)), + base::BindOnce(&FullScreenshotter::OnBitmapEncoded, + weak_ptr_factory_.GetWeakPtr(), std::move(pending), + index)); +} + +// static +base::expected, std::string> +FullScreenshotter::EncodeBitmap(const SkBitmap& bitmap) { + auto data = gfx::PNGCodec::EncodeBGRASkBitmap(bitmap, false); + if (!data) { + return base::unexpected("Failed to encode the bitmap"); + } + return base::ok(*data); +} + +void FullScreenshotter::OnBitmapEncoded( + std::unique_ptr pending, + size_t index, + base::expected, std::string> result) { + if (!result.has_value()) { + std::move(pending->callback).Run(base::unexpected(result.error())); + return; + } + pending->completed_images[index] = std::move(*result); + CaptureNextScreenshot(std::move(pending)); +} + +} // namespace ai_chat diff --git a/components/ai_chat/content/browser/full_screenshotter.h b/components/ai_chat/content/browser/full_screenshotter.h new file mode 100644 index 00000000000..c6cd0f4fa57 --- /dev/null +++ b/components/ai_chat/content/browser/full_screenshotter.h @@ -0,0 +1,105 @@ +/* Copyright (c) 2024 The Brave Authors. All rights reserved. + * This Source Code Form is subject to the terms of the Mozilla Public + * License, v. 2.0. If a copy of the MPL was not distributed with this file, + * You can obtain one at https://mozilla.org/MPL/2.0/. */ + +#ifndef BRAVE_COMPONENTS_AI_CHAT_CONTENT_BROWSER_FULL_SCREENSHOTTER_H_ +#define BRAVE_COMPONENTS_AI_CHAT_CONTENT_BROWSER_FULL_SCREENSHOTTER_H_ + +#include +#include +#include +#include + +#include "base/memory/raw_ptr.h" +#include "base/memory/weak_ptr.h" +#include "base/task/sequenced_task_runner.h" +#include "base/types/expected.h" +#include "components/paint_preview/browser/paint_preview_base_service.h" +#include "components/paint_preview/public/paint_preview_compositor_service.h" +#include "content/public/browser/web_contents.h" +#include "ui/gfx/geometry/rect.h" +#include "ui/gfx/geometry/size.h" + +namespace ai_chat { + +// This class uses paint preview service and compositor service to capture +// screenshot of a web_contents and split it into multiple ones based on the +// viewport height. If a single screenshot is larger than 1024x768, it will be +// scaled down to that resolution. +class FullScreenshotter : public paint_preview::PaintPreviewBaseService { + public: + FullScreenshotter(); + FullScreenshotter(const FullScreenshotter&) = delete; + FullScreenshotter& operator=(const FullScreenshotter&) = delete; + ~FullScreenshotter() override; + + // + using CaptureScreenshotsCallback = base::OnceCallback>, std::string>)>; + void CaptureScreenshots(const raw_ptr web_contents, + CaptureScreenshotsCallback callback); + + void InitCompositorServiceForTest( + std::unique_ptr service); + + paint_preview::PaintPreviewCompositorClient* GetCompositorClientForTest() { + return paint_preview_compositor_client_.get(); + } + + private: + void OnScreenshotCaptured( + CaptureScreenshotsCallback callback, + paint_preview::PaintPreviewBaseService::CaptureStatus status, + std::unique_ptr result); + void OnCompositorServiceDisconnected(); + paint_preview::mojom::PaintPreviewBeginCompositeRequestPtr + PrepareCompositeRequest( + std::unique_ptr capture_result); + void SendCompositeRequest( + CaptureScreenshotsCallback callback, + paint_preview::mojom::PaintPreviewBeginCompositeRequestPtr + begin_composite_request); + void RequestBitmapForMainFrame(CaptureScreenshotsCallback callback); + void OnCompositeFinished( + CaptureScreenshotsCallback callback, + paint_preview::mojom::PaintPreviewCompositor::BeginCompositeStatus status, + paint_preview::mojom::PaintPreviewBeginCompositeResponsePtr response); + + struct PendingScreenshots { + PendingScreenshots(); + ~PendingScreenshots(); + std::queue remaining_rects; + std::vector> completed_images; + CaptureScreenshotsCallback callback; + }; + + static base::expected, std::string> EncodeBitmap( + const SkBitmap& bitmap); + void OnBitmapReceived( + std::unique_ptr pending, + size_t index, + paint_preview::mojom::PaintPreviewCompositor::BitmapStatus status, + const SkBitmap& bitmap); + void OnBitmapEncoded( + std::unique_ptr pending, + size_t index, + base::expected, std::string> result); + + void CaptureNextScreenshot(std::unique_ptr pending); + + std::unique_ptr + paint_preview_compositor_service_; + std::unique_ptr + paint_preview_compositor_client_; + + gfx::Size viewport_bounds_; + + base::WeakPtrFactory weak_ptr_factory_{this}; +}; + +} // namespace ai_chat +#endif // BRAVE_COMPONENTS_AI_CHAT_CONTENT_BROWSER_FULL_SCREENSHOTTER_H_ diff --git a/components/ai_chat/core/browser/DEPS b/components/ai_chat/core/browser/DEPS index bdb704bc29c..79d54ad0128 100644 --- a/components/ai_chat/core/browser/DEPS +++ b/components/ai_chat/core/browser/DEPS @@ -18,6 +18,7 @@ include_rules = [ "+third_party/tflite", "+third_party/tflite_support", "+ui/base", + "+ui/gfx/image", ] specific_include_rules = { diff --git a/components/ai_chat/core/browser/ai_chat_database.cc b/components/ai_chat/core/browser/ai_chat_database.cc index ddd9f0f63de..01300185b8d 100644 --- a/components/ai_chat/core/browser/ai_chat_database.cc +++ b/components/ai_chat/core/browser/ai_chat_database.cc @@ -69,6 +69,15 @@ bool MigrateFrom2To3(sql::Database* db) { trimmed_tokens_statement.Run(); } +bool MigrateFrom3to4(sql::Database* db) { + static constexpr char kAddTypeColumnQuery[] = + "ALTER TABLE conversation_entry_uploaded_files ADD COLUMN type INTEGER " + "DEFAULT 0"; + sql::Statement statement(db->GetUniqueStatement(kAddTypeColumnQuery)); + + return statement.is_valid() && statement.Run(); +} + void SerializeWebSourcesEvent(const mojom::WebSourcesEventPtr& mojom_event, store::WebSourcesEventProto* proto_event) { proto_event->clear_sources(); @@ -125,7 +134,7 @@ constexpr int kLowestSupportedDatabaseVersion = 1; constexpr int kCompatibleDatabaseVersionNumber = 1; // Current version of the database. Increase if breaking changes are made. -constexpr int kCurrentDatabaseVersion = 3; +constexpr int kCurrentDatabaseVersion = 4; AIChatDatabase::AIChatDatabase(const base::FilePath& db_file_path, os_crypt_async::Encryptor encryptor) @@ -204,6 +213,15 @@ sql::InitStatus AIChatDatabase::InitInternal() { } current_version = 3; } + if (migration_success && current_version == 3) { + migration_success = MigrateFrom3to4(&GetDB()); + if (migration_success) { + migration_success = meta_table.SetCompatibleVersionNumber( + kCompatibleDatabaseVersionNumber) && + meta_table.SetVersionNumber(4); + } + current_version = 4; + } // Migration unsuccessful, raze the database and re-init if (!migration_success) { if (db_.Raze()) { @@ -447,9 +465,9 @@ std::vector AIChatDatabase::GetConversationEntries( } } - // Uploaded images + // Uploaded files sql::Statement uploaded_file_statement( - GetDB().GetUniqueStatement("SELECT filename, filesize, data" + GetDB().GetUniqueStatement("SELECT filename, filesize, data, type" " FROM conversation_entry_uploaded_files" " WHERE conversation_entry_uuid=?" " ORDER BY file_order ASC")); @@ -458,16 +476,18 @@ std::vector AIChatDatabase::GetConversationEntries( while (uploaded_file_statement.Step()) { auto filename = DecryptColumnToString(uploaded_file_statement, 0); int64_t filesize = uploaded_file_statement.ColumnInt64(1); - auto decrypted_image_str = + auto decrypted_bytes_str = DecryptColumnToString(uploaded_file_statement, 2); - base::span image_bytes = - base::as_byte_span(decrypted_image_str); - std::vector image_data(image_bytes.begin(), image_bytes.end()); - if (!entry->uploaded_images) { - entry->uploaded_images = std::vector{}; + base::span raw_bytes = + base::as_byte_span(decrypted_bytes_str); + std::vector data(raw_bytes.begin(), raw_bytes.end()); + auto type = static_cast( + uploaded_file_statement.ColumnInt(3)); + if (!entry->uploaded_files) { + entry->uploaded_files = std::vector{}; } - entry->uploaded_images->emplace_back(mojom::UploadedImage::New( - std::move(filename), filesize, std::move(image_data))); + entry->uploaded_files->emplace_back(mojom::UploadedFile::New( + std::move(filename), filesize, std::move(data), type)); } // root entry or edited entry @@ -834,29 +854,30 @@ bool AIChatDatabase::AddConversationEntry( } } - if (entry->uploaded_images.has_value()) { - for (size_t i = 0; i < entry->uploaded_images->size(); ++i) { - const mojom::UploadedImagePtr& uploaded_image = - entry->uploaded_images->at(i); + if (entry->uploaded_files.has_value()) { + for (size_t i = 0; i < entry->uploaded_files->size(); ++i) { + const mojom::UploadedFilePtr& uploaded_file = + entry->uploaded_files->at(i); sql::Statement uploaded_file_statement(GetDB().GetCachedStatement( SQL_FROM_HERE, "INSERT INTO conversation_entry_uploaded_files" - "(file_order, filename, filesize, data," + "(file_order, filename, filesize, data, type," " conversation_entry_uuid)" - " VALUES(?, ?, ?, ?, ?)")); + " VALUES(?, ?, ?, ?, ?, ?)")); CHECK(uploaded_file_statement.is_valid()); uploaded_file_statement.BindInt(0, static_cast(i)); if (!BindAndEncryptString(uploaded_file_statement, 1, - uploaded_image->filename)) { + uploaded_file->filename)) { return false; } - uploaded_file_statement.BindInt64(2, uploaded_image->filesize); + uploaded_file_statement.BindInt64(2, uploaded_file->filesize); if (!BindAndEncryptString( uploaded_file_statement, 3, - base::as_string_view(base::span(uploaded_image->image_data)))) { + base::as_string_view(base::span(uploaded_file->data)))) { return false; } - uploaded_file_statement.BindString(4, entry->uuid.value()); + uploaded_file_statement.BindInt(4, static_cast(uploaded_file->type)); + uploaded_file_statement.BindString(5, entry->uuid.value()); uploaded_file_statement.Run(); } } @@ -1376,6 +1397,8 @@ bool AIChatDatabase::CreateSchema() { "filesize INTEGER NOT NULL," // encrypted file byte data "data BLOB NOT NULL," + // mojom::UploadedFileType + "type INTEGER NOT NULL," "PRIMARY KEY(conversation_entry_uuid, file_order)" ")"; CHECK(GetDB().IsSQLValid(kCreateUploadedFilesTableQuery)); diff --git a/components/ai_chat/core/browser/ai_chat_database_unittest.cc b/components/ai_chat/core/browser/ai_chat_database_unittest.cc index b4c5de67b0d..45bcef353ab 100644 --- a/components/ai_chat/core/browser/ai_chat_database_unittest.cc +++ b/components/ai_chat/core/browser/ai_chat_database_unittest.cc @@ -359,8 +359,8 @@ TEST_P(AIChatDatabaseTest, WebSourcesEvent_Invalid) { history); } -TEST_P(AIChatDatabaseTest, UploadImage) { - constexpr char kUUID[] = "upload_image_uuid"; +TEST_P(AIChatDatabaseTest, UploadFile) { + constexpr char kUUID[] = "upload_file_uuid"; mojom::ConversationPtr metadata = mojom::Conversation::New( kUUID, "title", base::Time::Now() - base::Hours(2), true, std::nullopt, 0, 0, nullptr); @@ -692,10 +692,15 @@ class AIChatDatabaseMigrationTest : public testing::Test, std::unique_ptr db_; }; -INSTANTIATE_TEST_SUITE_P(, - AIChatDatabaseMigrationTest, - testing::Range(kLowestSupportedDatabaseVersion, - kCurrentDatabaseVersion)); +INSTANTIATE_TEST_SUITE_P( + , + AIChatDatabaseMigrationTest, + testing::Range(kLowestSupportedDatabaseVersion, kCurrentDatabaseVersion), + [](const testing::TestParamInfo& + info) { + return base::StringPrintf("From_v%d_to_v%d", info.param, + kCurrentDatabaseVersion); + }); // Tests the migration of the database from version() to kCurrentVersionNumber TEST_P(AIChatDatabaseMigrationTest, MigrationToVCurrent) { @@ -785,6 +790,48 @@ TEST_P(AIChatDatabaseMigrationTest, MigrationToVCurrent) { EXPECT_EQ(test_conversation->total_tokens, expected_total_tokens); EXPECT_EQ(test_conversation->trimmed_tokens, expected_trimmed_tokens); } + + // V4 Specific Migration checks + { + if (version() == 3) { + auto conversation_with_image = + db_->GetConversationData("1ae484fe-ab33-4f42-8813-14080e4addc1"); + ASSERT_TRUE(conversation_with_image); + ASSERT_EQ(conversation_with_image->entries.size(), 2u); + ASSERT_TRUE(conversation_with_image->entries[0]->uploaded_files); + ASSERT_TRUE(conversation_with_image->entries[1]->uploaded_files); + ASSERT_EQ(conversation_with_image->entries[0]->uploaded_files->size(), + 2u); + ASSERT_EQ(conversation_with_image->entries[1]->uploaded_files->size(), + 1u); + EXPECT_EQ( + conversation_with_image->entries[0]->uploaded_files->at(0)->type, + mojom::UploadedFileType::kImage); + EXPECT_EQ( + conversation_with_image->entries[0]->uploaded_files->at(1)->type, + mojom::UploadedFileType::kImage); + EXPECT_EQ( + conversation_with_image->entries[1]->uploaded_files->at(0)->type, + mojom::UploadedFileType::kImage); + } + + // Verify the newly added entry with files after migration have `type` + // persisted. + auto history = CreateSampleChatHistory(1u, 0, 3u); + EXPECT_TRUE( + db_->AddConversationEntry("migrationtest2", history[0]->Clone())); + auto test_conversation = db_->GetConversationData("migrationtest2"); + ASSERT_EQ(test_conversation->entries.size(), 2u); + ASSERT_TRUE(test_conversation->entries[1]->uploaded_files); + ASSERT_TRUE(history[0]->uploaded_files); + EXPECT_EQ(test_conversation->entries[1]->uploaded_files->size(), + history[0]->uploaded_files->size()); + for (size_t i = 0; + i < test_conversation->entries[1]->uploaded_files->size(); ++i) { + EXPECT_EQ(test_conversation->entries[1]->uploaded_files->at(i)->type, + history[0]->uploaded_files->at(i)->type); + } + } } } // namespace ai_chat diff --git a/components/ai_chat/core/browser/ai_chat_service_unittest.cc b/components/ai_chat/core/browser/ai_chat_service_unittest.cc index f4fe5bdacee..c5bbd6db6a7 100644 --- a/components/ai_chat/core/browser/ai_chat_service_unittest.cc +++ b/components/ai_chat/core/browser/ai_chat_service_unittest.cc @@ -117,7 +117,10 @@ class MockConversationHandlerClient : public mojom::ConversationUI { conversation_ui_receiver_.reset(); } - MOCK_METHOD(void, OnConversationHistoryUpdate, (), (override)); + MOCK_METHOD(void, + OnConversationHistoryUpdate, + (mojom::ConversationTurnPtr), + (override)); MOCK_METHOD(void, OnAPIRequestInProgress, (bool), (override)); diff --git a/components/ai_chat/core/browser/conversation_handler.cc b/components/ai_chat/core/browser/conversation_handler.cc index adc2d0de474..000f40cee60 100644 --- a/components/ai_chat/core/browser/conversation_handler.cc +++ b/components/ai_chat/core/browser/conversation_handler.cc @@ -101,6 +101,11 @@ bool AssociatedContentDelegate::HasOpenAIChatPermission() const { return false; } +void AssociatedContentDelegate::GetScreenshots( + mojom::ConversationHandler::GetScreenshotsCallback callback) { + std::move(callback).Run(std::nullopt); +} + void AssociatedContentDelegate::GetTopSimilarityWithPromptTilContextLimit( const std::string& prompt, const std::string& text, @@ -653,12 +658,17 @@ void ConversationHandler::GetIsRequestInProgress( void ConversationHandler::SubmitHumanConversationEntry( const std::string& input, - std::optional> uploaded_images) { + std::optional> uploaded_files) { DCHECK(!is_request_in_progress_) << "Should not be able to submit more" << "than a single human conversation turn at a time."; - if (uploaded_images && !uploaded_images->empty()) { + if (uploaded_files && !uploaded_files->empty() && + std::ranges::any_of( + uploaded_files.value(), [](const mojom::UploadedFilePtr& file) { + return file->type == mojom::UploadedFileType::kImage || + file->type == mojom::UploadedFileType::kScreenshot; + })) { auto* current_model = model_service_->GetModel(metadata_->model_key.value_or("").empty() ? model_service_->GetDefaultModelKey() @@ -673,7 +683,7 @@ void ConversationHandler::SubmitHumanConversationEntry( std::nullopt, CharacterType::HUMAN, mojom::ActionType::QUERY, input, std::nullopt /* prompt */, std::nullopt /* selected_text */, std::nullopt /* events */, base::Time::Now(), std::nullopt /* edits */, - std::move(uploaded_images), false); + std::move(uploaded_files), false); SubmitHumanConversationEntry(std::move(turn)); } @@ -704,7 +714,7 @@ void ConversationHandler::SubmitHumanConversationEntry( pending_conversation_entry_ = std::move(turn); // Pending entry is added to conversation history when asked for // so notify observers. - OnHistoryUpdate(); + OnHistoryUpdate(nullptr); return; } DCHECK(latest_turn->character_type == mojom::CharacterType::HUMAN); @@ -1280,7 +1290,7 @@ void ConversationHandler::UpdateOrCreateLastAssistantEntry( entry->events->push_back(std::move(event)); // Update clients for partial entries but not observers, who will get notified // when we know this is a complete entry. - OnHistoryUpdate(); + OnHistoryUpdate(entry.Clone()); } void ConversationHandler::MaybeSeedOrClearSuggestions() { @@ -1373,7 +1383,7 @@ void ConversationHandler::MaybeFetchOrClearContentStagedConversation() { return turn->from_brave_search_SERP; }); if (num_entries != chat_history_.size()) { - OnHistoryUpdate(); + OnHistoryUpdate(nullptr); } return; } @@ -1475,7 +1485,7 @@ void ConversationHandler::OnGetRefinedPageContent( if (last_turn->events && !last_turn->events->empty() && last_turn->events->back()->is_page_content_refine_event()) { chat_history_.pop_back(); - OnHistoryUpdate(); + OnHistoryUpdate(nullptr); } else { VLOG(1) << "last entry should be page content refine event"; } @@ -1608,20 +1618,20 @@ void ConversationHandler::OnModelDataChanged() { OnStateForConversationEntriesChanged(); } -void ConversationHandler::OnHistoryUpdate() { +void ConversationHandler::OnHistoryUpdate(mojom::ConversationTurnPtr entry) { // TODO(petemill): Provide the updated converation history item so that // we don't need to clone every entry. for (auto& client : conversation_ui_handlers_) { - client->OnConversationHistoryUpdate(); + client->OnConversationHistoryUpdate(entry ? entry.Clone() : nullptr); } for (auto& client : untrusted_conversation_ui_handlers_) { - client->OnConversationHistoryUpdate(); + client->OnConversationHistoryUpdate(entry ? entry.Clone() : nullptr); } } void ConversationHandler::OnConversationEntryRemoved( std::optional entry_uuid) { - OnHistoryUpdate(); + OnHistoryUpdate(nullptr); if (!entry_uuid.has_value()) { return; } @@ -1634,7 +1644,7 @@ void ConversationHandler::OnConversationEntryAdded( mojom::ConversationTurnPtr& entry) { // Only notify about staged entries once we have the first staged entry if (entry->from_brave_search_SERP) { - OnHistoryUpdate(); + OnHistoryUpdate(nullptr); return; } std::optional associated_content_value; @@ -1658,13 +1668,13 @@ void ConversationHandler::OnConversationEntryAdded( associated_content_value); } } - OnHistoryUpdate(); + OnHistoryUpdate(nullptr); return; } for (auto& observer : observers_) { observer.OnConversationEntryAdded(this, entry, associated_content_value); } - OnHistoryUpdate(); + OnHistoryUpdate(entry.Clone()); } int ConversationHandler::GetContentUsedPercentage() { @@ -1820,6 +1830,14 @@ size_t ConversationHandler::GetConversationHistorySize() { return GetConversationHistory().size(); } +void ConversationHandler::GetScreenshots(GetScreenshotsCallback callback) { + if (associated_content_delegate_) { + associated_content_delegate_->GetScreenshots(std::move(callback)); + } else { + std::move(callback).Run(std::nullopt); + } +} + bool ConversationHandler::should_send_page_contents() const { return should_send_page_contents_; } diff --git a/components/ai_chat/core/browser/conversation_handler.h b/components/ai_chat/core/browser/conversation_handler.h index cbcfdf00584..9cd783ffca5 100644 --- a/components/ai_chat/core/browser/conversation_handler.h +++ b/components/ai_chat/core/browser/conversation_handler.h @@ -116,6 +116,8 @@ class ConversationHandler : public mojom::ConversationHandler, // Signifies whether the content has permission to open a conversation's UI // within the browser. virtual bool HasOpenAIChatPermission() const; + virtual void GetScreenshots( + mojom::ConversationHandler::GetScreenshotsCallback callback); void GetTopSimilarityWithPromptTilContextLimit( const std::string& prompt, @@ -259,7 +261,7 @@ class ConversationHandler : public mojom::ConversationHandler, void GetIsRequestInProgress(GetIsRequestInProgressCallback callback) override; void SubmitHumanConversationEntry( const std::string& input, - std::optional> uploaded_images) + std::optional> uploaded_files) override; void SubmitHumanConversationEntry(mojom::ConversationTurnPtr turn); void SubmitHumanConversationEntryWithAction( @@ -295,6 +297,7 @@ class ConversationHandler : public mojom::ConversationHandler, void OnAssociatedContentTitleChanged(); void OnUserOptedIn(); size_t GetConversationHistorySize() override; + void GetScreenshots(GetScreenshotsCallback callback) override; // Some associated content may provide some conversation that the user wants // to continue, e.g. Brave Search. @@ -425,7 +428,7 @@ class ConversationHandler : public mojom::ConversationHandler, void OnModelDataChanged(); void OnConversationDeleted(); - void OnHistoryUpdate(); + void OnHistoryUpdate(mojom::ConversationTurnPtr entry); void OnConversationEntryAdded(mojom::ConversationTurnPtr& entry); void OnConversationEntryRemoved(std::optional turn_id); void OnSuggestedQuestionsChanged(); diff --git a/components/ai_chat/core/browser/conversation_handler_unittest.cc b/components/ai_chat/core/browser/conversation_handler_unittest.cc index 1888465703a..10eea5c8b08 100644 --- a/components/ai_chat/core/browser/conversation_handler_unittest.cc +++ b/components/ai_chat/core/browser/conversation_handler_unittest.cc @@ -90,7 +90,10 @@ class MockConversationHandlerClient : public mojom::ConversationUI { conversation_ui_receiver_.reset(); } - MOCK_METHOD(void, OnConversationHistoryUpdate, (), (override)); + MOCK_METHOD(void, + OnConversationHistoryUpdate, + (const mojom::ConversationTurnPtr), + (override)); MOCK_METHOD(void, OnAPIRequestInProgress, (bool), (override)); @@ -268,8 +271,10 @@ class ConversationHandlerUnitTest : public testing::Test { } // Pair of text and whether it's from Brave Search SERP - void SetupHistory(std::vector> entries) { + std::vector SetupHistory( + std::vector> entries) { std::vector history; + std::vector expected_history; for (size_t i = 0; i < entries.size(); i++) { bool is_human = i % 2 == 0; @@ -290,9 +295,11 @@ class ConversationHandlerUnitTest : public testing::Test { base::Time::Now(), std::nullopt /* edits */, std::nullopt /* uploaed_images */, entries[i].second /* from_brave_search_SERP */); + expected_history.push_back(entry.Clone()); history.push_back(std::move(entry)); } conversation_handler_->SetChatHistoryForTesting(std::move(history)); + return expected_history; } protected: @@ -344,6 +351,24 @@ MATCHER_P(LastTurnHasSelectedText, expected_text, "") { return !arg.empty() && arg.back()->selected_text == expected_text; } +// Can't use mojo::Equals with ::testing::Truly +// because we have uuid and created_time fields +MATCHER_P(TurnEq, expected_turn, "") { + if (!arg && !expected_turn) { + return true; + } + return arg && expected_turn && + arg->character_type == expected_turn->character_type && + arg->action_type == expected_turn->action_type && + arg->text == expected_turn->text && + arg->prompt == expected_turn->prompt && + arg->selected_text == expected_turn->selected_text && + arg->events == expected_turn->events && + arg->edits == expected_turn->edits && + arg->uploaded_files == expected_turn->uploaded_files && + arg->from_brave_search_SERP == expected_turn->from_brave_search_SERP; +} + TEST_F(ConversationHandlerUnitTest, GetState) { NiceMock client(conversation_handler_.get()); for (bool should_send_content : {false, true}) { @@ -417,13 +442,35 @@ TEST_F(ConversationHandlerUnitTest, SubmitSelectedText) { EXPECT_FALSE(should_send_page_contents); })); + std::vector expected_history; + + expected_history.push_back(mojom::ConversationTurn::New( + std::nullopt, mojom::CharacterType::HUMAN, + mojom::ActionType::SUMMARIZE_SELECTED_TEXT, expected_turn_text, + std::nullopt, selected_text, std::nullopt, base::Time::Now(), + std::nullopt, std::nullopt, false)); + + std::vector response_events; + response_events.push_back(mojom::ConversationEntryEvent::NewCompletionEvent( + mojom::CompletionEvent::New(expected_response))); + expected_history.push_back(mojom::ConversationTurn::New( + std::nullopt, mojom::CharacterType::ASSISTANT, + mojom::ActionType::RESPONSE, expected_response, std::nullopt, + std::nullopt, std::move(response_events), base::Time::Now(), std::nullopt, + std::nullopt, false)); + // Should never ask for page content EXPECT_CALL(*associated_content_, GetContent).Times(0); NiceMock client(conversation_handler_.get()); EXPECT_CALL(client, OnAPIRequestInProgress(true)).Times(1); // Human, AI entries and content event for AI response. - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(3); + EXPECT_CALL(client, OnConversationHistoryUpdate( + TurnEq(mojom::ConversationTurnPtr().get()))) + .Times(1); + EXPECT_CALL(client, + OnConversationHistoryUpdate(TurnEq(expected_history[1].get()))) + .Times(2); // Fired from OnEngineCompletionComplete. EXPECT_CALL(client, OnAPIRequestInProgress(false)).Times(1); // Ensure everything is sanitized @@ -457,22 +504,6 @@ TEST_F(ConversationHandlerUnitTest, SubmitSelectedText) { EXPECT_TRUE(conversation_handler_->HasAnyHistory()); const auto& history = conversation_handler_->GetConversationHistory(); - std::vector expected_history; - - expected_history.push_back(mojom::ConversationTurn::New( - std::nullopt, mojom::CharacterType::HUMAN, - mojom::ActionType::SUMMARIZE_SELECTED_TEXT, expected_turn_text, - std::nullopt, selected_text, std::nullopt, base::Time::Now(), - std::nullopt, std::nullopt, false)); - - std::vector response_events; - response_events.push_back(mojom::ConversationEntryEvent::NewCompletionEvent( - mojom::CompletionEvent::New(expected_response))); - expected_history.push_back(mojom::ConversationTurn::New( - std::nullopt, mojom::CharacterType::ASSISTANT, - mojom::ActionType::RESPONSE, expected_response, std::nullopt, - std::nullopt, std::move(response_events), base::Time::Now(), std::nullopt, - std::nullopt, false)); ExpectConversationHistoryEquals(FROM_HERE, history, expected_history, false); } @@ -512,10 +543,31 @@ TEST_F(ConversationHandlerUnitTest, SubmitSelectedText_WithAssociatedContent) { EXPECT_EQ(site_info->url, GURL("https://www.brave.com/")); })); + std::vector expected_history; + expected_history.push_back(mojom::ConversationTurn::New( + std::nullopt, mojom::CharacterType::HUMAN, + mojom::ActionType::SUMMARIZE_SELECTED_TEXT, expected_turn_text, + std::nullopt, selected_text, std::nullopt, base::Time::Now(), + std::nullopt, std::nullopt, false)); + + std::vector response_events; + response_events.push_back(mojom::ConversationEntryEvent::NewCompletionEvent( + mojom::CompletionEvent::New(expected_response))); + expected_history.push_back(mojom::ConversationTurn::New( + std::nullopt, mojom::CharacterType::ASSISTANT, + mojom::ActionType::RESPONSE, expected_response, std::nullopt, + std::nullopt, std::move(response_events), base::Time::Now(), std::nullopt, + std::nullopt, false)); + NiceMock client(conversation_handler_.get()); EXPECT_CALL(client, OnAPIRequestInProgress(true)).Times(1); // Human and AI entries, and content event for AI response. - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(3); + EXPECT_CALL(client, OnConversationHistoryUpdate( + TurnEq(mojom::ConversationTurnPtr().get()))) + .Times(1); + EXPECT_CALL(client, + OnConversationHistoryUpdate(TurnEq(expected_history[1].get()))) + .Times(2); // Fired from OnEngineCompletionComplete. EXPECT_CALL(client, OnAPIRequestInProgress(false)).Times(1); // Ensure everything is sanitized @@ -546,24 +598,8 @@ TEST_F(ConversationHandlerUnitTest, SubmitSelectedText_WithAssociatedContent) { EXPECT_EQ(1u, questions.size()); EXPECT_EQ(questions[0], "Summarize this page"); - const auto& history2 = conversation_handler_->GetConversationHistory(); - std::vector expected_history2; - expected_history2.push_back(mojom::ConversationTurn::New( - std::nullopt, mojom::CharacterType::HUMAN, - mojom::ActionType::SUMMARIZE_SELECTED_TEXT, expected_turn_text, - std::nullopt, selected_text, std::nullopt, base::Time::Now(), - std::nullopt, std::nullopt, false)); - - std::vector response_events; - response_events.push_back(mojom::ConversationEntryEvent::NewCompletionEvent( - mojom::CompletionEvent::New(expected_response))); - expected_history2.push_back(mojom::ConversationTurn::New( - std::nullopt, mojom::CharacterType::ASSISTANT, - mojom::ActionType::RESPONSE, expected_response, std::nullopt, - std::nullopt, std::move(response_events), base::Time::Now(), std::nullopt, - std::nullopt, false)); - ExpectConversationHistoryEquals(FROM_HERE, history2, expected_history2, - false); + const auto& history = conversation_handler_->GetConversationHistory(); + ExpectConversationHistoryEquals(FROM_HERE, history, expected_history, false); } TEST_F(ConversationHandlerUnitTest, UpdateOrCreateLastAssistantEntry_Delta) { @@ -942,7 +978,9 @@ TEST_F(ConversationHandlerUnitTest, EXPECT_TRUE(conversation_handler_->IsAnyClientConnected()); // History update notification once for each entry - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(2); + EXPECT_CALL(client, OnConversationHistoryUpdate( + TurnEq(mojom::ConversationTurnPtr().get()))) + .Times(2); conversation_handler_->GetAssociatedContentInfo( base::BindLambdaForTesting([&](mojom::AssociatedContentPtr site_info, @@ -978,7 +1016,9 @@ TEST_F(ConversationHandlerUnitTest, EXPECT_FALSE(conversation_handler_->HasAnyHistory()); // Verify turning off content association clears the conversation history. - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(1); + EXPECT_CALL(client, OnConversationHistoryUpdate( + TurnEq(mojom::ConversationTurnPtr().get()))) + .Times(1); // Shouldn't ask for staged entries if user doesn't want to be associated // with content. This verifies that even with existing staged entries, // MaybeFetchOrClearContentStagedConversation will always early return. @@ -1004,7 +1044,9 @@ TEST_F(ConversationHandlerUnitTest, EXPECT_CALL(observer, OnConversationEntryAdded).Times(0); EXPECT_CALL(*associated_content_, GetStagedEntriesFromContent).Times(1); NiceMock client(conversation_handler_.get()); - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(4); + EXPECT_CALL(client, OnConversationHistoryUpdate( + TurnEq(mojom::ConversationTurnPtr().get()))) + .Times(4); EXPECT_TRUE(conversation_handler_->IsAnyClientConnected()); conversation_handler_->GetAssociatedContentInfo( base::BindLambdaForTesting([&](mojom::AssociatedContentPtr site_info, @@ -1068,7 +1110,18 @@ TEST_F(ConversationHandlerUnitTest, base::test::RunOnceCallback<5>(base::ok("")))); EXPECT_CALL(observer, OnConversationEntryAdded).Times(6); - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(3); + EXPECT_CALL(client, OnConversationHistoryUpdate( + TurnEq(mojom::ConversationTurnPtr().get()))) + .Times(1); + std::vector events3; + events3.push_back(mojom::ConversationEntryEvent::NewCompletionEvent( + mojom::CompletionEvent::New("new answer"))); + auto expected_turn = mojom::ConversationTurn::New( + std::nullopt, mojom::CharacterType::ASSISTANT, + mojom::ActionType::RESPONSE, "new answer", std::nullopt, std::nullopt, + std::move(events3), base::Time::Now(), std::nullopt, std::nullopt, false); + EXPECT_CALL(client, OnConversationHistoryUpdate(TurnEq(expected_turn.get()))) + .Times(2); conversation_handler_->SubmitHumanConversationEntry("query3", std::nullopt); @@ -1088,7 +1141,7 @@ TEST_F(ConversationHandlerUnitTest, EXPECT_CALL(*associated_content_, GetStagedEntriesFromContent).Times(1); NiceMock client(conversation_handler_.get()); // Should not notify of new history - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(0); + EXPECT_CALL(client, OnConversationHistoryUpdate(_)).Times(0); EXPECT_TRUE(conversation_handler_->IsAnyClientConnected()); task_environment_.RunUntilIdle(); @@ -1108,15 +1161,24 @@ TEST_F( // MaybeFetchOrClearContentStagedConversation should clear old staged entries // and fetch new ones. EXPECT_CALL(*associated_content_, GetStagedEntriesFromContent).Times(1); - // 4 from SetupHistory and 4 from adding - // new entries in OnGetStagedEntriesFromContent. - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(8); // Fill history with staged and non-staged entries. - SetupHistory({{"old query" /* text */, true /*from_brave_search_SERP */}, - {"old summary", "true"}, - {"normal query", false}, - {"normal response", false}}); + auto expected_history = + SetupHistory({{"old query" /* text */, true /*from_brave_search_SERP */}, + {"old summary", "true"}, + {"normal query", false}, + {"normal response", false}}); + // 4 from SetupHistory and 4 from adding + // new entries in OnGetStagedEntriesFromContent. + EXPECT_CALL(client, OnConversationHistoryUpdate( + TurnEq(mojom::ConversationTurnPtr().get()))) + .Times(6); + EXPECT_CALL(client, + OnConversationHistoryUpdate(TurnEq(expected_history[2].get()))) + .Times(1); + EXPECT_CALL(client, + OnConversationHistoryUpdate(TurnEq(expected_history[3].get()))) + .Times(1); // Setting mock return values for GetStagedEntriesFromContent. SetAssociatedContentStagedEntries(/*empty=*/false, /*multi=*/true); @@ -1165,12 +1227,21 @@ TEST_F(ConversationHandlerUnitTest, OnGetStagedEntriesFromContent) { NiceMock client(conversation_handler_.get()); ASSERT_TRUE(conversation_handler_->IsAnyClientConnected()); - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(8); // Fill history with staged and non-staged entries. - SetupHistory({{"q1" /* text */, true /*from_brave_search_SERP */}, - {"s1", "true"}, - {"q2", false}, - {"r1", false}}); + auto expected_history = + SetupHistory({{"q1" /* text */, true /*from_brave_search_SERP */}, + {"s1", "true"}, + {"q2", false}, + {"r1", false}}); + EXPECT_CALL(client, OnConversationHistoryUpdate( + TurnEq(mojom::ConversationTurnPtr().get()))) + .Times(6); + EXPECT_CALL(client, + OnConversationHistoryUpdate(TurnEq(expected_history[2].get()))) + .Times(1); + EXPECT_CALL(client, + OnConversationHistoryUpdate(TurnEq(expected_history[3].get()))) + .Times(1); std::vector entries = {{"query", "summary"}, {"query2", "summary2"}}; @@ -1231,7 +1302,7 @@ TEST_F(ConversationHandlerUnitTest, // Client connecting will trigger content staging NiceMock client(conversation_handler_.get()); - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(0); + EXPECT_CALL(client, OnConversationHistoryUpdate(_)).Times(0); EXPECT_TRUE(conversation_handler_->IsAnyClientConnected()); task_environment_.RunUntilIdle(); @@ -1241,7 +1312,7 @@ TEST_F(ConversationHandlerUnitTest, EXPECT_TRUE(conversation_handler_->GetConversationHistory().empty()); } -TEST_F(ConversationHandlerUnitTest, UploadImage) { +TEST_F(ConversationHandlerUnitTest, UploadFile) { conversation_handler_->SetShouldSendPageContents(false); constexpr char kTestPrompt[] = "What is this?"; MockEngineConsumer* engine = static_cast( @@ -1252,7 +1323,7 @@ TEST_F(ConversationHandlerUnitTest, UploadImage) { base::ok("This is a lion."))); ASSERT_FALSE(conversation_handler_->GetCurrentModel().vision_support); - // No uploaded images + // No uploaded files base::RunLoop loop; EXPECT_CALL(client, OnModelDataChanged).Times(0); EXPECT_CALL(client, OnAPIRequestInProgress(true)).Times(1); @@ -1262,48 +1333,58 @@ TEST_F(ConversationHandlerUnitTest, UploadImage) { std::nullopt); loop.Run(); EXPECT_FALSE( - conversation_handler_->GetConversationHistory().back()->uploaded_images); + conversation_handler_->GetConversationHistory().back()->uploaded_files); testing::Mock::VerifyAndClearExpectations(&client); - // Empty uploaded images + // Empty uploaded files base::RunLoop loop2; EXPECT_CALL(client, OnModelDataChanged).Times(0); EXPECT_CALL(client, OnAPIRequestInProgress(true)).Times(1); EXPECT_CALL(client, OnAPIRequestInProgress(false)) .WillOnce(testing::InvokeWithoutArgs(&loop2, &base::RunLoop::Quit)); conversation_handler_->SubmitHumanConversationEntry( - kTestPrompt, std::vector()); + kTestPrompt, std::vector()); loop2.Run(); EXPECT_FALSE( - conversation_handler_->GetConversationHistory().back()->uploaded_images); + conversation_handler_->GetConversationHistory().back()->uploaded_files); testing::Mock::VerifyAndClearExpectations(&client); - auto uploaded_images = CreateSampleUploadedImages(3); + auto uploaded_files = CreateSampleUploadedFiles(3); // There are uploaded images. // Note that this will need to be put at the end of this test suite // because currently there is no perfect timing to call // SetEngineForTesting() after auto model switch. base::RunLoop loop3; - EXPECT_CALL(client, OnModelDataChanged) - .WillOnce(base::test::RunClosure(base::BindLambdaForTesting([&]() { - // verify auto switched to vision support model - EXPECT_TRUE(conversation_handler_->GetCurrentModel().vision_support); - loop3.Quit(); - }))); + if (std::ranges::any_of( + uploaded_files, [](const mojom::UploadedFilePtr& file) { + return file->type == mojom::UploadedFileType::kImage || + file->type == mojom::UploadedFileType::kScreenshot; + })) { + EXPECT_CALL(client, OnModelDataChanged) + .WillOnce(base::test::RunClosure(base::BindLambdaForTesting([&]() { + // verify auto switched to vision support model + EXPECT_TRUE(conversation_handler_->GetCurrentModel().vision_support); + loop3.Quit(); + }))); + } else { + EXPECT_CALL(client, OnAPIRequestInProgress(false)) + .WillOnce(testing::InvokeWithoutArgs(&loop3, &base::RunLoop::Quit)); + } conversation_handler_->SubmitHumanConversationEntry(kTestPrompt, - Clone(uploaded_images)); + Clone(uploaded_files)); loop3.Run(); testing::Mock::VerifyAndClearExpectations(&client); // verify image in history auto& last_entry = conversation_handler_->GetConversationHistory().back(); - EXPECT_TRUE(last_entry->uploaded_images); - const auto& images = last_entry->uploaded_images.value(); - for (size_t i = 0; i < images.size(); ++i) { - EXPECT_EQ(images[i]->filename, uploaded_images[i]->filename); - EXPECT_EQ(images[i]->filesize, uploaded_images[i]->filesize); - EXPECT_EQ(images[i]->image_data, uploaded_images[i]->image_data); + EXPECT_TRUE(last_entry->uploaded_files); + const auto& files = last_entry->uploaded_files.value(); + for (size_t i = 0; i < files.size(); ++i) { + EXPECT_EQ(files[i]->filename, uploaded_files[i]->filename); + EXPECT_EQ(files[i]->filesize, uploaded_files[i]->filesize); + EXPECT_EQ(files[i]->data, uploaded_files[i]->data); + EXPECT_EQ(files[i]->type, uploaded_files[i]->type); } } @@ -1315,7 +1396,7 @@ TEST_F(ConversationHandlerUnitTest_NoAssociatedContent, bool should_send_page_contents) { EXPECT_FALSE(site_info); })); // Client connecting would trigger content staging NiceMock client(conversation_handler_.get()); - EXPECT_CALL(client, OnConversationHistoryUpdate()).Times(0); + EXPECT_CALL(client, OnConversationHistoryUpdate(_)).Times(0); EXPECT_TRUE(conversation_handler_->IsAnyClientConnected()); task_environment_.RunUntilIdle(); diff --git a/components/ai_chat/core/browser/engine/conversation_api_client.cc b/components/ai_chat/core/browser/engine/conversation_api_client.cc index 7107b1632fa..9d3fc21b98e 100644 --- a/components/ai_chat/core/browser/engine/conversation_api_client.cc +++ b/components/ai_chat/core/browser/engine/conversation_api_client.cc @@ -114,7 +114,8 @@ base::Value::List ConversationEventsToList( {ConversationEventType::GetSuggestedAndDedupeTopicsForFocusTabs, "suggestAndDedupeFocusTopics"}, {ConversationEventType::GetFocusTabsForTopic, "classifyTabs"}, - {ConversationEventType::UploadImage, "uploadImage"}}); + {ConversationEventType::UploadImage, "uploadImage"}, + {ConversationEventType::PageScreenshot, "pageScreenshot"}}); base::Value::List events; for (const auto& event : conversation) { diff --git a/components/ai_chat/core/browser/engine/conversation_api_client.h b/components/ai_chat/core/browser/engine/conversation_api_client.h index b4e94eaebd8..f2523ea8be4 100644 --- a/components/ai_chat/core/browser/engine/conversation_api_client.h +++ b/components/ai_chat/core/browser/engine/conversation_api_client.h @@ -67,6 +67,7 @@ class ConversationAPIClient { DedupeTopics, GetSuggestedAndDedupeTopicsForFocusTabs, GetFocusTabsForTopic, + PageScreenshot, // TODO(petemill): // - Search in-progress? // - Sources? diff --git a/components/ai_chat/core/browser/engine/engine_consumer_conversation_api.cc b/components/ai_chat/core/browser/engine/engine_consumer_conversation_api.cc index 4377b94741f..5ede27084a5 100644 --- a/components/ai_chat/core/browser/engine/engine_consumer_conversation_api.cc +++ b/components/ai_chat/core/browser/engine/engine_consumer_conversation_api.cc @@ -196,14 +196,26 @@ void EngineConsumerConversationAPI::GenerateAssistantResponse( } // history for (const auto& message : conversation_history) { - if (message->uploaded_images) { - std::vector images; - for (const auto& uploaded_image : message->uploaded_images.value()) { - images.emplace_back(GetImageDataURL(uploaded_image->image_data)); + if (message->uploaded_files) { + std::vector uploaded_images; + std::vector screenshot_images; + for (const auto& uploaded_file : message->uploaded_files.value()) { + if (uploaded_file->type == mojom::UploadedFileType::kScreenshot) { + screenshot_images.emplace_back(GetImageDataURL(uploaded_file->data)); + } else if (uploaded_file->type == mojom::UploadedFileType::kImage) { + uploaded_images.emplace_back(GetImageDataURL(uploaded_file->data)); + } + } + if (!uploaded_images.empty()) { + conversation.push_back({mojom::CharacterType::HUMAN, + ConversationEventType::UploadImage, + std::move(uploaded_images)}); + } + if (!screenshot_images.empty()) { + conversation.push_back({mojom::CharacterType::HUMAN, + ConversationEventType::PageScreenshot, + std::move(screenshot_images)}); } - conversation.push_back({mojom::CharacterType::HUMAN, - ConversationEventType::UploadImage, - std::move(images)}); } if (message->selected_text.has_value() && !message->selected_text->empty()) { diff --git a/components/ai_chat/core/browser/engine/engine_consumer_conversation_api_unittest.cc b/components/ai_chat/core/browser/engine/engine_consumer_conversation_api_unittest.cc index 5cacb061978..c7e2a3262c6 100644 --- a/components/ai_chat/core/browser/engine/engine_consumer_conversation_api_unittest.cc +++ b/components/ai_chat/core/browser/engine/engine_consumer_conversation_api_unittest.cc @@ -29,6 +29,7 @@ #include "base/values.h" #include "brave/components/ai_chat/core/browser/engine/conversation_api_client.h" #include "brave/components/ai_chat/core/browser/engine/engine_consumer.h" +#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-shared.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h" #include "brave/components/ai_chat/core/common/test_utils.h" #include "services/network/public/cpp/shared_url_loader_factory.h" @@ -473,9 +474,17 @@ TEST_F(EngineConsumerConversationAPIUnitTest, GenerateEvents_SummarizePage) { } TEST_F(EngineConsumerConversationAPIUnitTest, GenerateEvents_UploadImage) { - auto uploaded_images = CreateSampleUploadedImages(3); - constexpr char kTestPrompt[] = "Tell the user what is in the image?"; - constexpr char kAssistantResponse[] = "It's a lion!"; + auto uploaded_images = + CreateSampleUploadedFiles(3, mojom::UploadedFileType::kImage); + auto screenshot_images = + CreateSampleUploadedFiles(3, mojom::UploadedFileType::kScreenshot); + uploaded_images.insert(uploaded_images.end(), + std::make_move_iterator(screenshot_images.begin()), + std::make_move_iterator(screenshot_images.end())); + constexpr char kTestPrompt[] = "Tell the user what these images are?"; + constexpr char kAssistantResponse[] = + "There are images of a lion, a dragon and a stag. And screenshots appear " + "to be telling the story of Game of Thrones"; auto* mock_api_client = GetMockConversationAPIClient(); base::RunLoop run_loop; EXPECT_CALL(*mock_api_client, PerformRequest(_, _, _, _)) @@ -483,24 +492,32 @@ TEST_F(EngineConsumerConversationAPIUnitTest, GenerateEvents_UploadImage) { const std::string& selected_language, EngineConsumer::GenerationDataCallback data_callback, EngineConsumer::GenerationCompletedCallback callback) { - // Only support one image for now. - ASSERT_EQ(conversation.size(), 2u); + ASSERT_EQ(conversation.size(), 3u); EXPECT_EQ(conversation[0].role, mojom::CharacterType::HUMAN); - EXPECT_EQ( - conversation[0].content[0], - base::StrCat({"data:image/png;base64,", - base::Base64Encode(uploaded_images[0]->image_data)})); + for (size_t i = 0; i < 3; ++i) { + EXPECT_EQ( + conversation[0].content[i], + base::StrCat({"data:image/png;base64,", + base::Base64Encode(uploaded_images[i]->data)})); + } EXPECT_EQ(conversation[0].type, ConversationAPIClient::UploadImage); - EXPECT_EQ(conversation[1].role, mojom::CharacterType::HUMAN); - EXPECT_EQ(conversation[1].content[0], kTestPrompt); - EXPECT_EQ(conversation[1].type, ConversationAPIClient::ChatMessage); + for (size_t i = 3; i < uploaded_images.size(); ++i) { + EXPECT_EQ( + conversation[1].content[i - 3], + base::StrCat({"data:image/png;base64,", + base::Base64Encode(uploaded_images[i]->data)})); + } + EXPECT_EQ(conversation[1].type, ConversationAPIClient::PageScreenshot); + EXPECT_EQ(conversation[2].role, mojom::CharacterType::HUMAN); + EXPECT_EQ(conversation[2].content[0], kTestPrompt); + EXPECT_EQ(conversation[2].type, ConversationAPIClient::ChatMessage); std::move(callback).Run(kAssistantResponse); }); std::vector history; history.push_back(mojom::ConversationTurn::New( std::nullopt, mojom::CharacterType::HUMAN, mojom::ActionType::UNSPECIFIED, - "What is this image?", kTestPrompt, std::nullopt, std::nullopt, + "What are these images?", kTestPrompt, std::nullopt, std::nullopt, base::Time::Now(), std::nullopt, Clone(uploaded_images), false)); base::test::TestFuture future; diff --git a/components/ai_chat/core/browser/engine/engine_consumer_oai.cc b/components/ai_chat/core/browser/engine/engine_consumer_oai.cc index 4dbd068fb16..ca6f8924deb 100644 --- a/components/ai_chat/core/browser/engine/engine_consumer_oai.cc +++ b/components/ai_chat/core/browser/engine/engine_consumer_oai.cc @@ -26,6 +26,7 @@ #include "base/values.h" #include "brave/components/ai_chat/core/browser/engine/engine_consumer.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-forward.h" +#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-shared.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h" #include "components/grit/brave_components_strings.h" #include "services/network/public/cpp/shared_url_loader_factory.h" @@ -92,31 +93,45 @@ base::Value::List BuildMessages( } for (const mojom::ConversationTurnPtr& turn : conversation_history) { - if (turn->uploaded_images) { - base::Value::Dict message; - message.Set("role", "user"); - base::Value::List content; - base::Value::Dict user_message; - user_message.Set("type", "text"); - user_message.Set("text", "These images are uploaded by the users"); - content.Append(std::move(user_message)); - size_t counter = 0; - // Only send the first uploaded_image becasue llama-vision seems to take - // the last one if there are multiple uploaded_images - for (const auto& uploaded_image : turn->uploaded_images.value()) { - if (counter++ > 0) { - break; + if (turn->uploaded_files) { + base::Value::List content_uploaded_images; + base::Value::List content_screenshots; + content_uploaded_images.Append( + base::Value::Dict() + .Set("type", "text") + .Set("text", "These images are uploaded by the user")); + content_screenshots.Append( + base::Value::Dict() + .Set("type", "text") + .Set("text", "These images are screenshots")); + for (const auto& uploaded_file : turn->uploaded_files.value()) { + if (uploaded_file->type != mojom::UploadedFileType::kImage && + uploaded_file->type != mojom::UploadedFileType::kScreenshot) { + continue; } base::Value::Dict image; image.Set("type", "image_url"); base::Value::Dict image_url_dict; image_url_dict.Set( - "url", EngineConsumer::GetImageDataURL(uploaded_image->image_data)); + "url", EngineConsumer::GetImageDataURL(uploaded_file->data)); image.Set("image_url", std::move(image_url_dict)); - content.Append(std::move(image)); + if (uploaded_file->type == mojom::UploadedFileType::kImage) { + content_uploaded_images.Append(std::move(image)); + } else { + content_screenshots.Append(std::move(image)); + } + } + if (content_uploaded_images.size() > 1) { + messages.Append( + base::Value::Dict() + .Set("role", "user") + .Set("content", std::move(content_uploaded_images))); + } + if (content_screenshots.size() > 1) { + messages.Append(base::Value::Dict() + .Set("role", "user") + .Set("content", std::move(content_screenshots))); } - message.Set("content", std::move(content)); - messages.Append(std::move(message)); } base::Value::Dict message; message.Set("role", turn->character_type == CharacterType::HUMAN diff --git a/components/ai_chat/core/browser/engine/engine_consumer_oai_unittest.cc b/components/ai_chat/core/browser/engine/engine_consumer_oai_unittest.cc index 848cf1d375c..03ce3e9c7ea 100644 --- a/components/ai_chat/core/browser/engine/engine_consumer_oai_unittest.cc +++ b/components/ai_chat/core/browser/engine/engine_consumer_oai_unittest.cc @@ -30,6 +30,7 @@ #include "brave/components/ai_chat/core/browser/engine/engine_consumer.h" #include "brave/components/ai_chat/core/browser/engine/test_utils.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-forward.h" +#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-shared.h" #include "brave/components/ai_chat/core/common/test_utils.h" #include "components/grit/brave_components_strings.h" #include "services/network/public/cpp/shared_url_loader_factory.h" @@ -449,9 +450,17 @@ TEST_F(EngineConsumerOAIUnitTest, GenerateAssistantResponseEarlyReturn) { TEST_F(EngineConsumerOAIUnitTest, GenerateAssistantResponseUploadImage) { EngineConsumer::ConversationHistory history; auto* client = GetClient(); - auto uploaded_images = CreateSampleUploadedImages(3); - constexpr char kTestPrompt[] = "Tell the user what is in the image?"; - constexpr char kAssistantResponse[] = "It's a lion!"; + auto uploaded_images = + CreateSampleUploadedFiles(3, mojom::UploadedFileType::kImage); + auto screenshot_images = + CreateSampleUploadedFiles(3, mojom::UploadedFileType::kScreenshot); + uploaded_images.insert(uploaded_images.end(), + std::make_move_iterator(screenshot_images.begin()), + std::make_move_iterator(screenshot_images.end())); + constexpr char kTestPrompt[] = "Tell the user what these images are?"; + constexpr char kAssistantResponse[] = + "There are images of a lion, a dragon and a stag. And screenshots appear " + "to be telling the story of Game of Thrones"; EXPECT_CALL(*client, PerformRequest(_, _, _, _)) .WillOnce( [kTestPrompt, kAssistantResponse, &uploaded_images]( @@ -461,27 +470,50 @@ TEST_F(EngineConsumerOAIUnitTest, GenerateAssistantResponseUploadImage) { EXPECT_EQ(*messages[0].GetDict().Find("role"), "system"); constexpr char kJsonTemplate[] = R"({ - "content": [ { - "text": "These images are uploaded by the users", + "content": [{ + "text": "$1", "type": "text" }, { "image_url": { - "url": "data:image/png;base64,$1" + "url": "data:image/png;base64,$2" }, "type": "image_url" - } ], + }, { + "image_url": { + "url": "data:image/png;base64,$3" + }, + "type": "image_url" + }, { + "image_url": { + "url": "data:image/png;base64,$4" + }, + "type": "image_url" + }], "role": "user" } )"; - const std::string json_str = base::ReplaceStringPlaceholders( + ASSERT_EQ(uploaded_images.size(), 6u); + const std::string image_json_str = base::ReplaceStringPlaceholders( kJsonTemplate, - {base::Base64Encode(uploaded_images[0]->image_data)}, nullptr); - auto expected_dict = ParseJsonDict(json_str); + {"These images are uploaded by the user", + base::Base64Encode(uploaded_images[0]->data), + base::Base64Encode(uploaded_images[1]->data), + base::Base64Encode(uploaded_images[2]->data)}, + nullptr); + EXPECT_EQ(messages[1].GetDict(), ParseJsonDict(image_json_str)); + const std::string screenshot_json_str = + base::ReplaceStringPlaceholders( + kJsonTemplate, + {"These images are screenshots", + base::Base64Encode(uploaded_images[3]->data), + base::Base64Encode(uploaded_images[4]->data), + base::Base64Encode(uploaded_images[5]->data)}, + nullptr); + EXPECT_EQ(messages[2].GetDict(), + ParseJsonDict(screenshot_json_str)); - EXPECT_EQ(messages[1].GetDict(), expected_dict); - - EXPECT_EQ(*messages[2].GetDict().Find("role"), "user"); - EXPECT_EQ(*messages[2].GetDict().Find("content"), kTestPrompt); + EXPECT_EQ(*messages[3].GetDict().Find("role"), "user"); + EXPECT_EQ(*messages[3].GetDict().Find("content"), kTestPrompt); std::move(completed_callback) .Run(EngineConsumer::GenerationResult(kAssistantResponse)); @@ -489,7 +521,7 @@ TEST_F(EngineConsumerOAIUnitTest, GenerateAssistantResponseUploadImage) { history.push_back(mojom::ConversationTurn::New( std::nullopt, mojom::CharacterType::HUMAN, mojom::ActionType::UNSPECIFIED, - "What is this image?", kTestPrompt, std::nullopt, std::nullopt, + "What are these images?", kTestPrompt, std::nullopt, std::nullopt, base::Time::Now(), std::nullopt, Clone(uploaded_images), false)); base::test::TestFuture future; engine_->GenerateAssistantResponse(false, "", history, "", base::DoNothing(), diff --git a/components/ai_chat/core/browser/test_utils.cc b/components/ai_chat/core/browser/test_utils.cc index 815f9a8c8e5..e2a13fc77eb 100644 --- a/components/ai_chat/core/browser/test_utils.cc +++ b/components/ai_chat/core/browser/test_utils.cc @@ -170,18 +170,19 @@ void ExpectConversationEntryEquals(base::Location location, } } - // compare uploaded images - EXPECT_EQ(a->uploaded_images.has_value(), b->uploaded_images.has_value()); - if (a->uploaded_images.has_value()) { - EXPECT_EQ(a->uploaded_images->size(), b->uploaded_images->size()); - for (size_t i = 0; i < a->uploaded_images->size(); ++i) { + // compare uploaded files + EXPECT_EQ(a->uploaded_files.has_value(), b->uploaded_files.has_value()); + if (a->uploaded_files.has_value()) { + EXPECT_EQ(a->uploaded_files->size(), b->uploaded_files->size()); + for (size_t i = 0; i < a->uploaded_files->size(); ++i) { SCOPED_TRACE(testing::Message() - << "Comparing uplodaed images at index " << i); - const auto& uploaded_image_a = a->uploaded_images->at(i); - const auto& uploaded_image_b = b->uploaded_images->at(i); - EXPECT_EQ(uploaded_image_a->filename, uploaded_image_b->filename); - EXPECT_EQ(uploaded_image_a->filesize, uploaded_image_b->filesize); - EXPECT_EQ(uploaded_image_a->image_data, uploaded_image_b->image_data); + << "Comparing uplodaed files at index " << i); + const auto& uploaded_file_a = a->uploaded_files->at(i); + const auto& uploaded_file_b = b->uploaded_files->at(i); + EXPECT_EQ(uploaded_file_a->filename, uploaded_file_b->filename); + EXPECT_EQ(uploaded_file_a->filesize, uploaded_file_b->filesize); + EXPECT_EQ(uploaded_file_a->data, uploaded_file_b->data); + EXPECT_EQ(uploaded_file_a->type, uploaded_file_b->type); } } @@ -214,15 +215,14 @@ mojom::Conversation* GetConversation( std::vector CreateSampleChatHistory( size_t num_query_pairs, int32_t future_hours, - size_t num_uploaded_images_per_query) { + size_t num_uploaded_files_per_query) { std::vector history; base::Time now = base::Time::Now(); for (size_t i = 0; i < num_query_pairs; i++) { // query - std::optional> uploaded_images; - if (num_uploaded_images_per_query) { - uploaded_images = - CreateSampleUploadedImages(num_uploaded_images_per_query); + std::optional> uploaded_files; + if (num_uploaded_files_per_query) { + uploaded_files = CreateSampleUploadedFiles(num_uploaded_files_per_query); } history.push_back(mojom::ConversationTurn::New( base::Uuid::GenerateRandomV4().AsLowercaseString(), @@ -230,7 +230,7 @@ std::vector CreateSampleChatHistory( base::StrCat({"query", base::NumberToString(i)}), std::nullopt /* prompt */, std::nullopt, std::nullopt, now + base::Seconds(i * 60) + base::Hours(future_hours), std::nullopt, - std::move(uploaded_images), false)); + std::move(uploaded_files), false)); // response std::vector events; events.emplace_back(mojom::ConversationEntryEvent::NewCompletionEvent( diff --git a/components/ai_chat/core/browser/test_utils.h b/components/ai_chat/core/browser/test_utils.h index f2bab4adf9c..473fc376c0d 100644 --- a/components/ai_chat/core/browser/test_utils.h +++ b/components/ai_chat/core/browser/test_utils.h @@ -41,7 +41,7 @@ mojom::Conversation* GetConversation( std::vector CreateSampleChatHistory( size_t num_query_pairs, int32_t future_hours = 0, - size_t num_uploaded_images_per_query = 0); + size_t num_uploaded_files_per_query = 0); std::vector CloneHistory( std::vector& history); diff --git a/components/ai_chat/core/browser/utils.cc b/components/ai_chat/core/browser/utils.cc index 074f334732e..1c88356773a 100644 --- a/components/ai_chat/core/browser/utils.cc +++ b/components/ai_chat/core/browser/utils.cc @@ -25,6 +25,8 @@ #include "components/prefs/pref_service.h" #include "mojo/public/cpp/bindings/struct_ptr.h" #include "third_party/re2/src/re2/re2.h" +#include "third_party/skia/include/core/SkCanvas.h" +#include "third_party/skia/include/core/SkImage.h" #include "ui/base/l10n/l10n_util.h" #include "url/gurl.h" #include "url/url_constants.h" @@ -224,4 +226,47 @@ EngineConsumer::GenerationDataCallback BindParseRewriteReceivedData( std::move(callback)); } +SkBitmap ScaleDownBitmap(const SkBitmap& bitmap) { + constexpr int kTargetWidth = 1024; + constexpr int kTargetHeight = 768; + + // Don't need to scale if dimensions are already smaller than target + // dimensions + if (bitmap.width() <= kTargetWidth && bitmap.height() <= kTargetHeight) { + return bitmap; + } + + SkBitmap scaled_bitmap; + scaled_bitmap.allocN32Pixels(kTargetWidth, kTargetHeight); + + SkCanvas canvas(scaled_bitmap); + canvas.clear(SK_ColorTRANSPARENT); + + // Use high-quality scaling options + SkSamplingOptions sampling_options(SkFilterMode::kLinear, + SkMipmapMode::kLinear); + + // Maintain aspect ratio while fitting within target dimensions + float src_aspect = static_cast(bitmap.width()) / bitmap.height(); + float dst_aspect = static_cast(kTargetWidth) / kTargetHeight; + + SkRect dst_rect; + if (src_aspect > dst_aspect) { + // Source is wider - fit to width + float scaled_height = kTargetWidth / src_aspect; + float y_offset = (kTargetHeight - scaled_height) / 2; + dst_rect = SkRect::MakeXYWH(0, y_offset, kTargetWidth, scaled_height); + } else { + // Source is taller - fit to height + float scaled_width = kTargetHeight * src_aspect; + float x_offset = (kTargetWidth - scaled_width) / 2; + dst_rect = SkRect::MakeXYWH(x_offset, 0, scaled_width, kTargetHeight); + } + + // Draw scaled bitmap with high-quality sampling + canvas.drawImageRect(bitmap.asImage(), dst_rect, sampling_options); + + return scaled_bitmap; +} + } // namespace ai_chat diff --git a/components/ai_chat/core/browser/utils.h b/components/ai_chat/core/browser/utils.h index aa3e8de5c93..a71df5d7033 100644 --- a/components/ai_chat/core/browser/utils.h +++ b/components/ai_chat/core/browser/utils.h @@ -43,6 +43,10 @@ const std::string& GetActionTypeQuestion(mojom::ActionType action_type); EngineConsumer::GenerationDataCallback BindParseRewriteReceivedData( ConversationHandler::GeneratedTextCallback callback); +// Only scales down to target dimension when input bitmap is larger than +// 1024x768 +SkBitmap ScaleDownBitmap(const SkBitmap& bitmap); + } // namespace ai_chat #endif // BRAVE_COMPONENTS_AI_CHAT_CORE_BROWSER_UTILS_H_ diff --git a/components/ai_chat/core/browser/utils_unittest.cc b/components/ai_chat/core/browser/utils_unittest.cc index 4533cc9003d..63b727e9c48 100644 --- a/components/ai_chat/core/browser/utils_unittest.cc +++ b/components/ai_chat/core/browser/utils_unittest.cc @@ -5,10 +5,12 @@ #include "brave/components/ai_chat/core/browser/utils.h" -#include #include +#include +#include #include "testing/gtest/include/gtest/gtest.h" +#include "ui/gfx/image/image_unittest_util.h" #include "url/gurl.h" namespace ai_chat { @@ -28,4 +30,25 @@ TEST(AIChatUtilsUnitTest, IsBraveSearchSERP) { EXPECT_FALSE(IsBraveSearchSERP(GURL("https://brave.com/search?q=foo"))); } +TEST(AIChatUtilsUnitTest, ScaleDownBitmap) { + const std::vector> large_test_dimensions = { + {2560, 1440}, {1024, 1440}, {2560, 768}}; + for (auto& [width, height] : large_test_dimensions) { + SCOPED_TRACE(testing::Message() << width << "x" << height); + const auto bitmap = gfx::test::CreateBitmap(width, height); + const auto scaled_bitmap = ScaleDownBitmap(bitmap); + EXPECT_EQ(scaled_bitmap.width(), 1024); + EXPECT_EQ(scaled_bitmap.height(), 768); + } + + const std::vector> no_change_test_dimensions = { + {1024, 768}, {1024, 720}, {960, 768}, {960, 720}}; + for (auto& [width, height] : no_change_test_dimensions) { + SCOPED_TRACE(testing::Message() << width << "x" << height); + const auto bitmap = gfx::test::CreateBitmap(width, height); + const auto scaled_bitmap = ScaleDownBitmap(bitmap); + EXPECT_TRUE(gfx::test::AreBitmapsEqual(bitmap, scaled_bitmap)); + } +} + } // namespace ai_chat diff --git a/components/ai_chat/core/common/mojom/ai_chat.mojom b/components/ai_chat/core/common/mojom/ai_chat.mojom index 4cf5d8cf62b..1a074a13cd4 100644 --- a/components/ai_chat/core/common/mojom/ai_chat.mojom +++ b/components/ai_chat/core/common/mojom/ai_chat.mojom @@ -191,10 +191,16 @@ struct ConversationTitleEvent { string title; }; -struct UploadedImage { +enum UploadedFileType { + kImage = 0, + kScreenshot, +}; + +struct UploadedFile { string filename; uint32 filesize; - array image_data; + array data; + UploadedFileType type; }; // Events that occur during a conversation turn (only assistant for now) @@ -250,9 +256,8 @@ struct ConversationTurn { // engine or displaying the most recent text to users. array? edits; - // Due to server limitation, we only support one image in one turn for the - // entire conversation. Each turns will share the same image. - array? uploaded_images; + // User uploaded files, currently including image and screenshot + array? uploaded_files; // Whether the turn was generated from Brave Search SERP. bool from_brave_search_SERP = false; @@ -404,7 +409,7 @@ interface AIChatUIHandler { // use_media_capture is used on Android only currently. It could // be used on desktop as well in the future when a camera support is added. UploadImage(bool use_media_capture) - => (array? uploaded_images); + => (array? uploaded_images); // This might be a no-op if the UI isn't closeable CloseUI(); @@ -494,7 +499,7 @@ interface ConversationHandler { GenerateQuestions(); SubmitHumanConversationEntry( - string input, array? uploaded_images); + string input, array? uploaded_files); SubmitHumanConversationEntryWithAction(string input, ActionType action_type); ModifyConversation(uint32 turn_index, string new_text); SubmitSummarizationRequest(); @@ -530,6 +535,8 @@ interface ConversationHandler { string category, string feedback, string rating_id, bool send_hostname) => (bool is_success); + + GetScreenshots() => (array? screenshots); }; // Browser-side handler for a Conversation's UI responsible for displaying @@ -548,18 +555,12 @@ interface UntrustedConversationHandler { // Untrusted-UI-side handler for a Conversation, responsible for displaying // content generated by the AI engine. interface UntrustedConversationUI { - // TODO(petemill): Provide single entry that's been updated so that we don't - // need to fetch (and clone) all conversation entries each time text is added - // to the most recent entry. - OnConversationHistoryUpdate(); + OnConversationHistoryUpdate(ConversationTurn? entry); OnEntriesUIStateChanged(ConversationEntriesState state); }; interface ConversationUI { - // TODO(petemill): Provide single entry that's been updated so that we don't - // need to fetch (and clone) all conversation entries each time text is added - // to the most recent entry. - OnConversationHistoryUpdate(); + OnConversationHistoryUpdate(ConversationTurn? entry); OnAPIRequestInProgress(bool is_request_in_progress); OnAPIResponseError(APIError error); // Usually the model is changed from the UI client, but occassionally diff --git a/components/ai_chat/core/common/test_utils.cc b/components/ai_chat/core/common/test_utils.cc index 231dc84edd2..5af04121abd 100644 --- a/components/ai_chat/core/common/test_utils.cc +++ b/components/ai_chat/core/common/test_utils.cc @@ -6,20 +6,32 @@ #include "brave/components/ai_chat/core/common/test_utils.h" #include "base/rand_util.h" +#include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-shared.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h" #include "crypto/random.h" namespace ai_chat { -std::vector CreateSampleUploadedImages(size_t number) { - std::vector uploaded_images; +std::vector CreateSampleUploadedFiles( + size_t number, + std::optional type) { + std::vector uploaded_files; for (size_t i = 0; i < number; ++i) { - std::vector image_data(base::RandGenerator(64)); - crypto::RandBytes(image_data); - uploaded_images.emplace_back(mojom::UploadedImage::New( - "filename" + base::NumberToString(i), sizeof(image_data), image_data)); + std::vector file_data(base::RandGenerator(64)); + crypto::RandBytes(file_data); + mojom::UploadedFileType type_to_use; + if (!type) { + type_to_use = static_cast( + base::RandInt(static_cast(mojom::UploadedFileType::kMinValue), + static_cast(mojom::UploadedFileType::kMaxValue))); + } else { + type_to_use = type.value(); + } + uploaded_files.emplace_back( + mojom::UploadedFile::New("filename" + base::NumberToString(i), + sizeof(file_data), file_data, type_to_use)); } - return uploaded_images; + return uploaded_files; } } // namespace ai_chat diff --git a/components/ai_chat/core/common/test_utils.h b/components/ai_chat/core/common/test_utils.h index 286634d4312..af64ec8258b 100644 --- a/components/ai_chat/core/common/test_utils.h +++ b/components/ai_chat/core/common/test_utils.h @@ -6,13 +6,16 @@ #ifndef BRAVE_COMPONENTS_AI_CHAT_CORE_COMMON_TEST_UTILS_H_ #define BRAVE_COMPONENTS_AI_CHAT_CORE_COMMON_TEST_UTILS_H_ +#include #include #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-forward.h" namespace ai_chat { -std::vector CreateSampleUploadedImages(size_t number); +std::vector CreateSampleUploadedFiles( + size_t number, + std::optional type = std::nullopt); } // namespace ai_chat diff --git a/components/ai_chat/resources/common/conversation_history_utils.ts b/components/ai_chat/resources/common/conversation_history_utils.ts new file mode 100644 index 00000000000..729f1fc388f --- /dev/null +++ b/components/ai_chat/resources/common/conversation_history_utils.ts @@ -0,0 +1,51 @@ +// Copyright (c) 2025 The Brave Authors. All rights reserved. +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this file, +// You can obtain one at https://mozilla.org/MPL/2.0/. + +import * as Mojom from './mojom' + +/** + * Updates the conversation history by either merging a new entry with an + * existing one or appending it if it doesn't exist. + * + * @param currentHistory - The current conversation history + * @param newEntry - The new entry to be merged or appended + * @returns Updated conversation history + */ +export function updateConversationHistory( + currentHistory: Mojom.ConversationTurn[], + newEntry: Mojom.ConversationTurn +): Mojom.ConversationTurn[] { + // Check if an entry with the same UUID already exists + const existingEntryIndex = currentHistory.findIndex( + (existingEntry) => existingEntry.uuid === newEntry.uuid + ) + + if (existingEntryIndex !== -1) { + // If entry exists, merge it with the existing one + const updatedHistory = [...currentHistory] + updatedHistory[existingEntryIndex] = { + ...updatedHistory[existingEntryIndex], + ...newEntry + } + return updatedHistory + } else { + // If entry doesn't exist, append it + return [...currentHistory, newEntry] + } +} + +/** + * Filters uploaded files to only include images and screenshots + * + * @param files - The array of uploaded files to filter + * @returns Filtered array containing only image and screenshot files + */ +export function getImageFiles( + files?: Mojom.UploadedFile[]): Mojom.UploadedFile[] | undefined { + return files?.filter(file => + file.type === Mojom.UploadedFileType.kImage || + file.type === Mojom.UploadedFileType.kScreenshot + ); +} diff --git a/components/ai_chat/resources/page/components/attachment_button_menu/index.tsx b/components/ai_chat/resources/page/components/attachment_button_menu/index.tsx index 13c878a05cc..3a88e74111b 100644 --- a/components/ai_chat/resources/page/components/attachment_button_menu/index.tsx +++ b/components/ai_chat/resources/page/components/attachment_button_menu/index.tsx @@ -10,6 +10,7 @@ import Icon from '@brave/leo/react/icon' import { ConversationContext } from '../../state/conversation_context' import { MAX_IMAGES } from '../../../common/constants' import { AIChatContext } from '../../state/ai_chat_context' +import { getImageFiles } from '../../../common/conversation_history_utils' // Utils import { getLocale } from '$web-common/locale' @@ -17,12 +18,14 @@ import { getLocale } from '$web-common/locale' // Styles import styles from './style.module.scss' -type Props = Pick & +type Props = Pick & Pick export default function AttachmentButtonMenu(props: Props) { const totalUploadedImages = props.conversationHistory.reduce( - (total, turn) => total + (turn.uploadedImages?.length || 0), + (total, turn) => total + + (getImageFiles(turn.uploadedFiles)?.length || 0), 0 ) @@ -50,6 +53,17 @@ export default function AttachmentButtonMenu(props: Props) { {getLocale('uploadFileButtonLabel')} + {!!props.associatedContentInfo && + props.getScreenshots()}> +
+ + {getLocale('screenshotButtonLabel')} +
+ + } {props.isMobile && props.uploadImage(true)}>
diff --git a/components/ai_chat/resources/page/components/input_box/index.tsx b/components/ai_chat/resources/page/components/input_box/index.tsx index 166a933ed81..0f31bdeeba7 100644 --- a/components/ai_chat/resources/page/components/input_box/index.tsx +++ b/components/ai_chat/resources/page/components/input_box/index.tsx @@ -32,9 +32,11 @@ type Props = Pick< | 'isGenerating' | 'handleStopGenerating' | 'uploadImage' + | 'getScreenshots' | 'pendingMessageImages' | 'removeImage' | 'conversationHistory' + | 'associatedContentInfo' > & Pick @@ -175,7 +177,9 @@ function InputBox(props: InputBoxProps) { )}
diff --git a/components/ai_chat/resources/page/components/input_box/style.module.scss b/components/ai_chat/resources/page/components/input_box/style.module.scss index 18f5c1df911..4bc12e423b7 100644 --- a/components/ai_chat/resources/page/components/input_box/style.module.scss +++ b/components/ai_chat/resources/page/components/input_box/style.module.scss @@ -127,4 +127,7 @@ .attachmentWrapper { width: 100%; padding: var(--leo-spacing-m) var(--leo-spacing-m) 0px var(--leo-spacing-m); + max-height: 300px; + overflow-y: auto; + gap: var(--leo-spacing-m); } \ No newline at end of file diff --git a/components/ai_chat/resources/page/components/uploaded_img_item/index.tsx b/components/ai_chat/resources/page/components/uploaded_img_item/index.tsx index da3e7ab60e9..94eb76b71f8 100644 --- a/components/ai_chat/resources/page/components/uploaded_img_item/index.tsx +++ b/components/ai_chat/resources/page/components/uploaded_img_item/index.tsx @@ -14,7 +14,7 @@ import * as Mojom from '../../../common/mojom' import styles from './style.module.scss' type Props = { - uploadedImage: Mojom.UploadedImage + uploadedImage: Mojom.UploadedFile // removeImage is optional here so we can also reuse // this component in the conversation thread where remove // is not needed. @@ -24,7 +24,7 @@ type Props = { export default function UploadedImgItem(props: Props) { // Memos const dataUrl = React.useMemo(() => { - const blob = new Blob([new Uint8Array(props.uploadedImage.imageData)], { + const blob = new Blob([new Uint8Array(props.uploadedImage.data)], { type: 'image/*' }) return URL.createObjectURL(blob) diff --git a/components/ai_chat/resources/page/state/conversation_context.tsx b/components/ai_chat/resources/page/state/conversation_context.tsx index 45415fb79d3..09df78ce553 100644 --- a/components/ai_chat/resources/page/state/conversation_context.tsx +++ b/components/ai_chat/resources/page/state/conversation_context.tsx @@ -16,6 +16,9 @@ import getAPI from '../api' import { IGNORE_EXTERNAL_LINK_WARNING_KEY, MAX_IMAGES // } from '../../common/constants' +import { + updateConversationHistory, getImageFiles +} from '../../common/conversation_history_utils' const MAX_INPUT_CHAR = 2000 const CHAR_LIMIT_THRESHOLD = MAX_INPUT_CHAR * 0.8 @@ -26,7 +29,7 @@ export interface CharCountContext { inputTextCharCountDisplay: string } -export type UploadedImageData = Mojom.UploadedImage +export type UploadedImageData = Mojom.UploadedFile export type ConversationContext = SendFeedbackState & CharCountContext & { historyInitialized: boolean @@ -70,10 +73,11 @@ export type ConversationContext = SendFeedbackState & CharCountContext & { showAttachments: boolean setShowAttachments: (show: boolean) => void uploadImage: (useMediaCapture: boolean) => void + getScreenshots: () => void removeImage: (index: number) => void setGeneratedUrlToBeOpened: (url?: Url) => void setIgnoreExternalLinkWarning: () => void - pendingMessageImages: Mojom.UploadedImage[] | null + pendingMessageImages: Mojom.UploadedFile[] | null } export const defaultCharCountContext: CharCountContext = { @@ -117,6 +121,7 @@ const defaultContext: ConversationContext = { showAttachments: false, setShowAttachments: () => { }, uploadImage: (useMediaCapture: boolean) => { }, + getScreenshots: () => {}, removeImage: () => { }, setGeneratedUrlToBeOpened: () => { }, setIgnoreExternalLinkWarning: () => { }, @@ -216,13 +221,24 @@ export function ConversationContextProvider(props: React.PropsWithChildren) { // Initialization React.useEffect(() => { - async function updateHistory() { - const { conversationHistory } = - await conversationHandler.getConversationHistory() - setPartialContext({ - conversationHistory, - historyInitialized: true - }) + async function updateHistory(entry?: Mojom.ConversationTurn) { + if (entry) { + // Use the shared utility function to update the history + const updatedHistory = + updateConversationHistory(context.conversationHistory, entry) + setPartialContext({ + conversationHistory: updatedHistory, + historyInitialized: true + }) + } else { + // When no entry is provided, fetch the full history + const { conversationHistory } = + await conversationHandler.getConversationHistory() + setPartialContext({ + conversationHistory, + historyInitialized: true + }) + } } async function initialize() { @@ -516,17 +532,15 @@ export function ConversationContextProvider(props: React.PropsWithChildren) { aiChatContext.uiHandler?.handleVoiceRecognition(context.conversationUuid) } - const uploadImage = (useMediaCapture: boolean) => { - aiChatContext.uiHandler?.uploadImage(useMediaCapture) - .then(({uploadedImages}) => { - if (uploadedImages) { + const processUploadedImage = (images: Mojom.UploadedFile[]) => { const totalUploadedImages = context.conversationHistory.reduce( - (total, turn) => total + (turn.uploadedImages?.length || 0), + (total, turn) => total + + (getImageFiles(turn.uploadedFiles)?.length || 0), 0 ) const currentPendingImages = context.pendingMessageImages?.length || 0 const maxNewImages = MAX_IMAGES - totalUploadedImages - currentPendingImages - const newImages = uploadedImages.slice(0, Math.max(0, maxNewImages)) + const newImages = images.slice(0, Math.max(0, maxNewImages)) if (newImages.length > 0) { setPartialContext({ @@ -535,6 +549,22 @@ export function ConversationContextProvider(props: React.PropsWithChildren) { : [...newImages] }) } + } + + const getScreenshots = () => { + conversationHandler.getScreenshots() + .then(({screenshots}) => { + if (screenshots) { + processUploadedImage(screenshots) + } + }) + } + + const uploadImage = (useMediaCapture: boolean) => { + aiChatContext.uiHandler?.uploadImage(useMediaCapture) + .then(({uploadedImages}) => { + if (uploadedImages) { + processUploadedImage(uploadedImages) } }) } @@ -645,6 +675,7 @@ export function ConversationContextProvider(props: React.PropsWithChildren) { setIsToolsMenuOpen: (isToolsMenuOpen) => setPartialContext({ isToolsMenuOpen }), handleVoiceRecognition, uploadImage, + getScreenshots, removeImage, conversationHandler, setGeneratedUrlToBeOpened: diff --git a/components/ai_chat/resources/page/stories/components_panel.tsx b/components/ai_chat/resources/page/stories/components_panel.tsx index 91e6dfcc154..ea5dd38403d 100644 --- a/components/ai_chat/resources/page/stories/components_panel.tsx +++ b/components/ai_chat/resources/page/stories/components_panel.tsx @@ -124,7 +124,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -137,7 +137,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [getCompletionEvent('The ways that animals move are just about as myriad as the animal kingdom itself. They walk, run, swim, crawl, fly and slither — and within each of those categories lies a tremendous number of subtly different movement types. A seagull and a *hummingbird* both have wings, but otherwise their flight techniques and abilities are poles apart. Orcas and **piranhas** both have tails, but they accomplish very different types of swimming. Even a human walking or running is moving their body in fundamentally different ways.')], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -150,7 +150,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -163,7 +163,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [getCompletionEvent(`## How We Created an Accessible, Scalable Color Palette\n\nDuring the latter part of 2021, I reflected on the challenges we were facing at Modern Health. One recurring problem that stood out was our struggle to create new products with an unstructured color palette. This resulted in poor [communication](https://www.google.com) between designers and developers, an inconsistent product brand, and increasing accessibility problems.\n\n1. Inclusivity: our palette provides easy ways to ensure our product uses accessible contrasts.\n 2. Efficiency: our palette is diverse enough for our current and future product design, yet values are still predictable and constrained.\n 3. Reusability: our palette is on-brand but versatile. There are very few one-offs that fall outside the palette.\n\n This article shares the process I followed to apply these principles to develop a more adaptable color palette that prioritizes accessibility and is built to scale into all of our future product **design** needs.`)], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -176,7 +176,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [getPageContentRefineEvent()], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -189,7 +189,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -202,7 +202,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [getCompletionEvent('The partial sum formed by the first n + 1 terms of a Taylor series is a polynomial of degree n that is called the nth Taylor polynomial of the function. Taylor polynomials are approximations of a function, which become generally better as n increases.')], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -215,7 +215,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -228,7 +228,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [getCompletionEvent("Hello! As a helpful and respectful AI assistant, I'd be happy to assist you with your question. However, I'm a text-based AI and cannot provide code in a specific programming language like C++. Instead, I can offer a brief explanation of how to write a \"hello world\" program in C++.\n\nTo write a \"hello world\" program in C++, you can use the following code:\n\n```c++\n#include \n\nint main() {\n std::cout << \"Hello, world!\" << std::endl;\n return 0;\n}\n```\nThis code will print \"Hello, world!\" and uses `iostream` std library. If you have any further questions or need more information, please don't hesitate to ask!")], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -241,7 +241,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -254,7 +254,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [getCompletionEvent('Pointer compression is a memory optimization technique where pointers are stored in a compressed format to save memory.')], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -267,7 +267,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -289,7 +289,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ { title: 'Tesla Model Y', faviconUrl: { url: 'https://www.tesla.com/favicon.ico' }, url: { url: 'https://www.tesla.com/modely' } } ]) ], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -309,12 +309,12 @@ const HISTORY: Mojom.ConversationTurn[] = [ createdTime: { internalValue: BigInt('13278618001000000') }, edits: [], events: [], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }], createdTime: { internalValue: BigInt('13278618001000000') }, events: [], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -327,7 +327,7 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [getSearchStatusEvent(), getSearchEvent(['LTT store backpack dimensions', 'Tesla Model Y frunk dimensions'])], - uploadedImages : [], + uploadedFiles : [], fromBraveSearchSERP: false }, { @@ -340,9 +340,11 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [], - uploadedImages : [ + uploadedFiles : [ { filename: 'lion.png', filesize: 128, - imageData: Array.from(new Uint8Array(128)) } + data: Array.from(new Uint8Array(128)), + type: Mojom.UploadedFileType.kImage + } ], fromBraveSearchSERP: false }, @@ -356,7 +358,87 @@ const HISTORY: Mojom.ConversationTurn[] = [ edits: [], createdTime: { internalValue: BigInt('13278618001000000') }, events: [getCompletionEvent('It is a lion!')], - uploadedImages : [], + uploadedFiles : [], + fromBraveSearchSERP: false + }, + { + uuid: undefined, + text: 'Summarize this page', + characterType: Mojom.CharacterType.HUMAN, + actionType: Mojom.ActionType.QUERY, + prompt: undefined, + selectedText: undefined, + edits: [], + createdTime: { internalValue: BigInt('13278618001000000') }, + events: [], + uploadedFiles : [ + { filename: 'full_screenshot_0.png', filesize: 128, + data: Array.from(new Uint8Array(128)), + type: Mojom.UploadedFileType.kScreenshot + }, + { filename: 'full_screenshot_1.png', filesize: 128, + data: Array.from(new Uint8Array(128)), + type: Mojom.UploadedFileType.kScreenshot + }, + ], + fromBraveSearchSERP: false + }, + { + uuid: undefined, + text: '', + characterType: Mojom.CharacterType.ASSISTANT, + actionType: Mojom.ActionType.UNSPECIFIED, + prompt: undefined, + selectedText: undefined, + edits: [], + createdTime: { internalValue: BigInt('13278618001000000') }, + events: [ + getCompletionEvent( + 'This website compares differences between Juniper Model Y and legacy one.' + )], + uploadedFiles : [], + fromBraveSearchSERP: false + }, + { + uuid: undefined, + text: 'Summarize these', + characterType: Mojom.CharacterType.HUMAN, + actionType: Mojom.ActionType.QUERY, + prompt: undefined, + selectedText: undefined, + edits: [], + createdTime: { internalValue: BigInt('13278618001000000') }, + events: [], + uploadedFiles : [ + { filename: 'full_screenshot_0.png', filesize: 128, + data: Array.from(new Uint8Array(128)), + type: Mojom.UploadedFileType.kScreenshot + }, + { filename: 'full_screenshot_1.png', filesize: 128, + data: Array.from(new Uint8Array(128)), + type: Mojom.UploadedFileType.kScreenshot + }, + { filename: 'lion.png', filesize: 128, + data: Array.from(new Uint8Array(128)), + type: Mojom.UploadedFileType.kImage + } + ], + fromBraveSearchSERP: false + }, + { + uuid: undefined, + text: '', + characterType: Mojom.CharacterType.ASSISTANT, + actionType: Mojom.ActionType.UNSPECIFIED, + prompt: undefined, + selectedText: undefined, + edits: [], + createdTime: { internalValue: BigInt('13278618001000000') }, + events: [ + getCompletionEvent( + 'According to screenshots, this website compares differences between Juniper Model Y and legacy one. And a lion image.' + )], + uploadedFiles : [], fromBraveSearchSERP: false } ] @@ -711,6 +793,7 @@ function StoryContext(props: React.PropsWithChildren<{ args: CustomArgs, setArgs showAttachments: options.args.showAttachments, removeImage: () => {}, uploadImage: (useMediaCapture: boolean) => {}, + getScreenshots: () => {}, setGeneratedUrlToBeOpened: (url?: Url) => setArgs({ generatedUrlToBeOpened: url }), setIgnoreExternalLinkWarning: () => { } diff --git a/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/index.tsx b/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/index.tsx index 634084680fb..3726ec19cd8 100644 --- a/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/index.tsx +++ b/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/index.tsx @@ -20,6 +20,7 @@ import AssistantResponse from '../assistant_response' import EditInput from '../edit_input' import EditIndicator from '../edit_indicator' import { getReasoningText } from './conversation_entries_utils' +import { getImageFiles } from '../../../common/conversation_history_utils' import styles from './style.module.scss' function ConversationEntries() { @@ -175,13 +176,19 @@ function ConversationEntries() { )} - {!!latestTurn.uploadedImages?.length && - latestTurn.uploadedImages.map((img) => ( - - ))} +
+ {(() => { + const imageFiles = + getImageFiles(latestTurn.uploadedFiles) || []; + return imageFiles.length > 0 && + imageFiles.map((img) => ( + + )); + })()} +
)} diff --git a/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/style.module.scss b/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/style.module.scss index d00853c3225..794cc583a7c 100644 --- a/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/style.module.scss +++ b/components/ai_chat/resources/untrusted_conversation_frame/components/conversation_entries/style.module.scss @@ -78,3 +78,9 @@ .divToKeepGap { min-width: 16px; } + +.uploadedImages { + max-height: 300px; + overflow-y: auto; + gap: var(--leo-spacing-m); +} diff --git a/components/ai_chat/resources/untrusted_conversation_frame/untrusted_conversation_frame_api.ts b/components/ai_chat/resources/untrusted_conversation_frame/untrusted_conversation_frame_api.ts index 5d462133163..75b65c268b3 100644 --- a/components/ai_chat/resources/untrusted_conversation_frame/untrusted_conversation_frame_api.ts +++ b/components/ai_chat/resources/untrusted_conversation_frame/untrusted_conversation_frame_api.ts @@ -6,6 +6,7 @@ import { loadTimeData } from '$web-common/loadTimeData' import API from '../common/api' import * as Mojom from '../common/mojom' +import { updateConversationHistory } from '../common/conversation_history_utils' // Global state for this UI export type ConversationEntriesUIState = Mojom.ConversationEntriesState & { @@ -60,7 +61,21 @@ export default class UntrustedConversationFrameAPI extends API this.setPartialState(await this.conversationHandler.getConversationHistory()) + async (entry?: Mojom.ConversationTurn) => { + if (entry) { + // Use the shared utility function to update the history + const updatedHistory = + updateConversationHistory(this.state.conversationHistory, entry) + this.setPartialState({ + conversationHistory: updatedHistory + }) + } else { + // When no entry is provided, fetch the full history + const { conversationHistory } = + await this.conversationHandler.getConversationHistory() + this.setPartialState({ conversationHistory }) + } + } ) this.conversationObserver.onEntriesUIStateChanged.addListener((state: Mojom.ConversationEntriesState) => { diff --git a/components/ai_rewriter/resources/page/components/BeginGeneration.tsx b/components/ai_rewriter/resources/page/components/BeginGeneration.tsx index 861342b197a..61a2d2080a3 100644 --- a/components/ai_rewriter/resources/page/components/BeginGeneration.tsx +++ b/components/ai_rewriter/resources/page/components/BeginGeneration.tsx @@ -56,6 +56,7 @@ export default function BeginGeneration() { handleStopGenerating: async () => {}, removeImage: () => {}, uploadImage: (useMediaCapture: boolean) => {}, + getScreenshots: () => {}, conversationHistory: [], pendingMessageImages: null }} /> diff --git a/ios/brave-ios/Sources/AIChat/Components/AIChatView.swift b/ios/brave-ios/Sources/AIChat/Components/AIChatView.swift index 6587e7e8e6c..16a9537405a 100644 --- a/ios/brave-ios/Sources/AIChat/Components/AIChatView.swift +++ b/ios/brave-ios/Sources/AIChat/Components/AIChatView.swift @@ -783,7 +783,7 @@ struct AIChatView_Preview: PreviewProvider { events: nil, createdTime: Date.now, edits: nil, - uploadedImages: nil, + uploadedFiles: nil, fromBraveSearchSerp: false ), isEntryInProgress: false, diff --git a/ios/brave-ios/Sources/AIChat/Components/Messages/AIChatResponseMessageView.swift b/ios/brave-ios/Sources/AIChat/Components/Messages/AIChatResponseMessageView.swift index 8475954809d..778ccc2866e 100644 --- a/ios/brave-ios/Sources/AIChat/Components/Messages/AIChatResponseMessageView.swift +++ b/ios/brave-ios/Sources/AIChat/Components/Messages/AIChatResponseMessageView.swift @@ -357,7 +357,7 @@ struct AIChatResponseMessageView_Previews: PreviewProvider { events: nil, createdTime: Date.now, edits: nil, - uploadedImages: nil, + uploadedFiles: nil, fromBraveSearchSerp: false ), isEntryInProgress: false, diff --git a/ios/browser/api/ai_chat/conversation_client.h b/ios/browser/api/ai_chat/conversation_client.h index 58d49976baa..49f6a80355d 100644 --- a/ios/browser/api/ai_chat/conversation_client.h +++ b/ios/browser/api/ai_chat/conversation_client.h @@ -34,7 +34,7 @@ class ConversationClient : public mojom::ConversationUI, protected: // mojom::ConversationUI - void OnConversationHistoryUpdate() override; + void OnConversationHistoryUpdate(mojom::ConversationTurnPtr entry) override; void OnAPIRequestInProgress(bool is_request_in_progress) override; void OnAPIResponseError(mojom::APIError error) override; void OnModelDataChanged( diff --git a/ios/browser/api/ai_chat/conversation_client.mm b/ios/browser/api/ai_chat/conversation_client.mm index eb8004ea9f7..83d3c43e932 100644 --- a/ios/browser/api/ai_chat/conversation_client.mm +++ b/ios/browser/api/ai_chat/conversation_client.mm @@ -35,7 +35,8 @@ void ConversationClient::ChangeConversation(ConversationHandler* conversation) { // MARK: - mojom::ConversationUI -void ConversationClient::OnConversationHistoryUpdate() { +void ConversationClient::OnConversationHistoryUpdate( + mojom::ConversationTurnPtr entry) { [bridge_ onHistoryUpdate]; } diff --git a/test/data/ai_chat/aichat_database_dump_version_1.sql b/test/data/ai_chat/aichat_database_dump_version_1.sql index bef117ef103..f4d74733b45 100644 --- a/test/data/ai_chat/aichat_database_dump_version_1.sql +++ b/test/data/ai_chat/aichat_database_dump_version_1.sql @@ -16,4 +16,5 @@ INSERT INTO conversation_entry VALUES('f761af26-4491-4787-ad53-82238b42f534','1a CREATE TABLE conversation_entry_event_completion(conversation_entry_uuid INTEGER NOT NULL,event_order INTEGER NOT NULL,text BLOB NOT NULL,PRIMARY KEY(conversation_entry_uuid, event_order)); INSERT INTO conversation_entry_event_completion VALUES('f761af26-4491-4787-ad53-82238b42f534',0,X'76313080a3088aa5874d731e9b5042cb565dab897e3af6560681aad24ef9c9379de805096188372c73372de010da70dd9231cfd70490087fd074004e3c826852cb400903a589baf91e4b3ad75acfe8b40b2bc6846b9c00edd6a3d32359c98614383d35'); CREATE TABLE conversation_entry_event_search_queries(conversation_entry_uuid INTEGER NOT NULL,event_order INTEGER NOT NULL,queries BLOB NOT NULL,PRIMARY KEY(conversation_entry_uuid, event_order)); +CREATE TABLE conversation_entry_uploaded_files(conversation_entry_uuid INTEGER NOT NULL,file_order INTEGER NOT NULL,filename BLOB NOT NULL,filesize INTEGER NOT NULL,data BLOB NOT NULL,PRIMARY KEY(conversation_entry_uuid, file_order)); COMMIT; diff --git a/test/data/ai_chat/aichat_database_dump_version_2.sql b/test/data/ai_chat/aichat_database_dump_version_2.sql index 306d684bcc0..74176e8331e 100644 --- a/test/data/ai_chat/aichat_database_dump_version_2.sql +++ b/test/data/ai_chat/aichat_database_dump_version_2.sql @@ -16,4 +16,5 @@ INSERT INTO conversation_entry VALUES('f761af26-4491-4787-ad53-82238b42f534','1a CREATE TABLE conversation_entry_event_completion(conversation_entry_uuid INTEGER NOT NULL,event_order INTEGER NOT NULL,text BLOB NOT NULL,PRIMARY KEY(conversation_entry_uuid, event_order)); INSERT INTO conversation_entry_event_completion VALUES('f761af26-4491-4787-ad53-82238b42f534',0,X'76313080a3088aa5874d731e9b5042cb565dab897e3af6560681aad24ef9c9379de805096188372c73372de010da70dd9231cfd70490087fd074004e3c826852cb400903a589baf91e4b3ad75acfe8b40b2bc6846b9c00edd6a3d32359c98614383d35'); CREATE TABLE conversation_entry_event_search_queries(conversation_entry_uuid INTEGER NOT NULL,event_order INTEGER NOT NULL,queries BLOB NOT NULL,PRIMARY KEY(conversation_entry_uuid, event_order)); +CREATE TABLE conversation_entry_uploaded_files(conversation_entry_uuid INTEGER NOT NULL,file_order INTEGER NOT NULL,filename BLOB NOT NULL,filesize INTEGER NOT NULL,data BLOB NOT NULL,PRIMARY KEY(conversation_entry_uuid, file_order)); COMMIT; diff --git a/test/data/ai_chat/aichat_database_dump_version_3.sql b/test/data/ai_chat/aichat_database_dump_version_3.sql new file mode 100644 index 00000000000..dbf6ed5dc7e --- /dev/null +++ b/test/data/ai_chat/aichat_database_dump_version_3.sql @@ -0,0 +1,24 @@ +PRAGMA foreign_keys=OFF; +BEGIN TRANSACTION; +CREATE TABLE meta(key LONGVARCHAR NOT NULL UNIQUE PRIMARY KEY, value LONGVARCHAR); +INSERT INTO meta VALUES('mmap_status','-1'); +INSERT INTO meta VALUES('version','3'); +INSERT INTO meta VALUES('last_compatible_version','1'); +CREATE TABLE conversation(uuid TEXT PRIMARY KEY NOT NULL,title BLOB,model_key TEXT,total_tokens INTEGER NOT NULL,trimmed_tokens INTEGER NOT NULL); +INSERT INTO conversation VALUES('1ae484fe-ab33-4f42-8813-14080e4addc1',X'763130a0cdc6b78d502119b80ada83b3b86e4ca7b0a7729e1b079a74c479ef1c3b4fb8',NULL,0,0); +INSERT INTO conversation VALUES('1ae484fe-ab33-4f42-8813-14080e4addc2',X'763130a0cdc6b78d502119b80ada83b3b86e4ca7b0a7729e1b079a74c479ef1c3b4fb8',NULL,0,0); +CREATE TABLE associated_content(uuid TEXT PRIMARY KEY NOT NULL,conversation_uuid TEXT NOT NULL,title BLOB,url BLOB,content_type INTEGER NOT NULL,last_contents BLOB,content_used_percentage INTEGER NOT NULL,is_content_refined INTEGER NOT NULL); +INSERT INTO associated_content (uuid, conversation_uuid, title, url, content_type, last_contents, content_used_percentage, is_content_refined) VALUES ('uuid1', '1ae484fe-ab33-4f42-8813-14080e4addc1', X'', X'', 1, X'', 50, 1); +INSERT INTO associated_content (uuid, conversation_uuid, title, url, content_type, last_contents, content_used_percentage, is_content_refined) VALUES ('uuid2', '1ae484fe-ab33-4f42-8813-14080e4addc2', X'', X'', 1, X'', 50, 1); +CREATE TABLE conversation_entry(uuid TEXT PRIMARY KEY NOT NULL,conversation_uuid STRING NOT NULL,date INTEGER NOT NULL,entry_text BLOB,prompt BLOB,character_type INTEGER NOT NULL,editing_entry_uuid TEXT,action_type INTEGER,selected_text BLOB); +INSERT INTO conversation_entry VALUES('5616a89c-7f56-4e7d-8e74-f882b76623a7','1ae484fe-ab33-4f42-8813-14080e4addc1',13379152669533821,X'763130d5244f5d32beef52f593859d25b33986df182571da857ea74ba49d4c14d72a2e84fc8c4c4df6bfb85728f5249fa9e22e',X'',0,NULL,5,NULL); +INSERT INTO conversation_entry VALUES('f761af26-4491-4787-ad53-82238b42f534','1ae484fe-ab33-4f42-8813-14080e4addc1',13379152671156556,X'76313080a3088aa5874d731e9b5042cb565dab897e3af6560681aad24ef9c9379de805096188372c73372de010da70dd9231cfd70490087fd074004e3c826852cb400903a589baf91e4b3ad75acfe8b40b2bc6846b9c00edd6a3d32359c98614383d35',X'',1,NULL,1,NULL); +CREATE TABLE conversation_entry_event_completion(conversation_entry_uuid INTEGER NOT NULL,event_order INTEGER NOT NULL,text BLOB NOT NULL,PRIMARY KEY(conversation_entry_uuid, event_order)); +INSERT INTO conversation_entry_event_completion VALUES('f761af26-4491-4787-ad53-82238b42f534',0,X'76313080a3088aa5874d731e9b5042cb565dab897e3af6560681aad24ef9c9379de805096188372c73372de010da70dd9231cfd70490087fd074004e3c826852cb400903a589baf91e4b3ad75acfe8b40b2bc6846b9c00edd6a3d32359c98614383d35'); +CREATE TABLE conversation_entry_event_search_queries(conversation_entry_uuid INTEGER NOT NULL,event_order INTEGER NOT NULL,queries BLOB NOT NULL,PRIMARY KEY(conversation_entry_uuid, event_order)); +CREATE TABLE conversation_entry_uploaded_files(conversation_entry_uuid INTEGER NOT NULL,file_order INTEGER NOT NULL,filename BLOB NOT NULL,filesize INTEGER NOT NULL,data BLOB NOT NULL,PRIMARY KEY(conversation_entry_uuid, file_order)); +INSERT INTO conversation_entry_uploaded_files VALUES('5616a89c-7f56-4e7d-8e74-f882b76623a7', 0, 'brave_logo.png', 4, X'deadbeef'); +INSERT INTO conversation_entry_uploaded_files VALUES('5616a89c-7f56-4e7d-8e74-f882b76623a7', 1, 'lion.png', 4, X'deadbeef'); +INSERT INTO conversation_entry_uploaded_files VALUES('f761af26-4491-4787-ad53-82238b42f534', 0, 'dragon.png', 4, X'deadbeef'); +COMMIT; +