Fix bugs (#226)

* test fix

* fix quip errors

* catch soql error

* security bug

---------

Co-authored-by: Vijay Swamidass <[email protected]>

Author: vswamidass-sfdc

Gemfile

@@ -135,4 +135,4 @@ gem 'omniauth-google-oauth2'
 gem 'google-api-client'
 gem 'googleauth'
 
-gem 'uri', '>= 1.0.3'
+gem 'uri', '>= 1.0.4'

Gemfile.lock

@@ -497,7 +497,7 @@ GEM
     unicode-display_width (3.1.4)
       unicode-emoji (~> 4.0, >= 4.0.4)
     unicode-emoji (4.0.4)
-    uri (1.0.3)
+    uri (1.0.4)
     useragent (0.16.11)
     version_gem (1.1.8)
     wasabi (5.1.0)
@@ -582,7 +582,7 @@ DEPENDENCIES
   tiktoken_ruby
   turbo-rails
   tzinfo-data
-  uri (>= 1.0.3)
+  uri (>= 1.0.4)
   web-console
   webdrivers
   websocket-client-simple

app/jobs/generate_message_response_job.rb

@@ -133,17 +133,20 @@ def perform(_message_id)
 
       # QUIP Doc
       if assistant.quip_url.present?
-        quip_client = Quip::Client.new(access_token: ENV.fetch('QUIP_TOKEN'))
-
         uri = URI.parse(assistant.quip_url)
-        path = uri.path.sub(%r{^/}, '') # Removes the leading /
-        quip_thread = quip_client.get_thread(path)
-
-        prompt += "# QUIP DOCUMENT\n\n"
-        # The quip api only returns html which has too much extra junk.
-        # Convert to md for smaller size
-        markdown_quip = ReverseMarkdown.convert quip_thread['html']
-        prompt += markdown_quip
+        # Some quip_urls have google and other things in the host.  skip this if it doesn't have quip.com in the host
+        if uri.host&.include?('quip.com')
+          quip_client = Quip::Client.new(access_token: ENV.fetch('QUIP_TOKEN'))
+          path = uri.path.sub(%r{^/}, '') # Removes the leading /
+
+          quip_thread = quip_client.get_thread(path)
+
+          prompt += "# QUIP DOCUMENT\n\n"
+          # The quip api only returns html which has too much extra junk.
+          # Convert to md for smaller size
+          markdown_quip = ReverseMarkdown.convert quip_thread['html']
+          prompt += markdown_quip
+        end
       end
 
       if assistant.confluence_spaces.present?
@@ -164,7 +167,11 @@ def perform(_message_id)
         salesforce_client = Salesforce::Client.new
 
         salesforce_results = salesforce_client.query(assistant.soql)
-        prompt += JSON.pretty_generate(salesforce_results.map(&:to_h))
+        prompt += if salesforce_results&.any?
+                    JSON.pretty_generate(salesforce_results.map(&:to_h))
+                  else
+                    'No results from query.'
+                  end
         prompt += '</SOQL>'
       end
 
@@ -217,12 +224,8 @@ def perform(_message_id)
       llm_message.from = :assistant
 
       begin
-        start_time = Time.now
         llm_message.generating!
         llm_message.content = get_generation(llm_message.prompt)
-        end_time = Time.now
-
-        generation_time = end_time - start_time
 
         llm_message.save
         llm_message.ready!

app/models/assistant.rb

@@ -7,10 +7,11 @@ class Assistant < ApplicationRecord
 
   has_many :assistant_users, dependent: :destroy
   has_many :users, through: :assistant_users
-  enum status: { development: 0, ready: 1 }
+  enum :status, { development: 0, ready: 1 }
   validates :name, presence: true
   validates :slack_channel_name, uniqueness: true, allow_blank: true
   validate :slack_channel_name_starts_with_unique
+  validate :quip_url_must_contain_quip_domain
 
   validate :libraries_must_be_csv_with_numbers
 
@@ -61,4 +62,12 @@ def libraries_must_be_csv_with_numbers
 
     errors.add(:libraries, 'must be a valid CSV format with only numbers')
   end
+
+  def quip_url_must_contain_quip_domain
+    return if quip_url.blank?
+
+    return if quip_url.include?('quip.com')
+
+    errors.add(:quip_url, 'Only quip urls are supported.')
+  end
 end

app/models/library_user.rb

@@ -2,5 +2,5 @@ class LibraryUser < ApplicationRecord
   belongs_to :user
   belongs_to :library
 
-  enum role: { admin: 0, editor: 1 }
+  enum :role, { admin: 0, editor: 1 }
 end

app/models/message.rb

@@ -6,8 +6,8 @@ class Message < ApplicationRecord
   validates :from, presence: true
   validates :content, presence: true
 
-  enum from: { user: 0, assistant: 1 }
-  enum status: { ready: 0, generating: 1 }
+  enum :from, { user: 0, assistant: 1 }
+  enum :status, { ready: 0, generating: 1 }
 
   after_commit :broadcast_message, on: %i[create update]
 

app/models/question.rb

@@ -11,7 +11,7 @@ class Question < ApplicationRecord
   # Tracking helpful answers
   acts_as_votable
 
-  enum status: { pending: 0, generating: 1, generated: 2, failed: 3 }
+  enum :status, { pending: 0, generating: 1, generated: 2, failed: 3 }
 
   validates :question, presence: true
   belongs_to :user, optional: true

app/models/webhook.rb

@@ -3,7 +3,7 @@ class Webhook < ApplicationRecord
   belongs_to :library
   has_many :chats
 
-  enum hook_type: { pagerduty: 0 }
+  enum :hook_type, { pagerduty: 0 }
   validates :hook_type, presence: true
 
   before_create :generate_secret_key

spec/models/assistant_spec.rb

@@ -22,5 +22,34 @@
         expect(assistant.errors[:libraries]).to include('must be a valid CSV format with only numbers')
       end
     end
+
+    context 'quip_url validation' do
+      it 'is valid with a quip.com URL' do
+        assistant = Assistant.new(name: 'test', quip_url: 'https://example.quip.com/document/123', input: 'input', user:, instructions: 'instructions', output: 'output')
+        expect(assistant).to be_valid
+      end
+
+      it 'is valid with blank quip_url' do
+        assistant = Assistant.new(name: 'test', quip_url: '', input: 'input', user:, instructions: 'instructions', output: 'output')
+        expect(assistant).to be_valid
+      end
+
+      it 'is valid with nil quip_url' do
+        assistant = Assistant.new(name: 'test', quip_url: nil, input: 'input', user:, instructions: 'instructions', output: 'output')
+        expect(assistant).to be_valid
+      end
+
+      it 'is not valid with non-quip URL' do
+        assistant = Assistant.new(name: 'test', quip_url: 'https://google.com/document/123', input: 'input', user:, instructions: 'instructions', output: 'output')
+        expect(assistant).not_to be_valid
+        expect(assistant.errors[:quip_url]).to include('Only quip urls are supported.')
+      end
+
+      it 'is not valid with URL that does not contain quip.com' do
+        assistant = Assistant.new(name: 'test', quip_url: 'https://example.com/quip/document/123', input: 'input', user:, instructions: 'instructions', output: 'output')
+        expect(assistant).not_to be_valid
+        expect(assistant.errors[:quip_url]).to include('Only quip urls are supported.')
+      end
+    end
   end
 end