From ed05f32feece16f28e894832a780f3d5ba1ececa Mon Sep 17 00:00:00 2001 From: euvre <93761161+euvre@users.noreply.github.com> Date: Fri, 31 Jul 2026 11:23:58 +0800 Subject: [PATCH] feat(go): implement memory extraction task consumer (#17404) --- cmd/ragflow_server.go | 7 + internal/dao/task.go | 8 + internal/engine/elasticsearch/chunk.go | 68 +++- .../elasticsearch/chunk_helpers_test.go | 64 ++++ internal/service/memory.go | 19 + internal/service/memory_extractor.go | 353 ++++++++++++++++++ internal/service/memory_message_service.go | 115 +++--- .../service/memory_message_service_test.go | 62 +-- 8 files changed, 592 insertions(+), 104 deletions(-) create mode 100644 internal/service/memory_extractor.go diff --git a/cmd/ragflow_server.go b/cmd/ragflow_server.go index 77f4bc5e16..897b0d6ce3 100644 --- a/cmd/ragflow_server.go +++ b/cmd/ragflow_server.go @@ -560,6 +560,13 @@ func runIngestor(ctx context.Context, cancel context.CancelFunc, args *serverArg common.Error("Failed to initialize ingestor", err) } + // Memory extraction consumer: drains task_type="memory" messages + // from the te.0.common Redis stream and runs LLM extraction. + memoryConsumerCtx, stopMemoryConsumer := context.WithCancel(ctx) + defer stopMemoryConsumer() + memoryMessageSvc := service.NewMemoryMessageService(service.NewMemoryService()) + go memoryMessageSvc.StartTaskConsumer(memoryConsumerCtx) + common.Info("\n ____ __ _\n" + " / _/___ ____ ____ _____/ /_(_)___ ____ ________ ______ _____ _____\n" + " / // __ \\/ __ `/ _ \\/ ___/ __/ / __ \\/ __ \\ / ___/ _ \\/ ___/ | / / _ \\/ ___/\n" + diff --git a/internal/dao/task.go b/internal/dao/task.go index b97a3412ca..bf52b8efe1 100644 --- 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 { diff --git a/internal/engine/elasticsearch/chunk.go b/internal/engine/elasticsearch/chunk.go index 1b1a4f1436..8a330bfecc 100644 --- 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", diff --git a/internal/engine/elasticsearch/chunk_helpers_test.go b/internal/engine/elasticsearch/chunk_helpers_test.go index ea000c4869..e77d5f151d 100644 --- 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) +} diff --git a/internal/service/memory.go b/internal/service/memory.go index 465c125060..f28be1b4a7 100644 --- 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: diff --git a/internal/service/memory_extractor.go b/internal/service/memory_extractor.go new file mode 100644 index 0000000000..3e07192729 --- /dev/null +++ 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..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 +} diff --git a/internal/service/memory_message_service.go b/internal/service/memory_message_service.go index d54510b60b..24799b8c03 100644 --- 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") diff --git a/internal/service/memory_message_service_test.go b/internal/service/memory_message_service_test.go index af55c17d96..832062a635 100644 --- 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.)