support JSON output from chat LLM

config/config.go

@@ -220,7 +220,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 {

config/config.toml

@@ -158,11 +158,31 @@ filter = "(now() - item.Timestamp).Hours() < 168"
 # #   embedding: recommend by Euclidean distance of embeddings.
 # #   tags: recommend by number of common tags.
 # #   users: recommend by number of common users.
+# #   chat: recommend by chat completion model.
 # type = "embedding"
 
 # # The column of the item embeddings. Leave blank if type is "users".
 # column = "item.Labels.embedding"
 
+# [[recommend.item-to-item]]
+
+# # The name of the item-to-item recommender.
+# name = "chat_recommend"
+
+# # The type of the item-to-item recommender.
+# type = "chat"
+
+# # The column of the item embeddings. Leave blank if type is "users".
+# column = "item.Labels.embedding"
+
+# # The prompt for the chat completion model.
+# prompt = """
+# This is the description of GitHub repository https://github.com/{{ item.ItemId | replace(':','/') }}:
+# {{ item.Comment }}
+# Please find some similar repositores on GitHub and provide a brief description for each of them.
+# The output should be a JSON array.
+# """
+
 [recommend.user_neighbors]
 
 # The type of neighbors for users. There are three types:
@@ -321,5 +341,8 @@ auth_token = "ollama"
 # Name of chat completion model.
 chat_completion_model = "qwen2.5"
 
-# Name of embeddings model.
-embeddings_model = "mxbai-embed-large"
+# Name of embedding model.
+embedding_model = "mxbai-embed-large"
+
+# Dimensions of embedding vectors.
+embedding_dimensions = 1024

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)
 		})
 	}
 }

go.mod

@@ -52,6 +52,7 @@ require (
 	github.com/steinfletcher/apitest v1.5.17
 	github.com/stretchr/testify v1.10.0
 	github.com/thoas/go-funk v0.9.2
+	github.com/yuin/goldmark v1.7.8
 	go.mongodb.org/mongo-driver v1.16.1
 	go.opentelemetry.io/contrib/instrumentation/github.com/emicklei/go-restful/otelrestful v0.36.4
 	go.opentelemetry.io/contrib/instrumentation/go.mongodb.org/mongo-driver/mongo/otelmongo v0.55.0

go.sum

@@ -654,6 +654,8 @@ github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9dec
 github.com/yuin/goldmark v1.3.5/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 github.com/yuin/goldmark v1.4.1/go.mod h1:mwnBkeHKe2W/ZEtQ+71ViKU8L12m81fl3OWwC1Zlc8k=
 github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
+github.com/yuin/goldmark v1.7.8 h1:iERMLn0/QJeHFhxSt3p6PeN9mGnvIKSpG9YYorDMnic=
+github.com/yuin/goldmark v1.7.8/go.mod h1:uzxRWxtg69N339t3louHJ7+O03ezfj6PlliRlaOzY1E=
 github.com/zenazn/goji v0.9.0/go.mod h1:7S9M489iMyHBNxwZnk9/EHS098H4/F6TATF2mIxtB1Q=
 go.mongodb.org/mongo-driver v1.16.1 h1:rIVLL3q0IHM39dvE+z2ulZLp9ENZKThVfuvN/IiN4l8=
 go.mongodb.org/mongo-driver v1.16.1/go.mod h1:oB6AhJQvFQL4LEHyXi6aJzQJtBiTQHiAd83l0GdFaiw=

logics/item_to_item.go

@@ -16,8 +16,8 @@ package logics
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
-	"fmt"
 	"sort"
 	"strings"
 	"time"
@@ -30,7 +30,11 @@ import (
 	"github.com/nikolalohinski/gonja/v2/exec"
 	"github.com/samber/lo"
 	"github.com/sashabaranov/go-openai"
+	"github.com/yuin/goldmark"
+	"github.com/yuin/goldmark/ast"
+	"github.com/yuin/goldmark/text"
 	"github.com/zhenghaoz/gorse/base/floats"
+	"github.com/zhenghaoz/gorse/base/heap"
 	"github.com/zhenghaoz/gorse/base/log"
 	"github.com/zhenghaoz/gorse/common/ann"
 	"github.com/zhenghaoz/gorse/config"
@@ -47,6 +51,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 +95,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
 }
@@ -353,10 +362,11 @@ func flatten(o any, tSet mapset.Set[dataset.ID]) {
 
 type chatItemToItem struct {
 	*embeddingItemToItem
-	template       *exec.Template
-	client         *openai.Client
-	chatModel      string
-	embeddingModel string
+	template            *exec.Template
+	client              *openai.Client
+	chatCompletionModel string
+	embeddingModel      string
+	embeddingDimensions int
 }
 
 func newChatItemToItem(cfg config.ItemToItemConfig, n int, timestamp time.Time, openaiConfig config.OpenAIConfig) (*chatItemToItem, error) {
@@ -377,12 +387,27 @@ func newChatItemToItem(cfg config.ItemToItemConfig, n int, timestamp time.Time,
 		embeddingItemToItem: embedding,
 		template:            template,
 		client:              openai.NewClientWithConfig(clientConfig),
-		chatModel:           openaiConfig.ChatCompletionModel,
-		embeddingModel:      openaiConfig.EmbeddingsModel,
+		chatCompletionModel: openaiConfig.ChatCompletionModel,
+		embeddingModel:      openaiConfig.EmbeddingModel,
+		embeddingDimensions: openaiConfig.EmbeddingDimensions,
 	}, nil
 }
 
 func (g *chatItemToItem) PopAll(i int) []cache.Score {
+	// evaluate column expression and get embedding vector
+	result, err := expr.Run(g.columnFunc, map[string]any{
+		"item": g.items[i],
+	})
+	if err != nil {
+		log.Logger().Error("failed to evaluate column expression",
+			zap.Any("item", g.items[i]), zap.Error(err))
+		return nil
+	}
+	embedding0, ok := result.([]float32)
+	if !ok {
+		log.Logger().Error("invalid column type", zap.Any("column", result))
+		return nil
+	}
 	// render template
 	var buf strings.Builder
 	ctx := exec.NewContext(map[string]any{
@@ -392,10 +417,9 @@ func (g *chatItemToItem) PopAll(i int) []cache.Score {
 		log.Logger().Error("failed to execute template", zap.Error(err))
 		return nil
 	}
-	fmt.Println(buf.String())
 	// chat completion
 	resp, err := g.client.CreateChatCompletion(context.Background(), openai.ChatCompletionRequest{
-		Model: g.chatModel,
+		Model: g.chatCompletionModel,
 		Messages: []openai.ChatCompletionMessage{{
 			Role:    openai.ChatMessageRoleUser,
 			Content: buf.String(),
@@ -405,29 +429,51 @@ func (g *chatItemToItem) PopAll(i int) []cache.Score {
 		log.Logger().Error("failed to chat completion", zap.Error(err))
 		return nil
 	}
-	message := stripThink(resp.Choices[0].Message.Content)
+	messages := parseMessage(resp.Choices[0].Message.Content)
+	log.Logger().Debug("chat based item-to-item recommendation",
+		zap.String("prompt", buf.String()), zap.Strings("response", messages))
 	// message embedding
-	resp2, err := g.client.CreateEmbeddings(context.Background(), openai.EmbeddingRequest{
-		Input: message,
-		Model: openai.EmbeddingModel(g.embeddingModel),
-	})
-	if err != nil {
-		log.Logger().Error("failed to create embeddings", zap.Error(err))
-		return nil
+	embeddings := make([][]float32, len(messages))
+	for i, message := range messages {
+		resp, err := g.client.CreateEmbeddings(context.Background(), openai.EmbeddingRequest{
+			Input:      message,
+			Model:      openai.EmbeddingModel(g.embeddingModel),
+			Dimensions: g.embeddingDimensions,
+		})
+		if err != nil {
+			log.Logger().Error("failed to create embeddings", zap.Error(err))
+			return nil
+		}
+		embeddings[i] = resp.Data[0].Embedding
 	}
-	embedding := resp2.Data[0].Embedding
 	// search index
-	scores := g.index.SearchVector(embedding, g.n+1, true)
-	return lo.Map(scores, func(v lo.Tuple2[int, float32], _ int) cache.Score {
-		return cache.Score{
-			Id:         g.items[v.A].ItemId,
-			Categories: g.items[v.A].Categories,
-			Score:      -float64(v.B),
+	pq := heap.NewPriorityQueue(true)
+	for _, embedding := range embeddings {
+		score0 := floats.Euclidean(embedding, embedding0)
+		scores := g.index.SearchVector(embedding, g.n+1, true)
+		for _, score := range scores {
+			if score.A != i {
+				pq.Push(int32(score.A), score.B*score0)
+				if pq.Len() > g.n {
+					pq.Pop()
+				}
+			}
+		}
+	}
+	scores := make([]cache.Score, pq.Len())
+	for i := 9; i >= 0; i-- {
+		id, score := pq.Pop()
+		scores[i] = cache.Score{
+			Id:         g.items[id].ItemId,
+			Categories: g.items[id].Categories,
+			Score:      -float64(score),
 			Timestamp:  g.timestamp,
 		}
-	})
+	}
+	return scores
 }
 
+// stripThink strips the <think> tag from the message.
 func stripThink(s string) string {
 	if len(s) < 7 || s[:7] != "<think>" {
 		return s
@@ -438,3 +484,39 @@ func stripThink(s string) string {
 	}
 	return s[end+8:]
 }
+
+// parseMessage parse message from chat completion response.
+// If there is any JSON in the message, it returns the JSON.
+// Otherwise, it returns the message.
+func parseMessage(message string) []string {
+	source := []byte(stripThink(message))
+	root := goldmark.DefaultParser().Parse(text.NewReader(source))
+	for n := root.FirstChild(); n != nil; n = n.NextSibling() {
+		if n.Kind() != ast.KindFencedCodeBlock {
+			continue
+		}
+		if codeBlock, ok := n.(*ast.FencedCodeBlock); ok {
+			if string(codeBlock.Language(source)) == "json" {
+				bytes := codeBlock.Text(source)
+				if bytes[0] == '[' {
+					var temp []any
+					err := json.Unmarshal(bytes, &temp)
+					if err != nil {
+						return []string{string(bytes)}
+					}
+					var result []string
+					for _, v := range temp {
+						bytes, err := json.Marshal(v)
+						if err != nil {
+							return []string{string(bytes)}
+						}
+						result = append(result, string(bytes))
+					}
+					return result
+				}
+				return []string{string(bytes)}
+			}
+		}
+	}
+	return []string{string(source)}
+}

logics/item_to_item_test.go

@@ -19,6 +19,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/suite"
 	"github.com/zhenghaoz/gorse/base/floats"
 	"github.com/zhenghaoz/gorse/common/mock"
@@ -212,13 +213,13 @@ 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)
 
 	for i := 0; i < 100; i++ {
 		embedding := mock.Hash("Please generate similar items for item_0.")
-		floats.AddConst(embedding, float32(i))
+		floats.AddConst(embedding, float32(i+1))
 		item2item.Push(&data.Item{
 			ItemId: strconv.Itoa(i),
 			Labels: map[string]any{
@@ -238,3 +239,25 @@ func (suite *ItemToItemTestSuite) TestChat() {
 func TestItemToItem(t *testing.T) {
 	suite.Run(t, new(ItemToItemTestSuite))
 }
+
+func TestParseMessage(t *testing.T) {
+	// parse JSON object
+	message := "```json\n{\"a\": 1, \"b\": 2}\n```"
+	contents := parseMessage(message)
+	assert.Equal(t, []string{"{\"a\": 1, \"b\": 2}\n"}, contents)
+
+	// parse JSON array
+	message = "```json\n[1, 2]\n```"
+	contents = parseMessage(message)
+	assert.Equal(t, []string{"1", "2"}, contents)
+
+	// parse text
+	message = "Hello, world!"
+	contents = parseMessage(message)
+	assert.Equal(t, []string{"Hello, world!"}, contents)
+
+	// strip think
+	message = "<think>hello</think>World!"
+	content := stripThink(message)
+	assert.Equal(t, "World!", content)
+}

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)
 		}