Add embedding dimensions to config

gorse-io · Feb 19, 2025 · 04301a3 · 04301a3
1 parent e088c82
commit 04301a3
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 4 deletions.
diff --git a/config/config.go b/config/config.go
@@ -233,7 +233,8 @@ type OpenAIConfig struct {
 	BaseURL             string `mapstructure:"base_url"`
 	AuthToken           string `mapstructure:"auth_token"`
 	ChatCompletionModel string `mapstructure:"chat_completion_model"`
-	EmbeddingsModel     string `mapstructure:"embeddings_model"`
+	EmbeddingModel      string `mapstructure:"embedding_model"`
+	EmbeddingDimensions int    `mapstructure:"embedding_dimensions"`
 }
 
 func GetDefaultConfig() *Config {

diff --git a/config/config_test.go b/config/config_test.go
@@ -161,7 +161,8 @@ func TestUnmarshal(t *testing.T) {
 			assert.Equal(t, "http://localhost:11434/v1", config.OpenAI.BaseURL)
 			assert.Equal(t, "ollama", config.OpenAI.AuthToken)
 			assert.Equal(t, "qwen2.5", config.OpenAI.ChatCompletionModel)
-			assert.Equal(t, "mxbai-embed-large", config.OpenAI.EmbeddingsModel)
+			assert.Equal(t, "mxbai-embed-large", config.OpenAI.EmbeddingModel)
+			assert.Equal(t, 1024, config.OpenAI.EmbeddingDimensions)
 		})
 	}
 }

diff --git a/logics/item_to_item.go b/logics/item_to_item.go
@@ -47,6 +47,7 @@ type ItemToItemOptions struct {
 }
 
 type ItemToItem interface {
+	Timestamp() time.Time
 	Items() []*data.Item
 	Push(item *data.Item, feedback []dataset.ID)
 	PopAll(i int) []cache.Score
@@ -90,6 +91,10 @@ type baseItemToItem[T any] struct {
 	items      []*data.Item
 }
 
+func (b *baseItemToItem[T]) Timestamp() time.Time {
+	return b.timestamp
+}
+
 func (b *baseItemToItem[T]) Items() []*data.Item {
 	return b.items
 }
@@ -419,7 +424,7 @@ func (g *chatItemToItem) PopAll(i int) []cache.Score {
 		embeddings[i] = resp2.Data[0].Embedding
 	}
 	// search index
-	pq := heap.NewPriorityQueue(false)
+	pq := heap.NewPriorityQueue(true)
 	for _, embedding := range embeddings {
 		scores := g.index.SearchVector(embedding, g.n+1, true)
 		for _, score := range scores {

diff --git a/logics/item_to_item_test.go b/logics/item_to_item_test.go
@@ -212,7 +212,7 @@ func (suite *ItemToItemTestSuite) TestChat() {
 		BaseURL:             mockAI.BaseURL(),
 		AuthToken:           mockAI.AuthToken(),
 		ChatCompletionModel: "deepseek-r1",
-		EmbeddingsModel:     "text-similarity-ada-001",
+		EmbeddingModel:      "text-similarity-ada-001",
 	})
 	suite.NoError(err)
 

diff --git a/master/tasks.go b/master/tasks.go
@@ -1060,6 +1060,15 @@ func (m *Master) updateItemToItem(dataset *dataset.Dataset) error {
 						zap.String("item_id", item.ItemId), zap.Error(err))
 					continue
 				}
+				// Remove stale item-to-item recommendation
+				if err := m.CacheClient.DeleteScores(ctx, []string{cache.ItemToItem}, cache.ScoreCondition{
+					Subset: lo.ToPtr(cache.Key(itemToItemConfig.Name, item.ItemId)),
+					Before: lo.ToPtr(recommender.Timestamp()),
+				}); err != nil {
+					log.Logger().Error("failed to remove stale item-to-item recommendation",
+						zap.String("item_id", item.ItemId), zap.Error(err))
+					continue
+				}
 			}
 			span.Add(1)
 		}