smart search: add v2 embedding strategy

Add `cohere.embed-multilingual-v3` as new embedding strategy for smart
search.

Test plan:
* Add bedrock creds to rails credential store
* Make sure that background jobs processor is running
* Open search page on existing indexed course
* Confirm that course is re-indexed and run searches
* Confirm that newly enabled courses are also index properly

flag=smart_search
closes ADV-98

Change-Id: I664a954f34e5c474db0b210f2ae092f891d60f89
Reviewed-on: https://gerrit.instructure.com/c/canvas-lms/+/346524
Reviewed-by: Jeremy Stanley <jeremy@instructure.com>
Reviewed-by: Jacob Windle <jacob.windle@instructure.com>
QA-Review: Jeremy Stanley <jeremy@instructure.com>
Tested-by: Service Cloud Jenkins <svc.cloudjenkins@instructure.com>
Product-Review: Jonathan Featherstone <jfeatherstone@instructure.com>
Build-Review: Aaron Ogata <aogata@instructure.com>
This commit is contained in:
Jonathan Featherstone 2024-04-30 17:00:28 -06:00
parent b41e5785e1
commit 95d9443b10
7 changed files with 50 additions and 45 deletions

View File

@ -33,6 +33,7 @@ gem "active_model_serializers", "~> 0.9.9"
gem "addressable", "~> 2.8", require: false
gem "authlogic", "~> 6.4"
gem "scrypt", "~> 3.0"
gem "aws-sdk-bedrockruntime", "~> 1.7", require: false
gem "aws-sdk-kinesis", "~> 1.45", require: false
gem "aws-sdk-s3", "~> 1.119", require: false
gem "aws-sdk-sns", "~> 1.60", require: false

View File

@ -384,8 +384,8 @@ GEM
aws-sdk-autoscaling (1.107.0)
aws-sdk-core (~> 3, >= 3.193.0)
aws-sigv4 (~> 1.1)
aws-sdk-bedrockruntime (1.5.0)
aws-sdk-core (~> 3, >= 3.191.0)
aws-sdk-bedrockruntime (1.7.0)
aws-sdk-core (~> 3, >= 3.193.0)
aws-sigv4 (~> 1.1)
aws-sdk-core (3.194.2)
aws-eventstream (~> 1, >= 1.3.0)
@ -1194,6 +1194,7 @@ DEPENDENCIES
attachment_fu!
authlogic (~> 6.4)
autoextend!
aws-sdk-bedrockruntime (~> 1.7)
aws-sdk-kinesis (~> 1.45)
aws-sdk-s3 (~> 1.119)
aws-sdk-sagemakerruntime (~> 1.61)

View File

@ -393,8 +393,8 @@ GEM
aws-sdk-autoscaling (1.107.0)
aws-sdk-core (~> 3, >= 3.193.0)
aws-sigv4 (~> 1.1)
aws-sdk-bedrockruntime (1.5.0)
aws-sdk-core (~> 3, >= 3.191.0)
aws-sdk-bedrockruntime (1.7.0)
aws-sdk-core (~> 3, >= 3.193.0)
aws-sigv4 (~> 1.1)
aws-sdk-core (3.194.2)
aws-eventstream (~> 1, >= 1.3.0)
@ -1211,6 +1211,7 @@ DEPENDENCIES
attachment_fu!
authlogic (~> 6.4)
autoextend!
aws-sdk-bedrockruntime (~> 1.7)
aws-sdk-kinesis (~> 1.45)
aws-sdk-s3 (~> 1.119)
aws-sdk-sagemakerruntime (~> 1.61)

View File

@ -18,24 +18,10 @@ GEM
minitest (>= 5.1)
tzinfo (~> 2.0)
concurrent-ruby (1.2.3)
diff-lcs (1.5.1)
i18n (1.14.5)
concurrent-ruby (~> 1.0)
minitest (5.22.3)
rake (13.2.1)
rspec (3.13.0)
rspec-core (~> 3.13.0)
rspec-expectations (~> 3.13.0)
rspec-mocks (~> 3.13.0)
rspec-core (3.13.0)
rspec-support (~> 3.13.0)
rspec-expectations (3.13.0)
diff-lcs (>= 1.2.0, < 2.0)
rspec-support (~> 3.13.0)
rspec-mocks (3.13.1)
diff-lcs (>= 1.2.0, < 2.0)
rspec-support (~> 3.13.0)
rspec-support (3.13.1)
tzinfo (2.0.6)
concurrent-ruby (~> 1.0)
@ -50,7 +36,6 @@ DEPENDENCIES
activerecord (>= 3.2)
bundler (~> 2.2)
rake
rspec (~> 3.12)
workflow!
RUBY VERSION

View File

@ -27,26 +27,12 @@ GEM
bigdecimal (3.1.8)
concurrent-ruby (1.2.3)
connection_pool (2.4.1)
diff-lcs (1.5.1)
drb (2.2.1)
i18n (1.14.5)
concurrent-ruby (~> 1.0)
minitest (5.22.3)
mutex_m (0.2.0)
rake (13.2.1)
rspec (3.13.0)
rspec-core (~> 3.13.0)
rspec-expectations (~> 3.13.0)
rspec-mocks (~> 3.13.0)
rspec-core (3.13.0)
rspec-support (~> 3.13.0)
rspec-expectations (3.13.0)
diff-lcs (>= 1.2.0, < 2.0)
rspec-support (~> 3.13.0)
rspec-mocks (3.13.1)
diff-lcs (>= 1.2.0, < 2.0)
rspec-support (~> 3.13.0)
rspec-support (3.13.1)
timeout (0.4.1)
tzinfo (2.0.6)
concurrent-ruby (~> 1.0)
@ -62,7 +48,6 @@ DEPENDENCIES
activerecord (>= 3.2)
bundler (~> 2.2)
rake
rspec (~> 3.12)
workflow!
RUBY VERSION

View File

@ -15,17 +15,34 @@
#
# You should have received a copy of the GNU Affero General Public License along
# with this program. If not, see <http://www.gnu.org/licenses/>.
require "aws-sdk-bedrockruntime"
module SmartSearch
EMBEDDING_VERSION = 1
EMBEDDING_VERSION = 2
CHUNK_MAX_LENGTH = 1500
class << self
def api_key
Rails.application.credentials.dig(:smart_search, :openai_api_token)
end
def bedrock_client
# for local dev, assume that we are using creds from inseng (us-west-2)
settings = YAML.safe_load(DynamicSettings.find(tree: :private)["bedrock.yml"] || "{}")
config = {
region: settings["bedrock_region"] || "us-west-2"
}
# Will load creds from vault (prod) or rails credential store (local / oss).
# Credentials stored in rails credential store in the `bedrock_creds` key
# with `aws_access_key_id` and `aws_secret_access_key` keys
config[:credentials] = Canvas::AwsCredentialProvider.new("bedrock_creds", settings["vault_credential_path"])
if config[:credentials].set?
Aws::BedrockRuntime::Client.new(config)
end
end
def smart_search_available?(context)
context&.feature_enabled?(:smart_search) && api_key.present?
context&.feature_enabled?(:smart_search) && bedrock_client.present?
end
def register_class(klass, index_scope_proc, search_scope_proc)
@ -45,16 +62,19 @@ module SmartSearch
end
end
def generate_embedding(input, version: EMBEDDING_VERSION)
case EMBEDDING_VERSION
def generate_embedding(input, query: false, version: EMBEDDING_VERSION)
case version
when 1
generate_embedding_v1(input)
when 2
generate_embedding_v2(input, query)
else
raise ArgumentError, "Unsupported embedding version #{version}"
end
end
def generate_embedding_v1(input)
# NOTE: openai does not differentiate between query and document embeddings
url = "https://api.openai.com/v1/embeddings"
headers = {
"Authorization" => "Bearer #{api_key}",
@ -72,9 +92,23 @@ module SmartSearch
response["data"].pluck("embedding")[0]
end
def generate_embedding_v2(input, query)
resp = bedrock_client.invoke_model({
content_type: "application/json",
accept: "application/json",
model_id: "cohere.embed-multilingual-v3",
body: {
texts: [input],
input_type: query ? "search_query" : "search_document"
}.to_json,
})
json = JSON.parse(resp.body.string)
json["embeddings"][0]
end
def perform_search(context, user, query, type_filter = [])
version = context.search_embedding_version || EMBEDDING_VERSION
embedding = SmartSearch.generate_embedding(query, version:)
embedding = SmartSearch.generate_embedding(query, version:, query: true)
collections = []
ActiveRecord::Base.with_pgvector do
SmartSearch.search_scopes(context, user).each do |klass, item_scope|
@ -132,7 +166,7 @@ module SmartSearch
def index_course(course)
return if course.search_embedding_version == EMBEDDING_VERSION
# TODO: investigate pipelining this after we switch to Bedrock
# TODO: investigate pipelining this
index_scopes(course).each do |scope|
scope.left_joins(:embeddings)
.group("#{scope.table_name}.id")
@ -143,9 +177,7 @@ module SmartSearch
end
course.update!(search_embedding_version: EMBEDDING_VERSION)
# TODO: implement this when we add a second embeddings version
# delay_if_production(priority: Delayed::LOW_PRIORITY).delete_old_embeddings(course)
delay_if_production(priority: Delayed::LOW_PRIORITY).delete_old_embeddings(course)
end
def delete_old_embeddings(course)

View File

@ -72,9 +72,9 @@ module SmartSearchable
def generate_embeddings
delete_embeddings
chunk_content do |chunk|
chunk_content(SmartSearch::CHUNK_MAX_LENGTH) do |chunk|
embedding = SmartSearch.generate_embedding(chunk)
embeddings.create!(embedding:)
embeddings.create!(embedding:, version: SmartSearch::EMBEDDING_VERSION)
end
end
handle_asynchronously :generate_embeddings, priority: Delayed::LOW_PRIORITY