Skip to content

Commit 98fa88b

Browse files
Aias00Aias00
committed
refactor: enhance logging and simplify memory management methods in Memory and MilvusVectorStoreService
1 parent 5e9087d commit 98fa88b

File tree

2 files changed

+98
-84
lines changed

2 files changed

+98
-84
lines changed

‎mem4j-core/src/main/java/io/github/mem4j/memory/Memory.java‎

Lines changed: 43 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,13 @@
2020
import io.github.mem4j.embeddings.EmbeddingService;
2121
import io.github.mem4j.llms.LLMService;
2222
import io.github.mem4j.vectorstores.VectorStoreService;
23+
import java.util.*;
24+
import java.util.stream.Collectors;
2325
import org.slf4j.Logger;
2426
import org.slf4j.LoggerFactory;
2527
import org.springframework.stereotype.Service;
2628

27-
import java.util.*;
28-
import java.util.stream.Collectors;
29-
30-
/**
31-
* Core memory management class for Java Mem4j
32-
*/
29+
/** Core memory management class for Java Mem4j */
3330
@Service
3431
public class Memory {
3532

@@ -52,16 +49,12 @@ public Memory(MemoryConfigurable config, VectorStoreService vectorStoreService,
5249
this.embeddingService = embeddingService;
5350
}
5451

55-
/**
56-
* Add memories from a conversation
57-
*/
52+
/** Add memories from a conversation */
5853
public void add(List<Message> messages, String userId) {
5954
add(messages, userId, null, true, MemoryType.FACTUAL);
6055
}
6156

62-
/**
63-
* Add memories with custom parameters
64-
*/
57+
/** Add memories with custom parameters */
6558
public void add(List<Message> messages, String userId, Map<String, Object> metadata, boolean infer,
6659
MemoryType memoryType) {
6760

@@ -88,47 +81,63 @@ public void add(List<Message> messages, String userId, Map<String, Object> metad
8881
logger.error("Error adding memories for user {}", userId, e);
8982
throw new RuntimeException("Failed to add memories", e);
9083
}
91-
9284
}
9385

94-
/**
95-
* Search for relevant memories
96-
*/
86+
/** Search for relevant memories */
9787
public List<MemoryItem> search(String query, String userId) {
9888
return search(query, userId, null, 10, null);
9989
}
10090

101-
/**
102-
* Search with custom parameters
103-
*/
91+
/** Search with custom parameters */
10492
public List<MemoryItem> search(String query, String userId, Map<String, Object> filters, int limit,
10593
Double threshold) {
10694

10795
try {
10896
// Generate embedding for query
10997
Double[] queryEmbedding = embeddingService.embed(query);
98+
logger.debug("Generated query embedding with {} dimensions for query: '{}'", queryEmbedding.length, query);
11099

111100
// Build search filters
112101
Map<String, Object> searchFilters = buildSearchFilters(userId, filters);
102+
logger.debug("Search filters: {}", searchFilters);
103+
104+
// Log the threshold being used
105+
Double actualThreshold = threshold != null ? threshold : config.getSimilarityThreshold();
106+
logger.debug("Using similarity threshold: {}", actualThreshold);
113107

114108
// Search vector store
115-
List<MemoryItem> results = vectorStoreService.search(queryEmbedding, searchFilters, limit,
116-
threshold != null ? threshold : config.getSimilarityThreshold());
109+
List<MemoryItem> results = vectorStoreService.search(queryEmbedding, searchFilters, limit, actualThreshold);
110+
111+
// If no results found and threshold > 0.2, try with lower threshold
112+
if (results.isEmpty() && actualThreshold > 0.2) {
113+
logger.debug("No results with threshold {}, retrying with 0.2", actualThreshold);
114+
results = vectorStoreService.search(queryEmbedding, searchFilters, limit, 0.2);
115+
actualThreshold = 0.2; // Update for logging
116+
}
117+
118+
logger.info("Found {} memories for query: '{}' with threshold: {}", results.size(), query, actualThreshold);
119+
120+
// Log found results for debugging
121+
if (!results.isEmpty()) {
122+
for (MemoryItem item : results) {
123+
logger.debug("Found memory: '{}' with score: {}", item.getContent(), item.getScore());
124+
}
125+
}
126+
else {
127+
logger.warn("No memories found for query: '{}' with user_id: '{}' and threshold: {}", query, userId,
128+
actualThreshold);
129+
}
117130

118-
logger.info("Found {} memories for query: {}", results.size(), query);
119131
return results;
120132
}
121133
catch (Exception e) {
122134

123135
logger.error("Error searching memories for query: {}", query, e);
124136
throw new RuntimeException("Failed to search memories", e);
125137
}
126-
127138
}
128139

129-
/**
130-
* Get all memories for a user
131-
*/
140+
/** Get all memories for a user */
132141
public List<MemoryItem> getAll(String userId, Map<String, Object> filters, int limit) {
133142

134143
try {
@@ -139,12 +148,9 @@ public List<MemoryItem> getAll(String userId, Map<String, Object> filters, int l
139148
logger.error("Error getting all memories for user {}", userId, e);
140149
throw new RuntimeException("Failed to get memories", e);
141150
}
142-
143151
}
144152

145-
/**
146-
* Update a memory item
147-
*/
153+
/** Update a memory item */
148154
public void update(String memoryId, Map<String, Object> data) {
149155

150156
try {
@@ -170,12 +176,9 @@ public void update(String memoryId, Map<String, Object> data) {
170176
logger.error("Error updating memory: {}", memoryId, e);
171177
throw new RuntimeException("Failed to update memory", e);
172178
}
173-
174179
}
175180

176-
/**
177-
* Delete a memory item
178-
*/
181+
/** Delete a memory item */
179182
public void delete(String memoryId) {
180183

181184
try {
@@ -186,12 +189,9 @@ public void delete(String memoryId) {
186189
logger.error("Error deleting memory: {}", memoryId, e);
187190
throw new RuntimeException("Failed to delete memory", e);
188191
}
189-
190192
}
191193

192-
/**
193-
* Delete all memories for a user
194-
*/
194+
/** Delete all memories for a user */
195195
public void deleteAll(String userId) {
196196

197197
try {
@@ -203,12 +203,9 @@ public void deleteAll(String userId) {
203203
logger.error("Error deleting all memories for user {}", userId, e);
204204
throw new RuntimeException("Failed to delete memories", e);
205205
}
206-
207206
}
208207

209-
/**
210-
* Extract memories from conversation using LLM
211-
*/
208+
/** Extract memories from conversation using LLM */
212209
private List<String> extractMemories(List<Message> messages, boolean infer) {
213210

214211
if (!infer) {
@@ -226,8 +223,6 @@ private List<String> extractMemories(List<Message> messages, boolean infer) {
226223

227224
String prompt = String.format("""
228225
Extract key memories from this conversation. Focus on:
229-
- Keep the original language and tone
230-
- Be concise and specific
231226
- Important facts about the user
232227
- User preferences and behaviors
233228
- Significant events or experiences
@@ -250,9 +245,7 @@ private List<String> extractMemories(List<Message> messages, boolean infer) {
250245
.collect(Collectors.toList());
251246
}
252247

253-
/**
254-
* Create a memory item from content
255-
*/
248+
/** Create a memory item from content */
256249
private MemoryItem createMemoryItem(String content, String userId, Map<String, Object> metadata,
257250
MemoryType memoryType) {
258251

@@ -263,9 +256,7 @@ private MemoryItem createMemoryItem(String content, String userId, Map<String, O
263256
return item;
264257
}
265258

266-
/**
267-
* Build search filters
268-
*/
259+
/** Build search filters */
269260
private Map<String, Object> buildSearchFilters(String userId, Map<String, Object> additionalFilters) {
270261

271262
Map<String, Object> filters = new HashMap<>();
@@ -278,9 +269,7 @@ private Map<String, Object> buildSearchFilters(String userId, Map<String, Object
278269
return filters;
279270
}
280271

281-
/**
282-
* Get memory by ID
283-
*/
272+
/** Get memory by ID */
284273
public MemoryItem get(String memoryId) {
285274

286275
try {
@@ -292,9 +281,7 @@ public MemoryItem get(String memoryId) {
292281
}
293282
}
294283

295-
/**
296-
* Reset all memories (for testing)
297-
*/
284+
/** Reset all memories (for testing) */
298285
public void reset() {
299286

300287
try {

‎mem4j-core/src/main/java/io/github/mem4j/vectorstores/MilvusVectorStoreService.java‎

Lines changed: 55 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,22 @@
11
/*
2-
* Copyright 2024-2025 the original author or authors.
3-
*
4-
* Licensed under the Apache License, Version 2.0 (the "License");
5-
* you may not use this file except in compliance with the Lic client.createCollection(createCollectionParam);
6-
logger.info("Milvus 集合已创建: {}", collectionName);
7-
8-
// 加载集合到内存中
9-
ensureCollectionLoaded();se.
10-
* You may obtain a copy of the License at
11-
*
12-
* https://www.apache.org/licenses/LICENSE-2.0
13-
*
14-
* Unless required by applicable law or agreed to in writing, software
15-
* distributed under the License is distributed on an "AS IS" BASIS,
16-
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17-
* See the License for the specific language governing permissions and
18-
* limitations under the License.
19-
*/
2+
* Copyright 2024-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the Lic client.createCollection(createCollectionParam);
6+
logger.info("Milvus 集合已创建: {}", collectionName);
7+
8+
// 加载集合到内存中
9+
ensureCollectionLoaded();se.
10+
* You may obtain a copy of the License at
11+
*
12+
* https://www.apache.org/licenses/LICENSE-2.0
13+
*
14+
* Unless required by applicable law or agreed to in writing, software
15+
* distributed under the License is distributed on an "AS IS" BASIS,
16+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
17+
* See the License for the specific language governing permissions and
18+
* limitations under the License.
19+
*/
2020

2121
package io.github.mem4j.vectorstores;
2222

@@ -34,22 +34,22 @@
3434
import io.milvus.param.collection.FieldType;
3535
import io.milvus.param.collection.HasCollectionParam;
3636
import io.milvus.param.collection.LoadCollectionParam;
37-
import io.milvus.param.index.CreateIndexParam;
3837
import io.milvus.param.dml.DeleteParam;
3938
import io.milvus.param.dml.InsertParam;
4039
import io.milvus.param.dml.QueryParam;
4140
import io.milvus.param.dml.SearchParam;
41+
import io.milvus.param.index.CreateIndexParam;
42+
import java.time.Instant;
43+
import java.util.*;
44+
import java.util.stream.Collectors;
4245
import lombok.AllArgsConstructor;
4346
import org.slf4j.Logger;
4447
import org.slf4j.LoggerFactory;
4548
import org.springframework.stereotype.Service;
4649

47-
import java.time.Instant;
48-
import java.util.*;
49-
import java.util.stream.Collectors;
50-
5150
/**
5251
* Milvus implementation of VectorStoreService
52+
*
5353
* <p>
5454
* This implementation provides vector storage capabilities using Milvus, supporting
5555
* operations like adding, searching, updating, and deleting memory items.
@@ -326,9 +326,13 @@ public List<MemoryItem> search(Double[] queryEmbedding, Map<String, Object> filt
326326

327327
// 构建搜索表达式
328328
String searchExpr = buildSearchExpression(filters);
329+
logger.debug("Milvus search expression: '{}'", searchExpr);
330+
logger.debug("Search filters: {}", filters);
331+
logger.debug("Search threshold: {}, limit: {}", threshold, limit);
329332

330333
// 转换查询向量为Float类型
331334
List<Float> queryVector = Arrays.stream(queryEmbedding).map(Double::floatValue).toList();
335+
logger.debug("Query vector length: {}", queryVector.size());
332336

333337
// 构建搜索参数
334338
SearchParam searchParam = SearchParam.newBuilder()
@@ -356,6 +360,7 @@ public List<MemoryItem> search(Double[] queryEmbedding, Map<String, Object> filt
356360
if (searchResults != null && searchResults.getResults() != null) {
357361
// 获取搜索结果的字段数据
358362
List<FieldData> fieldsData = searchResults.getResults().getFieldsDataList();
363+
logger.debug("Milvus returned {} field data entries", fieldsData.size());
359364

360365
if (!fieldsData.isEmpty()) {
361366
// 创建字段映射
@@ -382,10 +387,12 @@ else if (fieldData.getType() == DataType.FloatVector) {
382387
}
383388

384389
fieldMap.put(fieldName, values);
390+
logger.debug("Field '{}' has {} values", fieldName, values.size());
385391
}
386392

387393
// 获取结果数量
388-
int resultCount = fieldMap.get("id").size();
394+
int resultCount = fieldMap.get("id") != null ? fieldMap.get("id").size() : 0;
395+
logger.debug("Total result count from Milvus: {}", resultCount);
389396

390397
// 构建MemoryItem对象并计算相似度
391398
for (int i = 0; i < resultCount; i++) {
@@ -396,20 +403,42 @@ else if (fieldData.getType() == DataType.FloatVector) {
396403
if (itemEmbedding != null) {
397404
double similarity = cosineSimilarity(queryEmbedding, itemEmbedding);
398405
item.setScore(similarity);
406+
logger.debug("Item '{}' has similarity: {}, threshold: {}", item.getContent(), similarity,
407+
threshold);
399408

400409
// 根据阈值过滤结果
401410
if (threshold == null || similarity >= threshold) {
402411
results.add(item);
412+
logger.debug("Added item to results: '{}'", item.getContent());
413+
}
414+
else {
415+
logger.debug("Filtered out item due to low similarity: '{}' ({})", item.getContent(),
416+
similarity);
403417
}
404418
}
419+
else {
420+
logger.warn("Item has no embedding: '{}'", item.getContent());
421+
}
405422
}
406423

407424
// 按相似度分数降序排序
408425
results.sort((a, b) -> Double.compare(b.getScore(), a.getScore()));
409426
}
427+
else {
428+
logger.debug("No field data returned from Milvus search");
429+
}
430+
}
431+
else {
432+
logger.debug("No search results returned from Milvus");
433+
}
434+
435+
logger.debug("Found {} similar memories after filtering", results.size());
436+
437+
// 如果没有结果，尝试不使用阈值进行搜索以诊断问题
438+
if (results.isEmpty() && threshold != null && threshold > 0) {
439+
logger.warn("No results found with threshold {}, this might indicate a threshold issue", threshold);
410440
}
411441

412-
logger.debug("Found {} similar memories", results.size());
413442
return results;
414443
}
415444
catch (Exception e) {
@@ -642,9 +671,7 @@ public void reset() {
642671
}
643672
}
644673

645-
/**
646-
* 构建搜索表达式
647-
*/
674+
/** 构建搜索表达式 */
648675
private String buildSearchExpression(Map<String, Object> filters) {
649676
if (filters == null || filters.isEmpty()) {
650677
return "";

0 commit comments

Comments
 (0)