mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-02 13:57:30 +08:00
feat(go): implement memory extraction task consumer (#17404)
This commit is contained in:
@@ -34,6 +34,7 @@ import (
|
||||
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/engine/types"
|
||||
"ragflow/internal/tokenizer"
|
||||
|
||||
"github.com/elastic/go-elasticsearch/v8/esapi"
|
||||
"github.com/json-iterator/go"
|
||||
@@ -169,7 +170,8 @@ func (e *elasticsearchEngine) InsertChunks(ctx context.Context, chunks []map[str
|
||||
return nil, fmt.Errorf("index name cannot be empty")
|
||||
}
|
||||
|
||||
if strings.HasPrefix(baseName, "memory_") {
|
||||
isMemoryIndex := strings.HasPrefix(baseName, "memory_")
|
||||
if isMemoryIndex {
|
||||
if err := e.ensureMemoryMessageVectorMappingsForDocs(ctx, baseName, chunks); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -178,6 +180,11 @@ func (e *elasticsearchEngine) InsertChunks(ctx context.Context, chunks []map[str
|
||||
// Build bulk request body with index operations (upsert behavior: insert if not exists, update if exists)
|
||||
var buf bytes.Buffer
|
||||
for _, doc := range chunks {
|
||||
if isMemoryIndex {
|
||||
// Memory messages arrive with logical fields; map them to
|
||||
// the storage document (tokenization included) before write.
|
||||
doc = mapMemoryMessageToESDoc(doc)
|
||||
}
|
||||
docID, _ := doc["doc_id"].(string)
|
||||
chunkID, _ := doc["id"].(string)
|
||||
if docID == "" || chunkID == "" {
|
||||
@@ -440,6 +447,11 @@ func mapMemoryMessageESUpdateFields(newValue map[string]interface{}) map[string]
|
||||
doc[mapMemoryMessageESField(k, false)] = v
|
||||
}
|
||||
}
|
||||
// Mirror Python memory/utils/es_conn.py update: writing content
|
||||
// refreshes tokenized_content_ltks before write.
|
||||
if content, ok := newValue["content"].(string); ok {
|
||||
doc["tokenized_content_ltks"] = tokenizeMemoryMessageContent(content)
|
||||
}
|
||||
return doc
|
||||
}
|
||||
|
||||
@@ -1514,6 +1526,56 @@ func sortValuesEqual(a, b []interface{}) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// mapMemoryMessageToESDoc converts a logical memory message into the
|
||||
// Elasticsearch storage document, mirroring Python
|
||||
// memory/utils/es_conn.py:map_message_to_es_fields. Elasticsearch has no
|
||||
// server-side rag tokenizer, so tokenized_content_ltks is computed here
|
||||
// before write; on Infinity the equivalent insert path stores content
|
||||
// verbatim and lets Infinity tokenize on save.
|
||||
func mapMemoryMessageToESDoc(message map[string]interface{}) map[string]interface{} {
|
||||
doc := make(map[string]interface{}, len(message)+2)
|
||||
for k, v := range message {
|
||||
switch k {
|
||||
case "message_type":
|
||||
doc["message_type_kwd"] = v
|
||||
case "status":
|
||||
doc["status_int"] = memoryMessageStatusInt(v)
|
||||
case "content":
|
||||
content, _ := v.(string)
|
||||
doc["content_ltks"] = content
|
||||
doc["tokenized_content_ltks"] = tokenizeMemoryMessageContent(content)
|
||||
default:
|
||||
doc[k] = v
|
||||
}
|
||||
}
|
||||
return doc
|
||||
}
|
||||
|
||||
// tokenizeMemoryMessageContent runs the rag tokenizer pipeline
|
||||
// (tokenize + fine-grained tokenize) over message content. On tokenizer
|
||||
// failure it degrades to the untokenized text so the document stays
|
||||
// searchable under the whitespace analyzer.
|
||||
func tokenizeMemoryMessageContent(content string) string {
|
||||
tokenized, err := tokenizer.Tokenize(content)
|
||||
if err != nil {
|
||||
common.Warn("memory message tokenize failed, storing raw content", zap.Error(err))
|
||||
return content
|
||||
}
|
||||
fine, err := tokenizer.FineGrainedTokenize(tokenized)
|
||||
if err != nil {
|
||||
common.Warn("memory message fine-grained tokenize failed, storing coarse tokens", zap.Error(err))
|
||||
return tokenized
|
||||
}
|
||||
return fine
|
||||
}
|
||||
|
||||
func memoryMessageStatusInt(value interface{}) int {
|
||||
if memoryMessageStatusBool(value) {
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func mapMemoryMessageESField(field string, useTokenizedContent bool) string {
|
||||
name := field
|
||||
boost := ""
|
||||
@@ -2659,10 +2721,6 @@ func getMemoryMessageMapping(vectorSize int) map[string]interface{} {
|
||||
"status_int": map[string]interface{}{
|
||||
"type": "integer",
|
||||
},
|
||||
"content": map[string]interface{}{
|
||||
"type": "text",
|
||||
"index": false,
|
||||
},
|
||||
"content_ltks": map[string]interface{}{
|
||||
"type": "text",
|
||||
"analyzer": "whitespace",
|
||||
|
||||
@@ -281,3 +281,67 @@ func assertEqual(t *testing.T, got, want interface{}) {
|
||||
t.Fatalf("got %#v, want %#v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMapMemoryMessageToESDoc: logical message fields are mapped to the
|
||||
// storage document at insert time — message_type/status collapse to
|
||||
// their storage counterparts and content is tokenized before write,
|
||||
// mirroring Python memory/utils/es_conn.py:map_message_to_es_fields.
|
||||
func TestMapMemoryMessageToESDoc(t *testing.T) {
|
||||
doc := mapMemoryMessageToESDoc(map[string]interface{}{
|
||||
"id": "mem-1_42",
|
||||
"doc_id": "mem-1",
|
||||
"message_id": int64(42),
|
||||
"message_type": "raw",
|
||||
"source_id": 0,
|
||||
"memory_id": "mem-1",
|
||||
"agent_id": "a1",
|
||||
"content": "User Input: hello world\nAgent Response: hi there",
|
||||
"status": true,
|
||||
"forget_at": nil,
|
||||
})
|
||||
|
||||
assertEqual(t, doc["message_type_kwd"], "raw")
|
||||
if _, ok := doc["message_type"]; ok {
|
||||
t.Fatalf("message_type must collapse into message_type_kwd, got %#v", doc)
|
||||
}
|
||||
assertEqual(t, doc["status_int"], 1)
|
||||
if _, ok := doc["status"]; ok {
|
||||
t.Fatalf("status must collapse into status_int, got %#v", doc)
|
||||
}
|
||||
|
||||
raw := "User Input: hello world\nAgent Response: hi there"
|
||||
assertEqual(t, doc["content_ltks"], raw)
|
||||
tokenized, ok := doc["tokenized_content_ltks"].(string)
|
||||
if !ok || tokenized == "" {
|
||||
t.Fatalf("tokenized_content_ltks = %#v, want non-empty tokenized string", doc["tokenized_content_ltks"])
|
||||
}
|
||||
if _, ok := doc["content"]; ok {
|
||||
t.Fatalf("content must not be stored verbatim, got %#v", doc)
|
||||
}
|
||||
|
||||
// Untouched fields pass through.
|
||||
assertEqual(t, doc["id"], "mem-1_42")
|
||||
assertEqual(t, doc["doc_id"], "mem-1")
|
||||
assertEqual(t, doc["message_id"], int64(42))
|
||||
assertEqual(t, doc["agent_id"], "a1")
|
||||
if v, ok := doc["forget_at"]; !ok || v != nil {
|
||||
t.Fatalf("forget_at = %#v, want explicit nil", v)
|
||||
}
|
||||
}
|
||||
|
||||
// TestMapMemoryMessageESUpdateFieldsRefreshesTokens: writing content via
|
||||
// the update path recomputes tokenized_content_ltks before write, same
|
||||
// as the Python update path.
|
||||
func TestMapMemoryMessageESUpdateFieldsRefreshesTokens(t *testing.T) {
|
||||
doc := mapMemoryMessageESUpdateFields(map[string]interface{}{
|
||||
"content": "fresh content",
|
||||
"status": 0,
|
||||
})
|
||||
|
||||
assertEqual(t, doc["content_ltks"], "fresh content")
|
||||
tokenized, ok := doc["tokenized_content_ltks"].(string)
|
||||
if !ok || tokenized == "" {
|
||||
t.Fatalf("tokenized_content_ltks = %#v, want refreshed tokenized string", doc["tokenized_content_ltks"])
|
||||
}
|
||||
assertEqual(t, doc["status_int"], 0)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user