Compare commits

...
Author SHA1 Message Date
Sam Saffron 6e6ced4554 Work in progress... post embeddings and version
Feel free to continue any work here, will pick it up again next week
2023-05-12 15:28:12 +10:00
4 changed files with 44 additions and 13 deletions
+1
View File
@@ -47,6 +47,7 @@ en:
ai_embeddings_generate_for_pms: "Generate embeddings for personal messages." ai_embeddings_generate_for_pms: "Generate embeddings for personal messages."
ai_embeddings_semantic_related_topics_enabled: "Use Semantic Search for related topics." ai_embeddings_semantic_related_topics_enabled: "Use Semantic Search for related topics."
ai_embeddings_semantic_related_topics: "Maximum number of topics to show in related topic section." ai_embeddings_semantic_related_topics: "Maximum number of topics to show in related topic section."
ai_embeddings_semantic_related_topics_include_closed_topics: "Include closed topics in related topics."
ai_embeddings_pg_connection_string: "PostgreSQL connection string for the embeddings module. Needs pgvector extension enabled and a series of tables created. See docs for more info." ai_embeddings_pg_connection_string: "PostgreSQL connection string for the embeddings module. Needs pgvector extension enabled and a series of tables created. See docs for more info."
ai_embeddings_semantic_search_model: "Model to use for semantic search." ai_embeddings_semantic_search_model: "Model to use for semantic search."
ai_embeddings_semantic_search_enabled: "Enable full-page semantic search." ai_embeddings_semantic_search_enabled: "Enable full-page semantic search."
+3 -1
View File
@@ -156,7 +156,9 @@ plugins:
ai_embeddings_semantic_related_topics_enabled: false ai_embeddings_semantic_related_topics_enabled: false
ai_embeddings_semantic_related_topics: 5 ai_embeddings_semantic_related_topics: 5
ai_embeddings_semantic_related_include_closed_topics: true ai_embeddings_semantic_related_include_closed_topics: true
ai_embeddings_pg_connection_string: "" ai_embeddings_pg_connection_string:
default: ""
secret: true
ai_embeddings_semantic_search_enabled: ai_embeddings_semantic_search_enabled:
default: false default: false
client: true client: true
+27 -6
View File
@@ -3,7 +3,9 @@
module DiscourseAi module DiscourseAi
module Embeddings module Embeddings
class Topic class Topic
def generate_and_store_embeddings_for(topic) VERSION = 1
def generate_and_store_embeddings_for(topic, include_posts: true)
return unless SiteSetting.ai_embeddings_enabled return unless SiteSetting.ai_embeddings_enabled
return if topic.blank? || topic.first_post.blank? return if topic.blank? || topic.first_post.blank?
@@ -13,6 +15,18 @@ module DiscourseAi
enabled_models.each do |model| enabled_models.each do |model|
embedding = model.generate_embedding(topic.first_post.raw) embedding = model.generate_embedding(topic.first_post.raw)
persist_embedding(topic, model, embedding) if embedding persist_embedding(topic, model, embedding) if embedding
if include_posts
persist_embedding(topic.first_post, model, embedding) if embedding
topic
.posts
.where("post_number > 1 AND post_type = 1")
.each do |post|
embedding = model.generate_embedding(post.raw)
persist_embedding(post, model, embedding) if embedding
end
end
end end
end end
@@ -70,13 +84,20 @@ module DiscourseAi
private private
def persist_embedding(topic, model, embedding) def persist_embedding(topic_or_post, model, embedding)
DiscourseAi::Database::Connection.db.exec(<<~SQL, topic_id: topic.id, embedding: embedding) table = topic_or_post.is_a?(Topic) ? "topic" : "post"
INSERT INTO topic_embeddings_#{model.name.underscore} (topic_id, embedding)
VALUES (:topic_id, '[:embedding]') DiscourseAi::Database::Connection.db.exec(
ON CONFLICT (topic_id) <<~SQL,
INSERT INTO #{table}_embeddings_#{model.name.underscore} (#{table}_id, embedding, version)
VALUES (:id, '[:embedding]', :version)
ON CONFLICT (#{table}_id)
DO UPDATE SET embedding = '[:embedding]' DO UPDATE SET embedding = '[:embedding]'
SQL SQL
id: topic_or_post.id,
embedding: embedding,
version: VERSION,
)
end end
end end
end end
+13 -6
View File
@@ -7,17 +7,24 @@ task "ai:embeddings:create_table" => [:environment] do
SQL SQL
DiscourseAi::Embeddings::Model.enabled_models.each do |model| DiscourseAi::Embeddings::Model.enabled_models.each do |model|
DiscourseAi::Database::Connection.db.exec(<<~SQL) %w[topic post].each do |table|
CREATE TABLE IF NOT EXISTS topic_embeddings_#{model.name.underscore} ( table_name = "#{table}_embeddings_#{model.name.underscore}"
topic_id bigint PRIMARY KEY, DiscourseAi::Database::Connection.db.exec(<<~SQL)
embedding vector(#{model.dimensions}) CREATE TABLE IF NOT EXISTS #{table_name} (
); #{table}_id bigint PRIMARY KEY,
embedding vector(#{model.dimensions}),
version smallint
)
SQL SQL
DiscourseAi::Database::Connection.db.exec(<<~SQL)
ALTER TABLE #{table_name} ADD COLUMN IF NOT EXISTS version smallint
SQL
end
end end
end end
desc "Backfill embeddings for all topics" desc "Backfill embeddings for all topics"
task "ai:embeddings:backfill", [:start_topic] => [:environment] do task "ai:embeddings:backfill", [:start_topic] => [:environment] do |_, args|
public_categories = Category.where(read_restricted: false).pluck(:id) public_categories = Category.where(read_restricted: false).pluck(:id)
topic_embeddings = DiscourseAi::Embeddings::Topic.new topic_embeddings = DiscourseAi::Embeddings::Topic.new
Topic Topic