redis4rag/chathistory_test.go at dev · bitsark/redis4rag · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
package redis4rag

import (
	"cmp"
	"context"
	"encoding/json"
	"os"
	"testing"
	"time"

	"github.com/redis/go-redis/v9"
	should "github.com/stretchr/testify/assert"
)

func TestChatHistoryBasic(t *testing.T) {
	indexname := "idx:test_chat_history_basic"
	docprefix := "doc:test_chat_history_basic"
	redisCli := redis.NewClient(&redis.Options{
		Addr:     cmp.Or(os.Getenv("REDIS_URL"), "localhost:6379"),
		Protocol: 2,
	})
	redisCli.FTDropIndexWithArgs(context.Background(), indexname, &redis.FTDropIndexOptions{DeleteDocs: true})

	err := CreateIndex(redisCli, ChatHistorySchema, indexname, []interface{}{docprefix})
	should.Nil(t, err)
	t.Logf("index %s created", indexname)

	//TODO
	chatHistory := &ChatHistory{
		indexName: indexname,
		docPrefix: docprefix,
		redisCli:  redisCli,
	}

	t1 := time.Now()
	time.Sleep(1 * time.Millisecond)
	err = chatHistory.Add(context.Background(), &ChatMessage{
		Typ:       "user",
		UserId:    "user_id",
		SessionId: "session_id",
		Content:   "咱俩谁跟谁呀。",
		Timestamp: time.Now().UnixMilli(),
	}, localEmbedder.Embedding)
	should.Nil(t, err)

	t2 := time.Now()
	time.Sleep(1 * time.Millisecond)

	err = chatHistory.Add(context.Background(), &ChatMessage{
		Typ:       "user",
		UserId:    "user_id",
		SessionId: "session_id_2",
		Content:   "我俩谁跟谁呀。",
		Timestamp: time.Now().UnixMilli(),
	}, localEmbedder.Embedding)
	should.Nil(t, err)

	t3 := time.Now()
	// list
	{
		msgs, err := chatHistory.ListBySessionId(context.Background(), t1.UnixMilli(), "session_id")
		should.Nil(t, err)
		buf, err := json.Marshal(msgs)
		should.Nil(t, err)
		t.Logf("%s\n", string(buf))
	}
	{
		msgs, err := chatHistory.ListBySessionId(context.Background(), t2.UnixMilli(), "session_id")
		should.Nil(t, err)
		buf, err := json.Marshal(msgs)
		should.Nil(t, err)
		t.Logf("%s\n", string(buf))
	}
	{
		msgs, err := chatHistory.ListBySessionId(context.Background(), t3.UnixMilli(), "session_id")
		should.Nil(t, err)
		buf, err := json.Marshal(msgs)
		should.Nil(t, err)
		t.Logf("%s\n", string(buf))
	}

	//search
	{
		msgs, err := chatHistory.SearchWithUserId(context.Background(),
			t1.UnixMilli(), "user_id", "咱俩关系不错呀。", localEmbedder.Embedding)
		should.Nil(t, err)
		buf, err := json.Marshal(msgs)
		should.Nil(t, err)
		t.Logf("%s\n", string(buf))
	}

	{
		err := chatHistory.DeleteBySessionId(context.Background(), "session_id")
		should.Nil(t, err)
		msgs, err := chatHistory.SearchWithUserId(context.Background(),
			t1.UnixMilli(), "user_id", "咱俩关系不错呀。", localEmbedder.Embedding)
		should.Nil(t, err)
		buf, err := json.Marshal(msgs)
		should.Nil(t, err)
		t.Logf("%s\n", string(buf))

	}

	err = redisCli.FTDropIndexWithArgs(context.Background(), indexname, &redis.FTDropIndexOptions{DeleteDocs: true}).Err()
	should.Nil(t, err)
	t.Logf("index %s dropped", indexname)
}