Image generation with Dall-e

app/jobs/get_next_ai_message_job.rb

@@ -1,3 +1,4 @@
+require 'open-uri'
 include ActionView::RecordIdentifier
 require "nokogiri/xml/node"
 
@@ -190,6 +191,7 @@ def call_tools_before_wrapping_up
     end
 
     index = @message.index
+    url_of_dalle_generated_image = nil
     msgs.each do |tool_message| # one message for each tool executed
       @conversation.messages.create!(
         assistant: @assistant,
@@ -200,6 +202,13 @@ def call_tools_before_wrapping_up
         index: index += 1,
         processed_at: Time.current,
       )
+
+      parsed = JSON.parse(tool_message[:content]) rescue nil
+
+      if parsed.is_a?(Hash) && parsed.has_key?("url_of_dalle_generated_image")
+        url_of_dalle_generated_image = parsed["url_of_dalle_generated_image"]
+      end
+
     end
 
     assistant_reply = @conversation.messages.create!(
@@ -210,6 +219,12 @@ def call_tools_before_wrapping_up
       index: index += 1
     )
 
+    unless url_of_dalle_generated_image.nil?
+      d = Document.new
+      d.file.attach(io: URI.open(url_of_dalle_generated_image), filename: 'image.png')
+      assistant_reply.documents << d
+    end
+
     GetNextAIMessageJob.perform_later(
       @user.id,
       assistant_reply.id,

app/services/toolbox.rb

@@ -6,6 +6,7 @@ def self.descendants
     [
       test_env && Toolbox::HelloWorld,
       Toolbox::OpenMeteo,
+      Toolbox::Dalle,
       Toolbox::Memory,
       gmail_active && Toolbox::Gmail,
       tasks_active && Toolbox::GoogleTasks,

app/services/toolbox/dalle.rb

@@ -0,0 +1,36 @@
+class Toolbox::Dalle < Toolbox
+
+  describe :generate_an_image, <<~S
+    Generate an image based on what the user asks you to generate. You will pass the user's prompt and will get back a URL to an image.
+  S
+
+  def self.generate_an_image(image_generation_prompt_s:)
+
+    response = client.images.generate(
+      parameters: {
+        prompt: image_generation_prompt_s,
+        model: "dall-e-3",
+        size: "1024x1792",
+        quality: "standard"
+      }
+    )
+
+    dalle_url = response.dig("data", 0, "url")
+
+    {
+      prompt_given: image_generation_prompt_s,
+      url_of_dalle_generated_image: dalle_url,
+      note_to_assistant: "The image at the URL is already being shown on screen so reply with a nice message confirming the image has been generated, maybe re-describing it, but don't include the link to it."
+    }
+  end
+
+  class << self
+    private
+
+    def client
+      OpenAI::Client.new(
+        access_token: Current.user.openai_key,
+      )
+    end
+  end
+end

lib/active_storage/service/postgresql_service.rb

@@ -103,7 +103,7 @@ def generate_url(key, expires_in:, filename:, disposition:, content_type:)
         )
 
         generated_url = url_helpers.rails_postgresql_service_url(verified_key_with_expiration,
-          **url_options,
+          only_path: true,  # This fixes an exception with attachment URL generation from a worker: https://github.com/AllYourBot/hostedgpt/pull/398#issuecomment-2168135853
           disposition: content_disposition,
           content_type: content_type,
           filename: filename,

test/fixtures/conversations.yml

@@ -76,6 +76,12 @@ weather:
   title: Weather
   last_assistant_message: weather_explained
 
+image_generation:
+  user: christoph
+  assistant: samantha
+  title: Generating an image
+  last_assistant_message: image_generation_explained
+
 trees:
   user: keith
   assistant: keith_gpt3

test/fixtures/messages.yml

@@ -660,6 +660,45 @@ weather_explained:
   version: 1
 
 
+image_generation_tool_call:
+  assistant: samantha
+  conversation: image_generation
+  role: assistant
+  tool_call_id:
+  content_text:
+  content_tool_calls: '[{"index": 0, "id": "def456", "type": "function", "function": {"name": "generate_an_image", "arguments": {"image_generation_prompt_s": "World"}}}]'
+  content_document:
+  created_at: 2024-08-25 1:01:00
+  processed_at: 2024-08-25 1:01:00
+  index: 1
+  version: 1
+
+image_generation_tool_result:
+  assistant: samantha
+  conversation: image_generation
+  role: tool
+  tool_call_id: def456
+  content_text: weather is
+  content_tool_calls:
+  content_document:
+  created_at: 2024-08-25 1:02:00
+  processed_at: 2024-08-25 1:02:00
+  index: 2
+  version: 1
+
+image_generation_explained:
+  assistant: samantha
+  conversation: image_generation
+  role: assistant
+  tool_call_id:
+  content_text: The weather in Austin is
+  content_tool_calls:
+  content_document:
+  created_at: 2024-08-25 1:03:00
+  processed_at: 2024-08-25 1:03:00
+  index: 3
+  version: 1 
+
 # Next conversation
 
 trees_explained:

test/fixtures/users.yml

@@ -14,3 +14,8 @@ taylor:
   first_name: Taylor
   registered_at: 2024-05-31 08:40:05
   preferences: {}
+
+christoph:
+  first_name: Christoph
+  registered_at: 2024-08-25 08:08:08
+  preferences: {}

test/jobs/get_next_ai_message_job_openai_test.rb

@@ -3,10 +3,19 @@
 class GetNextAIMessageJobOpenaiTest < ActiveJob::TestCase
   setup do
     @conversation = conversations(:greeting)
-    @user = @conversation.user
     @conversation.messages.create! role: :user, content_text: "Still there?", assistant: @conversation.assistant
     @message = @conversation.latest_message_for_version(:latest)
-    @test_client = TestClient::OpenAI.new(access_token: "abc")
+
+    p "#########"
+    @image_generation = conversations(:image_generation)
+    p @image_generation
+    p @image_generation.messages.create! role: :user, content_text: "Generate an image", assistant: @image_generation.assistant
+    p @image_generation.messages.length
+    @image_generation_message = @image_generation.latest_message_for_version(:latest)
+    p @image_generation_message
+
+    @user = @conversation.user
+    @test_client = TestClient::OpenAI.new(access_token: 'abc')
   end
 
   test "populates the latest message from the assistant" do
@@ -56,6 +65,52 @@ class GetNextAIMessageJobOpenaiTest < ActiveJob::TestCase
     refute second_new_message.finished?, "This message SHOULD NOT be considered finished yet"
   end
 
+  test "properly handles a tool response call from the assistant when images are included" do
+    assert_difference "@image_generation.messages.reload.length", 2 do
+      TestClient::OpenAI.stub :function, "generate_an_image" do
+        TestClient::OpenAI.stub :arguments, { image_generation_prompt_s:  "Kitten" } do
+          TestClient::OpenAI.stub :api_response, TestClient::OpenAI.api_function_response do
+            assert GetNextAIMessageJob.perform_now(@user.id, @message.id, @image_generation.assistant.id)
+          end
+        end
+      end
+    end
+
+    p "###### Message:"
+    p @image_generation_message
+
+    @image_generation_message.reload
+    assert @image_generation_message.content_text.blank?
+    assert @image_generation_message.tool_call_id.nil?
+    assert @image_generation_message.content_tool_calls.present?, "Assistant should have decided to call a tool"
+
+    @new_messages = @image_generation.messages.where("id > ?", @message.id).order(:created_at)
+
+    p "###### Image generation:"
+    p @image_generation
+    @new_messages.messages.each_with_index do |message, index|
+      p "######## #{index}:"
+      p message
+    end
+
+    # first
+    first_new_message = @new_messages.first
+    assert first_new_message.tool?
+    # assert_equal "Hello, Keith!".to_json, first_new_message.content_text, "First new message should have the result of calling the tool"
+    # assert first_new_message.tool_call_id.present?
+    # assert first_new_message.content_tool_calls.blank?
+    # assert_equal @message.content_tool_calls.dig(0, :id), first_new_message.tool_call_id, "ID of tool execution should have matched decision to call the tool"
+    # assert first_new_message.finished?, "This message SHOULD HAVE been considered finished"
+
+    # second
+    # second_new_message = @new_messages.second
+    # assert second_new_message.assistant?, "Second new message should be queued up for the assistant to reply"
+    # assert second_new_message.content_text.nil?, "The content should be nil to indicate that it hasn't even started processing"
+    # assert second_new_message.tool_call_id.nil?
+    # assert second_new_message.content_tool_calls.blank?
+    # refute second_new_message.finished?, "This message SHOULD NOT be considered finished yet"
+  end
+
   test "returns early if the message id was invalid" do
     refute GetNextAIMessageJob.perform_now(@user.id, 0, @conversation.assistant.id)
   end

test/lib/active_storage/postgresql_test.rb

@@ -48,7 +48,7 @@ class ActiveStorage::PostgresqlTest < ActiveSupport::TestCase
   end
 
   test "url generation" do
-    assert_match(/^https:\/\/example.com\/rails\/active_storage\/postgresql\/.*\/avatar\.png\?content_type=image%2Fpng&disposition=inline/,
+    assert_match(/^\/rails\/active_storage\/postgresql\/.*\/avatar\.png\?content_type=image%2Fpng&disposition=inline/,
       @service.url(FIXTURE_KEY, expires_in: 5.minutes, disposition: :inline, filename: ActiveStorage::Filename.new("avatar.png"), content_type: "image/png"))
   end
 

test/lib/active_storage/service/public_postgresql_test.rb

@@ -8,6 +8,6 @@ class ActiveStorage::Service::PublicPostgresqlTest < ActiveSupport::TestCase
   test "public URL generation" do
     url = @service.url(@key, disposition: :inline, filename: ActiveStorage::Filename.new("avatar.png"), content_type: "image/png")
 
-    assert_match(/^https:\/\/example.com\/rails\/active_storage\/postgresql\/.*\/avatar\.png/, url)
+    assert_match(/^\/rails\/active_storage\/postgresql\/.*\/avatar\.png/, url)
   end
 end

test/models/document_test.rb

@@ -55,9 +55,8 @@ class DocumentTest < ActiveSupport::TestCase
   end
 
   test "fully_processed_url" do
-    assert documents(:cat_photo).fully_processed_url(:small).starts_with?("http")
-    assert documents(:cat_photo).fully_processed_url(:small).include?("rails/active_storage/postgresql")
-    assert documents(:cat_photo).fully_processed_url(:small).exclude?("/redirect")
+    assert documents(:cat_photo).fully_processed_url(:small).include?('rails/active_storage/postgresql')
+    assert documents(:cat_photo).fully_processed_url(:small).exclude?('/redirect')
   end
 
   test "redirect_to_processed_path" do