feat(go): implement memory extraction task consumer (#17404)

2026-08-01 13:33:48 +08:00 · 2026-07-31 11:23:58 +08:00
parent 569a29500f
commit ed05f32fee
8 changed files with 592 additions and 104 deletions
--- a/internal/dao/task.go
+++ b/internal/dao/task.go
@@ -54,6 +54,14 @@ func (dao *TaskDAO) GetByID(ctx context.Context, db *gorm.DB, id string) (*entit
 	return &task, nil
 }

+// UpdateProgress sets the progress and progress message of a task
+func (dao *TaskDAO) UpdateProgress(id string, progress float64, progressMsg string) error {
+	return DB.Model(&entity.Task{}).Where("id = ?", id).Updates(map[string]any{
+		"progress":     progress,
+		"progress_msg": progressMsg,
+	}).Error
+}
+
 // DeleteIngestionTasksByDocIDs deletes ingestion tasks by document IDs (hard delete)
 func (dao *TaskDAO) DeleteIngestionTasksByDocIDs(ctx context.Context, db *gorm.DB, docIDs []string) (int64, error) {
 	if len(docIDs) == 0 {
--- a/internal/engine/elasticsearch/chunk.go
+++ b/internal/engine/elasticsearch/chunk.go
@@ -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",
--- a/internal/engine/elasticsearch/chunk_helpers_test.go
+++ b/internal/engine/elasticsearch/chunk_helpers_test.go
@@ -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)
+}
--- a/internal/service/memory.go
+++ b/internal/service/memory.go
@@ -189,6 +189,25 @@ func (PromptAssembler) AssembleSystemPrompt(memoryTypes []string) string {
 	return fullPrompt
 }

+// baseUserPromptTemplate is the default user prompt for memory extraction,
+// matching Python PromptAssembler.BASE_USER_PROMPT.
+const baseUserPromptTemplate = `
+**CONVERSATION:**
+{conversation}
+
+**CONVERSATION TIME:** {conversation_time}
+**CURRENT TIME:** {current_time}
+`
+
+// AssembleUserPrompt renders the default extraction user prompt with the
+// conversation content and timestamps.
+func (PromptAssembler) AssembleUserPrompt(conversation, conversationTime, currentTime string) string {
+	prompt := strings.Replace(baseUserPromptTemplate, "{conversation}", conversation, 1)
+	prompt = strings.Replace(prompt, "{conversation_time}", conversationTime, 1)
+	prompt = strings.Replace(prompt, "{current_time}", currentTime, 1)
+	return prompt
+}
+
 // getTypesToExtract filters out "raw" type and returns valid memory types
 //
 // Parameters:
--- a/internal/service/memory_extractor.go
+++ b/internal/service/memory_extractor.go
@@ -0,0 +1,353 @@
+//
+//  Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+// memory_extractor.go — async memory extraction worker.
+//
+// Port of the Python task-executor memory path:
+//
+//	rag/svr/task_executor.py:handle_task (task_type == "memory")
+//	api/db/joint_services/memory_message_service.py:
+//	    handle_save_to_memory_task / save_extracted_to_memory_only / extract_by_llm
+//
+// QueueSaveToMemoryTask persists the raw message and enqueues a
+// task_type="memory" message on the Redis stream te.<priority>.common.
+// StartTaskConsumer drains that stream (consumer group
+// "rag_flow_svr_task_broker", same as the Python executor), runs LLM
+// extraction for the non-raw memory types configured on the memory, and
+// persists the extracted messages with source_id pointing at the raw
+// message so listMemoryMessages can aggregate them under `extract`.
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"strings"
+	"time"
+
+	"ragflow/internal/common"
+	"ragflow/internal/dao"
+	redisengine "ragflow/internal/engine/redis"
+	"ragflow/internal/entity"
+	models "ragflow/internal/entity/models"
+	"ragflow/internal/utility"
+
+	"go.uber.org/zap"
+)
+
+// memoryTaskConsumerGroup matches Python common.constants.SVR_CONSUMER_GROUP_NAME
+// so Go and Python executors never double-process the same stream entry.
+const memoryTaskConsumerGroup = "rag_flow_svr_task_broker"
+
+// memoryTimeLayout is the storage format for valid_at / invalid_at,
+// matching timestamp_to_date on the Python side.
+const memoryTimeLayout = "2006-01-02 15:04:05"
+
+// extractedMemory is one LLM-extracted memory item ready for persistence.
+type extractedMemory struct {
+	MessageType string
+	Content     string
+	ValidAt     string
+	InvalidAt   string // empty means still valid
+}
+
+// StartTaskConsumer is the long-running loop that drains memory
+// extraction tasks from the Redis stream. It returns when ctx is
+// cancelled. Per-message failures are logged and acked so one bad
+// message cannot stall the stream.
+func (s *MemoryMessageService) StartTaskConsumer(ctx context.Context) {
+	redisClient := redisengine.Get()
+	if redisClient == nil {
+		common.Error("memory task consumer: Redis is not available", nil)
+		return
+	}
+	consumerName := fmt.Sprintf("go_memory_extractor_%s", utility.GenerateUUID())
+	queueName := memoryTaskQueueName(0)
+	common.Info(fmt.Sprintf("Memory task consumer %s started on queue %s", consumerName, queueName))
+
+	for {
+		if ctx.Err() != nil {
+			return
+		}
+		msg, err := redisClient.QueueConsumer(queueName, memoryTaskConsumerGroup, consumerName, ">")
+		if err != nil {
+			common.Error("memory task consumer: consume error", err)
+			select {
+			case <-time.After(time.Second):
+			case <-ctx.Done():
+				return
+			}
+			continue
+		}
+		if msg == nil {
+			continue
+		}
+		payload := msg.GetMessage()
+		if taskType, _ := payload["task_type"].(string); taskType != "memory" {
+			common.Warn(fmt.Sprintf("memory task consumer: skip task_type %q", taskType))
+			msg.Ack()
+			continue
+		}
+		if err := s.HandleSaveToMemoryTask(ctx, payload); err != nil {
+			common.Error("memory task consumer: handle task failed", err)
+		}
+		msg.Ack()
+	}
+}
+
+// HandleSaveToMemoryTask processes one queued memory task. Mirrors
+// Python handle_save_to_memory_task: validate the task row, then
+// extract + persist, settling task progress on the way out.
+func (s *MemoryMessageService) HandleSaveToMemoryTask(ctx context.Context, payload map[string]any) error {
+	taskID, _ := payload["id"].(string)
+	if taskID == "" {
+		taskID, _ = payload["task_id"].(string)
+	}
+	memoryID, _ := payload["memory_id"].(string)
+	sourceID := payloadInt64(payload["source_id"])
+	msgDict, _ := payload["message_dict"].(map[string]any)
+	msg := MemoryMessage{
+		UserID:        payloadString(msgDict["user_id"]),
+		AgentID:       payloadString(msgDict["agent_id"]),
+		SessionID:     payloadString(msgDict["session_id"]),
+		UserInput:     payloadString(msgDict["user_input"]),
+		AgentResponse: payloadString(msgDict["agent_response"]),
+	}
+
+	task, err := s.taskDAO.GetByID(ctx, dao.DB, taskID)
+	if err != nil {
+		return fmt.Errorf("memory: task %s is not found", taskID)
+	}
+	if task.Progress == -1 {
+		return fmt.Errorf("memory: task %s is already failed", taskID)
+	}
+
+	if err := s.saveExtractedToMemory(ctx, memoryID, msg, sourceID, taskID); err != nil {
+		s.updateTaskProgress(taskID, -1, err.Error())
+		return err
+	}
+	return nil
+}
+
+// saveExtractedToMemory mirrors Python save_extracted_to_memory_only:
+// skip raw-only memories, run LLM extraction, embed and persist the
+// extracted messages under the raw message's id.
+func (s *MemoryMessageService) saveExtractedToMemory(ctx context.Context, memoryID string, msg MemoryMessage, sourceID int64, taskID string) error {
+	mem, err := s.memories.getMemoryConfig(ctx, memoryID)
+	if err != nil {
+		return err
+	}
+	memoryTypes := mem.MemoryType
+	if len(memoryTypes) == 0 {
+		memoryTypes = dao.GetMemoryTypeHuman(mem.Memory.MemoryType)
+	}
+	extractTypes := getTypesToExtract(memoryTypes)
+	if len(extractTypes) == 0 {
+		s.updateTaskProgress(taskID, 1.0, fmt.Sprintf("Memory '%s' don't need to extract.", memoryID))
+		return nil
+	}
+
+	extracted, err := s.extractByLLM(ctx, mem, extractTypes, msg, taskID)
+	if err != nil {
+		return err
+	}
+	if len(extracted) == 0 {
+		s.updateTaskProgress(taskID, 1.0, "No memory extracted from raw message.")
+		return nil
+	}
+	s.updateTaskProgress(taskID, 0.5, fmt.Sprintf("Extracted %d messages from raw dialogue.", len(extracted)))
+
+	now := time.Now().UTC()
+	messages := make([]map[string]any, 0, len(extracted))
+	for _, item := range extracted {
+		messages = append(messages, buildExtractedMessage(generateRawMessageID(), sourceID, memoryID, msg, item, now))
+	}
+	if err := s.embedAndSaveMessages(ctx, mem, messages); err != nil {
+		return err
+	}
+	s.updateTaskProgress(taskID, 1.0, "Message saved successfully.")
+	return nil
+}
+
+// extractByLLM mirrors Python extract_by_llm: build the system/user
+// prompts from the memory config, chat with the configured model, and
+// parse the JSON result into per-type extracted items.
+func (s *MemoryMessageService) extractByLLM(ctx context.Context, mem *CreateMemoryResponse, extractTypes []string, msg MemoryMessage, taskID string) ([]extractedMemory, error) {
+	systemPrompt := ""
+	if mem.SystemPrompt != nil {
+		systemPrompt = *mem.SystemPrompt
+	}
+	if strings.TrimSpace(systemPrompt) == "" {
+		systemPrompt = PromptAssembler{}.AssembleSystemPrompt(extractTypes)
+	}
+
+	conversation := fmt.Sprintf("User Input: %s\nAgent Response: %s", msg.UserInput, msg.AgentResponse)
+	now := time.Now().UTC().Format(memoryTimeLayout)
+	messages := []models.Message{{Role: "system", Content: systemPrompt}}
+	if mem.UserPrompt != nil && strings.TrimSpace(*mem.UserPrompt) != "" {
+		messages = append(messages,
+			models.Message{Role: "user", Content: *mem.UserPrompt},
+			models.Message{Role: "user", Content: fmt.Sprintf("Conversation: %s\nConversation Time: %s\nCurrent Time: %s", conversation, now, now)},
+		)
+	} else {
+		messages = append(messages, models.Message{Role: "user", Content: PromptAssembler{}.AssembleUserPrompt(conversation, now, now)})
+	}
+
+	// Python prefers tenant_llm_id and falls back to llm_id;
+	// ResolveModelConfig accepts both tenant-model ids and model names.
+	llmRef := mem.LLMID
+	if mem.TenantLLMID != nil && *mem.TenantLLMID != "" {
+		llmRef = *mem.TenantLLMID
+	}
+	driver, modelName, apiConfig, _, err := NewModelProviderService().ResolveModelConfig(ctx, mem.TenantID, entity.ModelTypeChat, llmRef)
+	if err != nil {
+		return nil, fmt.Errorf("resolve chat model: %w", err)
+	}
+	chatModel := models.NewChatModel(driver, &modelName, apiConfig)
+
+	s.updateTaskProgress(taskID, 0.15, "Prepared prompts and LLM.")
+	temperature := mem.Temperature
+	resp, err := chatModel.ModelDriver.ChatWithMessages(ctx, modelName, messages, apiConfig, &models.ChatConfig{Temperature: &temperature}, nil)
+	if err != nil {
+		return nil, fmt.Errorf("chat model: %w", err)
+	}
+	if resp == nil || resp.Answer == nil {
+		return nil, errors.New("empty response from chat model")
+	}
+	s.updateTaskProgress(taskID, 0.35, "Get extracted result from LLM.")
+
+	return parseMemoryExtraction(*resp.Answer, extractTypes), nil
+}
+
+// buildExtractedMessage builds the persisted envelope for one extracted
+// memory item. Field set matches buildRawMessage except message_type and
+// source_id, which listMemoryMessages uses to aggregate extracts. Only
+// logical message fields are set here; the doc engine maps them to
+// storage fields (including tokenization) at insert time.
+func buildExtractedMessage(messageID, sourceID int64, memoryID string, msg MemoryMessage, item extractedMemory, now time.Time) map[string]any {
+	var invalidAt any
+	if strings.TrimSpace(item.InvalidAt) != "" {
+		invalidAt = formatMemoryTime(item.InvalidAt, now)
+	}
+	return map[string]any{
+		"message_id":   messageID,
+		"message_type": item.MessageType,
+		"source_id":    sourceID,
+		"memory_id":    memoryID,
+		"user_id":      msg.UserID,
+		"agent_id":     msg.AgentID,
+		"session_id":   msg.SessionID,
+		"content":      item.Content,
+		"valid_at":     formatMemoryTime(item.ValidAt, now),
+		"invalid_at":   invalidAt,
+		"forget_at":    nil,
+		"status":       true,
+	}
+}
+
+// parseMemoryExtraction ports memory.utils.msg_util.get_json_result_from_llm_response
+// plus the per-type flattening in extract_by_llm. Only the configured
+// extract types are collected; unparseable responses yield an empty list.
+func parseMemoryExtraction(answer string, extractTypes []string) []extractedMemory {
+	clean := strings.TrimSpace(answer)
+	clean = strings.TrimPrefix(clean, "```json")
+	clean = strings.TrimPrefix(clean, "```")
+	clean = strings.TrimSuffix(clean, "```")
+	if start := strings.Index(clean, "{"); start >= 0 {
+		if end := strings.LastIndex(clean, "}"); end > start {
+			clean = clean[start : end+1]
+		}
+	}
+
+	var parsed map[string][]struct {
+		Content   string `json:"content"`
+		ValidAt   string `json:"valid_at"`
+		InvalidAt string `json:"invalid_at"`
+	}
+	if err := json.Unmarshal([]byte(strings.TrimSpace(clean)), &parsed); err != nil {
+		common.Warn("memory: failed to parse LLM extraction result", zap.Error(err))
+		return nil
+	}
+
+	var out []extractedMemory
+	for _, memoryType := range extractTypes {
+		for _, item := range parsed[memoryType] {
+			if strings.TrimSpace(item.Content) == "" {
+				continue
+			}
+			out = append(out, extractedMemory{
+				MessageType: memoryType,
+				Content:     item.Content,
+				ValidAt:     item.ValidAt,
+				InvalidAt:   item.InvalidAt,
+			})
+		}
+	}
+	return out
+}
+
+// formatMemoryTime normalizes an LLM-supplied timestamp (ISO 8601 or
+// already-formatted) into memoryTimeLayout. Unparseable or empty input
+// falls back to the supplied time.
+func formatMemoryTime(value string, fallback time.Time) string {
+	v := strings.TrimSpace(value)
+	if v != "" {
+		for _, layout := range []string{time.RFC3339, "2006-01-02T15:04:05", "2006-01-02 15:04:05", "2006-01-02"} {
+			if t, err := time.Parse(layout, v); err == nil {
+				return t.UTC().Format(memoryTimeLayout)
+			}
+		}
+	}
+	return fallback.UTC().Format(memoryTimeLayout)
+}
+
+// updateTaskProgress stamps and persists task progress, mirroring
+// Python TaskService.update_progress call sites. Failures are logged
+// and swallowed so progress reporting never breaks extraction.
+func (s *MemoryMessageService) updateTaskProgress(taskID string, progress float64, msg string) {
+	if s == nil || s.taskDAO == nil || taskID == "" {
+		return
+	}
+	stamped := time.Now().Format(memoryTimeLayout) + " " + msg
+	if err := s.taskDAO.UpdateProgress(taskID, progress, stamped); err != nil {
+		common.Warn("memory: update task progress failed", zap.Error(err))
+	}
+}
+
+func payloadInt64(v any) int64 {
+	switch n := v.(type) {
+	case float64:
+		return int64(n)
+	case int64:
+		return n
+	case int:
+		return int64(n)
+	case json.Number:
+		i, _ := n.Int64()
+		return i
+	case string:
+		var i int64
+		fmt.Sscanf(n, "%d", &i)
+		return i
+	}
+	return 0
+}
+
+func payloadString(v any) string {
+	s, _ := v.(string)
+	return s
+}
--- a/internal/service/memory_message_service.go
+++ b/internal/service/memory_message_service.go
@@ -65,17 +65,6 @@ import (
 	models "ragflow/internal/entity/models"
 )

-// ErrEmbedderNotWired is returned by QueueSaveToMemoryTask when
-// the embedding-model call is reached. The Go runtime has no
-// embedding model port yet; until one lands, callers see this
-// error and know to fall back to the Python Canvas.
-var ErrEmbedderNotWired = errors.New(
-	"memory: embedder not wired in Go — " +
-		"QueueSaveToMemoryTask runs the lookup + message construction " +
-		"but cannot embed / save until internal/rag/llm/embedding_model " +
-		"ships (Phase 8b follow-up)",
-)
-
 // MemoryMessage is the wire shape for QueueSaveToMemoryTask. It
 // mirrors the Python `message_dict` built in
 // agent/component/message.py:_save_to_memory:
@@ -157,10 +146,10 @@ func (s *MemoryMessageService) QueueSaveToMemoryTask(ctx context.Context, memory
 		}
 		// (2) + (3) build the raw_message envelope. The Go port
 		// keeps the same field set as Python:344-386 so the
-		// downstream extractor (also still on the Python side)
-		// can consume the row without schema changes.
+		// downstream extractor can consume the row without
+		// schema changes.
 		rawMessageID := generateRawMessageID()
-		rawMessage := buildRawMessage(rawMessageID, memoryID, mem, msg)
+		rawMessage := buildRawMessage(rawMessageID, memoryID, msg)

 		if err := s.embedAndSave(ctx, mem, rawMessage); err != nil {
 			res.Failed = append(res.Failed, MemoryFailure{
@@ -201,35 +190,32 @@ func generateRawMessageID() int64 {

 // buildRawMessage constructs the raw_message envelope that gets
 // passed to embed_and_save (and persisted in the message table
-// for the async extractor to read).
-func buildRawMessage(rawMessageID int64, memoryID string, mem *CreateMemoryResponse, msg MemoryMessage) map[string]any {
+// for the async extractor to read). Only logical message fields
+// are set here, mirroring Python queue_save_to_memory_task; the
+// doc engine maps them to storage fields at insert time
+// (Elasticsearch tokenizes content before write, Infinity
+// tokenizes on save).
+func buildRawMessage(
+	rawMessageID int64,
+	memoryID string,
+	msg MemoryMessage,
+) map[string]any {
 	content := fmt.Sprintf("User Input: %s\nAgent Response: %s",
 		msg.UserInput, msg.AgentResponse)
-	out := map[string]any{
-		"message_id":             rawMessageID,
-		"message_type":           "raw",
-		"message_type_kwd":       "raw",
-		"source_id":              0,
-		"memory_id":              memoryID,
-		"user_id":                msg.UserID,
-		"agent_id":               msg.AgentID,
-		"session_id":             msg.SessionID,
-		"content":                content,
-		"content_ltks":           content,
-		"tokenized_content_ltks": content,
-		"valid_at":               time.Now().UTC().Format("2006-01-02 15:04:05"),
-		"invalid_at":             nil,
-		"forget_at":              nil,
-		"status":                 true,
-		"status_int":             1,
+	return map[string]any{
+		"message_id":   rawMessageID,
+		"message_type": "raw",
+		"source_id":    0,
+		"memory_id":    memoryID,
+		"user_id":      msg.UserID,
+		"agent_id":     msg.AgentID,
+		"session_id":   msg.SessionID,
+		"content":      content,
+		"valid_at":     time.Now().UTC().Format("2006-01-02 15:04:05"),
+		"invalid_at":   nil,
+		"forget_at":    nil,
+		"status":       true,
 	}
-	if mem != nil {
-		// The embedder uses the memory's embd_id; keep the
-		// pointer on the envelope so embed_and_save can
-		// pick the right model when it lands.
-		out["_memory_embd_id"] = mem.EmbdID
-	}
-	return out
 }

 // buildTaskRow constructs the Task row the async extractor polls.
@@ -246,31 +232,51 @@ func buildTaskRow(rawMessageID int64, memoryID string) map[string]any {
 }

 func (s *MemoryMessageService) embedAndSave(ctx context.Context, mem *CreateMemoryResponse, rawMessage map[string]any) error {
+	return s.embedAndSaveMessages(ctx, mem, []map[string]any{rawMessage})
+}
+
+// embedAndSaveMessages embeds every message's content with the memory's
+// embedding model and inserts the batch into the memory's chunk store,
+// creating the index on first use. Mirrors Python embed_and_save.
+func (s *MemoryMessageService) embedAndSaveMessages(ctx context.Context, mem *CreateMemoryResponse, messages []map[string]any) error {
 	if mem == nil {
 		return errors.New("memory not found")
 	}
 	if s == nil || s.memories == nil || s.memories.docEngine == nil {
 		return errors.New("message store is not initialized")
 	}
+	if len(messages) == 0 {
+		return nil
+	}

-	content, _ := rawMessage["content"].(string)
+	contents := make([]string, len(messages))
+	for i, message := range messages {
+		contents[i], _ = message["content"].(string)
+	}
 	driver, modelName, apiConfig, maxTokens, err := NewModelProviderService().ResolveModelConfig(ctx, mem.TenantID, entity.ModelTypeEmbedding, mem.EmbdID)
 	if err != nil {
 		return err
 	}
 	embeddingModel := models.NewEmbeddingModel(driver, &modelName, apiConfig, maxTokens)
-	embeddings, err := embeddingModel.ModelDriver.Embed(ctx, embeddingModel.ModelName, []string{content}, embeddingModel.APIConfig, &models.EmbeddingConfig{Dimension: 0}, nil)
+	embeddings, err := embeddingModel.ModelDriver.Embed(ctx, embeddingModel.ModelName, contents, embeddingModel.APIConfig, &models.EmbeddingConfig{Dimension: 0}, nil)
 	if err != nil {
 		return err
 	}
-	if len(embeddings) == 0 || len(embeddings[0].Embedding) == 0 {
-		return errors.New("embedding response is empty")
+	if len(embeddings) != len(messages) {
+		return fmt.Errorf("embedding response count %d does not match message count %d", len(embeddings), len(messages))
 	}

-	vector := embeddings[0].Embedding
-	rawMessage[fmt.Sprintf("q_%d_vec", len(vector))] = vector
-	rawMessage["id"] = fmt.Sprintf("%s_%d", rawMessage["memory_id"], rawMessage["message_id"])
-	rawMessage["doc_id"] = rawMessage["memory_id"]
+	vectorDim := 0
+	for i, message := range messages {
+		vector := embeddings[i].Embedding
+		if len(vector) == 0 {
+			return errors.New("embedding response is empty")
+		}
+		vectorDim = len(vector)
+		message[fmt.Sprintf("q_%d_vec", len(vector))] = vector
+		message["id"] = fmt.Sprintf("%s_%v", message["memory_id"], message["message_id"])
+		message["doc_id"] = message["memory_id"]
+	}

 	indexName := memoryIndexName(mem.TenantID)
 	exists, err := s.memories.docEngine.ChunkStoreExists(ctx, indexName, mem.ID)
@@ -278,22 +284,21 @@ func (s *MemoryMessageService) embedAndSave(ctx context.Context, mem *CreateMemo
 		return fmt.Errorf("check message index: %w", err)
 	}
 	if !exists {
-		if err := s.memories.docEngine.CreateChunkStore(ctx, indexName, mem.ID, len(vector), ""); err != nil {
+		if err := s.memories.docEngine.CreateChunkStore(ctx, indexName, mem.ID, vectorDim, ""); err != nil {
 			return fmt.Errorf("create message index: %w", err)
 		}
 	}
-	if _, err := s.memories.docEngine.InsertChunks(ctx, []map[string]interface{}{mapStringAny(rawMessage)}, indexName, mem.ID); err != nil {
+	docs := make([]map[string]interface{}, len(messages))
+	for i, message := range messages {
+		docs[i] = mapStringAny(message)
+	}
+	if _, err := s.memories.docEngine.InsertChunks(ctx, docs, indexName, mem.ID); err != nil {
 		return fmt.Errorf("insert message into memory: %w", err)
 	}

 	return nil
 }

-// embedAndSave is kept for older unit tests; production uses the method above.
-func embedAndSave(_ context.Context, _ *CreateMemoryResponse, _ map[string]any) error {
-	return ErrEmbedderNotWired
-}
-
 func (s *MemoryMessageService) insertTask(ctx context.Context, row map[string]any) error {
 	if s == nil {
 		return errors.New("nil MemoryMessageService")
--- a/internal/service/memory_message_service_test.go
+++ b/internal/service/memory_message_service_test.go
@@ -14,20 +14,15 @@
 //  limitations under the License.
 //

-// memory_message_service_test.go — Phase 8b real MemorySaver port
-// tests.
+// memory_message_service_test.go — MemorySaver port tests.
 //
 // Coverage focuses on the synchronous parts (memory lookup + raw
-// message construction + result aggregation). The embedder call
-// is exercised only to verify it loud-fails with
-// ErrEmbedderNotWired — the embedder port itself is a separate
-// follow-up.
+// message construction + result aggregation).

 package service

 import (
 	"context"
-	"errors"
 	"strings"
 	"testing"
 )
@@ -73,11 +68,11 @@ func TestQueueSaveToMemoryTask_MissingAgentID(t *testing.T) {
 }

 // TestBuildRawMessage_EnvelopeShape: the row envelope carries
-// every field the Python extractor reads.
+// every logical field the Python extractor reads. Storage-level
+// fields (message_type_kwd, content_ltks, tokenized_content_ltks,
+// status_int) are the doc engine's concern and must not be set here.
 func TestBuildRawMessage_EnvelopeShape(t *testing.T) {
-	mem := &CreateMemoryResponse{}
-	mem.EmbdID = "embd-1"
-	raw := buildRawMessage(42, "mem-1", mem, MemoryMessage{
+	raw := buildRawMessage(42, "mem-1", MemoryMessage{
 		UserID:        "u1",
 		AgentID:       "a1",
 		SessionID:     "s1",
@@ -86,19 +81,24 @@ func TestBuildRawMessage_EnvelopeShape(t *testing.T) {
 	})

 	want := map[string]any{
-		"message_id":      int64(42),
-		"message_type":    "raw",
-		"memory_id":       "mem-1",
-		"user_id":         "u1",
-		"agent_id":        "a1",
-		"session_id":      "s1",
-		"_memory_embd_id": "embd-1",
+		"message_id":   int64(42),
+		"message_type": "raw",
+		"memory_id":    "mem-1",
+		"user_id":      "u1",
+		"agent_id":     "a1",
+		"session_id":   "s1",
+		"status":       true,
 	}
 	for k, want := range want {
 		if got := raw[k]; got != want {
 			t.Errorf("raw[%q] = %v (%T), want %v (%T)", k, got, got, want, want)
 		}
 	}
+	for _, storageField := range []string{"message_type_kwd", "content_ltks", "tokenized_content_ltks", "status_int"} {
+		if _, ok := raw[storageField]; ok {
+			t.Errorf("raw[%q] is a storage field and must not be set by the service layer", storageField)
+		}
+	}

 	content, _ := raw["content"].(string)
 	if !strings.Contains(content, "User Input: hi") {
@@ -112,20 +112,6 @@ func TestBuildRawMessage_EnvelopeShape(t *testing.T) {
 	}
 }

-// TestBuildRawMessage_NilMemory: when the lookup returned nil
-// (e.g. caller already filtered NotFound), buildRawMessage
-// doesn't panic and omits _memory_embd_id.
-func TestBuildRawMessage_NilMemory(t *testing.T) {
-	raw := buildRawMessage(1, "m", nil, MemoryMessage{AgentID: "a"})
-	if _, ok := raw["_memory_embd_id"]; ok {
-		t.Errorf("_memory_embd_id should be absent when mem is nil")
-	}
-	// The other fields still land.
-	if raw["agent_id"] != "a" {
-		t.Errorf("agent_id missing or wrong: %v", raw["agent_id"])
-	}
-}
-
 // TestBuildTaskRow_Shape: the task row mirrors the Python Task
 // entity's memory-task shape.
 func TestBuildTaskRow_Shape(t *testing.T) {
@@ -160,18 +146,6 @@ func TestTaskFromRow_InitializesProgressMessage(t *testing.T) {
 	}
 }

-// TestEmbedAndSave_LoudFails: until the embedder port lands,
-// the call must return ErrEmbedderNotWired — not panic, not
-// silently succeed. A successful embed_and_save would mask the
-// deferred state and let callers persist data that never gets
-// embedded.
-func TestEmbedAndSave_LoudFails(t *testing.T) {
-	err := embedAndSave(context.Background(), nil, nil)
-	if !errors.Is(err, ErrEmbedderNotWired) {
-		t.Errorf("err = %v, want ErrEmbedderNotWired", err)
-	}
-}
-
 // TestGenerateRawMessageID_Unique: two calls produce different
 // values. (Wall-clock based today; the Redis-backed counter will
 // be added when the project's Redis client lands.)