Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -13,12 +13,14 @@
import com.openjiuwen.core.memory.graph.extraction.prompts.TemplateManager;
import com.openjiuwen.core.memory.graph.extraction.prompts.entity_extraction.ExtractionPromptLanguageBase;

import java.math.BigInteger;
import java.time.Instant;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Optional;

/**
* Prompt building and entity extraction orchestration by episode type.
Expand Down Expand Up @@ -105,10 +107,41 @@ public static PromptCall extractEntityAttributes(Entity entity, String content,
if (extras != null) {
kwargs.putAll(extras);
}
applyHumanSummaryTarget(entity, kwargs);
return new PromptCall(kwargs, TemplateManager.getInstance().get(templateName),
MultilingualBaseModel.responseFormat(EntitySummary.class, language));
}

private static void applyHumanSummaryTarget(Entity entity, Map<String, Object> kwargs) {
if (!"Human".equalsIgnoreCase(entity.getObjType())) {
return;
}
doubleSummaryTarget(kwargs.get("summary_target"))
.ifPresent(doubledTarget -> kwargs.put("summary_target", doubledTarget));
}

private static Optional<Number> doubleSummaryTarget(Object summaryTarget) {
BigInteger target;
if (summaryTarget instanceof BigInteger bigInteger) {
target = bigInteger;
} else if (summaryTarget instanceof Number number) {
target = BigInteger.valueOf(number.longValue());
} else if (summaryTarget instanceof String value
&& !value.isEmpty() && value.codePoints().allMatch(Character::isDigit)) {
target = new BigInteger(value);
} else {
return Optional.empty();
}
BigInteger doubledTarget = target.shiftLeft(1);
if (doubledTarget.bitLength() < Integer.SIZE) {
return Optional.of(Integer.valueOf(doubledTarget.intValue()));
}
if (doubledTarget.bitLength() < Long.SIZE) {
return Optional.of(Long.valueOf(doubledTarget.longValue()));
}
return Optional.of(doubledTarget);
}

/**
* extractRelationDeclaration.
*
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,10 @@
package com.openjiuwen.core.memory.graph.graph_memory;

import com.openjiuwen.core.common.concurrent.OpenJiuwenExecutors;
import com.openjiuwen.core.common.exception.ErrorHelper;
import com.openjiuwen.core.common.exception.StatusCode;
import com.openjiuwen.core.common.logging.LoggerProtocol;
import com.openjiuwen.core.common.logging.Loggers;
import com.openjiuwen.core.foundation.llm.Model;
import com.openjiuwen.core.foundation.llm.schema.AssistantMessage;
import com.openjiuwen.core.foundation.store.base_embedding.Embedding;
Expand Down Expand Up @@ -60,6 +64,7 @@
*/
public class GraphMemory {
private static final String STORE_TYPE = "graph mem store";
private static final LoggerProtocol MEMORY_LOGGER = Loggers.MEMORY;

/**
* Public record SearchHit used by the Java parity implementation.
Expand Down Expand Up @@ -244,7 +249,9 @@ public void registerSearchStrategy(String name, SearchConfig searchEntity, Searc
throw new IllegalArgumentException("Search config cannot be registered as an empty value.");
}
if (searchStrategies.containsKey(name) && !isForceRegister) {
throw new IllegalArgumentException("Search config with name [" + name + "] already exists.");
throw ErrorHelper.buildError(StatusCode.MEMORY_STORE_VALIDATION_INVALID,
"store_type", STORE_TYPE,
"error_msg", "Search config with name [" + name + "] already exists.");
}
searchStrategies.put(name,
List.of(searchEntity != null ? copySearchConfig(searchEntity) : new SearchConfig(),
Expand Down Expand Up @@ -785,12 +792,35 @@ private AssistantMessage invokeLlm(ExtractionPrompts.PromptCall promptCall, Map<
if (extra != null) {
params.putAll(extra);
}
AssistantMessage response;
semaphore.acquire();
try {
return llmClient.invoke(params.get("messages"), null, null, null, null, null, null, null, null, params);
response = llmClient.invoke(params.get("messages"), null, null, null, null, null, null, null, null, params);
} finally {
semaphore.release();
}
if (isDebugEnabled) {
logLlmInvocation(promptCall.template().getName(), params.get("messages"), response);
}
return response;
}

private static void logLlmInvocation(String templateName, Object messages, AssistantMessage response) {
String separator = System.lineSeparator() + "=".repeat(60) + System.lineSeparator();
String debugMessage = "TEMPLATE " + templateName + separator + lastMessageContent(messages)
+ separator + String.valueOf(response.getContent());
MEMORY_LOGGER.info("Graph Memory LLM Invoke: {}", debugMessage);
}

private static String lastMessageContent(Object messages) {
if (!(messages instanceof List<?> messageList) || messageList.isEmpty()) {
return "";
}
Object lastMessage = messageList.get(messageList.size() - 1);
if (lastMessage instanceof Map<?, ?> messageMap) {
return String.valueOf(messageMap.get("content"));
}
return String.valueOf(lastMessage);
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

import static org.assertj.core.api.Assertions.assertThat;

import com.openjiuwen.core.foundation.store.graph.Entity;
import com.openjiuwen.core.memory.config.graph.AddMemStrategy;
import com.openjiuwen.core.memory.config.graph.EpisodeType;
import com.openjiuwen.core.memory.config.graph.GraphDefaults;
Expand Down Expand Up @@ -62,4 +63,21 @@ void templateManagerShouldLoadPromptResources() {
assertThat(template.toMessages()).isNotEmpty();
assertThat(manager.contains("entity_extraction_relation_en")).isTrue();
}

@Test
void humanSummaryTargetShouldBeTwiceTheConfiguredLimit() {
Entity entity = new Entity();
entity.setName("A公司");
Entity human = new Entity();
human.setName("张明");
human.setObjType("Human");

ExtractionPrompts.PromptCall entityPrompt = ExtractionPrompts.extractEntityAttributes(
entity, "content", "", "cn", Map.of("summary_target", "110"), 2);
ExtractionPrompts.PromptCall humanPrompt = ExtractionPrompts.extractEntityAttributes(
human, "content", "", "cn", Map.of("summary_target", "110"), 2);

assertThat(entityPrompt.kwargs()).containsEntry("summary_target", "110");
assertThat(humanPrompt.kwargs()).containsEntry("summary_target", 220);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

import com.openjiuwen.core.common.exception.BaseError;
import com.openjiuwen.core.common.exception.StatusCode;
import com.openjiuwen.core.foundation.llm.Model;
import com.openjiuwen.core.foundation.llm.schema.AssistantMessage;
import com.openjiuwen.core.foundation.llm.schema.ModelClientConfig;
Expand All @@ -17,8 +19,14 @@
import com.openjiuwen.core.foundation.store.graph.Relation;
import com.openjiuwen.core.memory.config.graph.SearchConfig;

import ch.qos.logback.classic.Level;
import ch.qos.logback.classic.Logger;
import ch.qos.logback.classic.spi.ILoggingEvent;
import ch.qos.logback.core.read.ListAppender;

import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import org.slf4j.LoggerFactory;

import java.lang.reflect.Method;
import java.nio.file.Path;
Expand Down Expand Up @@ -101,6 +109,21 @@ void shouldRejectUnknownStrategy() {
.isInstanceOf(IllegalArgumentException.class).hasMessageContaining("Strategy [missing] not found");
}

@Test
void shouldRejectDuplicateStrategyWithMemoryValidationError() {
GraphConfig config =
GraphConfig.builder().uri(tempDir.resolve("graph.db").toString()).backend("in_memory").build();
GraphMemory memory = new GraphMemory(config);

assertThatThrownBy(() -> memory.registerSearchStrategy(
"default", new SearchConfig(), new SearchConfig(), new SearchConfig(), false))
.isInstanceOfSatisfying(BaseError.class, error -> {
assertThat(error.getStatus()).isEqualTo(StatusCode.MEMORY_STORE_VALIDATION_INVALID);
assertThat(error.getCode()).isEqualTo(StatusCode.MEMORY_STORE_VALIDATION_INVALID.getCode());
})
.hasMessageContaining("Search config with name [default] already exists.");
}

@Test
void shouldPrepareConversationEpisodesAndBuildHistory() throws Exception {
GraphConfig config =
Expand Down Expand Up @@ -156,6 +179,37 @@ void shouldAddMemoryThroughBasicMainFlow() throws Exception {
assertThat(entity.getContent()).contains("Alice likes coffee");
}

@Test
void shouldLogLlmInvocationWhenDebugEnabled() throws Exception {
Logger memoryLogger = (Logger) LoggerFactory.getLogger("memory");
ListAppender<ILoggingEvent> appender = new ListAppender<>();
appender.start();
Level previousLevel = memoryLogger.getLevel();
memoryLogger.setLevel(Level.INFO);
memoryLogger.addAppender(appender);
try {
GraphConfig config =
GraphConfig.builder().uri(tempDir.resolve("debug-graph.db").toString()).backend("in_memory").build();
FakeModel fakeModel = new FakeModel("{\"extracted_relations\":[]}",
"{\"extracted_entities\":[{\"name\":\"Alice\",\"entityTypeId\":0}]}",
"{\"extracted_relations\":[]}",
"{\"summary\":\"Alice likes coffee\",\"attributes\":{}}");
GraphMemory memory = new GraphMemory(config, fakeModel, true, null, null, Map.of(), null, "cn", true);
memory.attachEmbedder(new DummyEmbedding());

memory.addMemory(com.openjiuwen.core.memory.config.graph.EpisodeType.DOCUMENT, "user-1",
"Alice likes coffee", null, OffsetDateTime.now());

assertThat(appender.list).extracting(ILoggingEvent::getFormattedMessage)
.anySatisfy(message -> assertThat(message)
.contains("Graph Memory LLM Invoke: TEMPLATE", "Alice likes coffee",
"extracted_relations"));
} finally {
memoryLogger.detachAppender(appender);
memoryLogger.setLevel(previousLevel);
}
}

private static final class DummyEmbedding extends Embedding {
@Override
public List<Float> embedQuery(String text) {
Expand Down