From 7d2207fe6ceae660e2a102a46fd817c163e5cf29 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Mon, 7 Sep 2026 13:36:34 +0800 Subject: [PATCH 01/30] feat(ai): define retrieval api contract and configuration --- .../ai/retrieval/codec/RetrievalApiJson.java | 66 ++----- .../retrieval/config/RetrievalProperties.java | 4 +- .../service/RetrievalRequestValidator.java | 33 +++- .../retrieval/api/RetrievalApiJsonTest.java | 178 +----------------- .../api/RetrievalPropertiesTest.java | 4 - 5 files changed, 43 insertions(+), 242 deletions(-) diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java index cb1a6b643..dfeeb650c 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java @@ -27,7 +27,6 @@ import com.google.gson.stream.JsonToken; import java.io.IOException; import java.io.StringReader; -import java.math.BigDecimal; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; @@ -39,41 +38,26 @@ import org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse; import org.apache.geaflow.ai.retrieval.api.model.RetrievalTrace; import org.apache.geaflow.ai.retrieval.api.model.TraceStage; -import org.apache.geaflow.ai.retrieval.config.RetrievalProperties; import org.apache.geaflow.ai.retrieval.model.document.SourceRef; import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; import org.apache.geaflow.ai.retrieval.model.graph.GraphPathRef; -import org.apache.geaflow.ai.retrieval.service.RetrievalErrorValidator; -import org.apache.geaflow.ai.retrieval.service.RetrievalRequestValidator; -import org.apache.geaflow.ai.retrieval.service.RetrievalResponseValidator; /** Strict, dependency-light JSON boundary for the v1 retrieval API. */ public final class RetrievalApiJson { private static final int MAX_JSON_LENGTH = 1024 * 1024; - private static final int MAX_JSON_DEPTH = 256; private static final Gson GSON = new Gson(); private RetrievalApiJson() { } public static RetrievalRequest parseRequest(String json) { - return parseRequest(json, defaultProperties()); - } - - public static RetrievalRequest parseRequest(String json, RetrievalProperties properties) { JsonObject object = object(json); validateRequestShape(object); - RetrievalRequest request = GSON.fromJson(object, RetrievalRequest.class); - new RetrievalRequestValidator(properties).validate(request); - return request; + return GSON.fromJson(object, RetrievalRequest.class); } public static RetrievalResponse parseResponse(String json) { - return parseResponse(json, defaultProperties()); - } - - public static RetrievalResponse parseResponse(String json, RetrievalProperties properties) { JsonObject object = object(json); requireString(object, "requestId"); requireString(object, "graphName"); @@ -97,7 +81,7 @@ public static RetrievalResponse parseResponse(String json, RetrievalProperties p response.setTrace(parseTrace(object.getAsJsonObject("trace"))); response.setEffectiveBudget(parseBudget(object.getAsJsonObject("effectiveBudget"), true)); response.setDegradedChannels(strings(object, "degradedChannels")); - return RetrievalResponseValidator.validate(response, properties); + return response; } public static RetrievalError parseError(String json) { @@ -123,30 +107,13 @@ public static RetrievalError parseError(String json) { RetrievalError error = new RetrievalError(object.get("requestId").getAsString(), errorCode, object.get("message").getAsString()); error.setRetriable(retriable); - return RetrievalErrorValidator.validate(error); + return error; } public static String toJson(Object value) { - return toJson(value, defaultProperties()); - } - - public static String toJson(Object value, RetrievalProperties properties) { - if (value instanceof RetrievalResponse) { - RetrievalResponseValidator.validate((RetrievalResponse) value, properties); - } else if (value instanceof RetrievalError) { - RetrievalErrorValidator.validate((RetrievalError) value); - } else if (value instanceof RetrievalRequest) { - new RetrievalRequestValidator(properties).validate((RetrievalRequest) value); - } return GSON.toJson(value); } - private static RetrievalProperties defaultProperties() { - RetrievalProperties properties = new RetrievalProperties(); - properties.validateConfiguration(); - return properties; - } - private static JsonObject object(String json) { if (json == null || json.trim().isEmpty()) { throw new JsonParseException("JSON object is required"); @@ -200,7 +167,8 @@ private static RetrievalBudget parseBudget(JsonObject object, boolean requiredVa Integer maxCandidates = optionalInt(object, "maxCandidates"); Integer tokenBudget = optionalInt(object, "tokenBudget"); if (requiredValues && (topK == null || timeoutMs == null || maxCandidates == null - || tokenBudget == null)) { + || tokenBudget == null || topK < 1 || timeoutMs < 1 || maxCandidates < 1 + || tokenBudget < 1 || topK > maxCandidates)) { throw new JsonParseException("effectiveBudget contains invalid values"); } return new RetrievalBudget(topK, timeoutMs, maxCandidates, tokenBudget); @@ -295,11 +263,12 @@ private static Integer optionalInt(JsonObject object, String name) { if (!element.isJsonPrimitive() || !element.getAsJsonPrimitive().isNumber()) { throw new JsonParseException(name + " must be an integer"); } - try { - return new BigDecimal(element.getAsString()).intValueExact(); - } catch (NumberFormatException | ArithmeticException e) { + double value = element.getAsDouble(); + if (!Double.isFinite(value) || value != Math.rint(value) + || value < Integer.MIN_VALUE || value > Integer.MAX_VALUE) { throw new JsonParseException(name + " must be an integer"); } + return (int) value; } private static String string(JsonObject object, String name) { @@ -354,7 +323,7 @@ private static void rejectDuplicateKeys(String json) { JsonReader reader = new JsonReader(new StringReader(json)); reader.setLenient(false); try { - consume(reader, 0); + consume(reader); if (reader.peek() != JsonToken.END_DOCUMENT) { throw new JsonParseException("trailing JSON content"); } @@ -363,10 +332,9 @@ private static void rejectDuplicateKeys(String json) { } } - private static void consume(JsonReader reader, int depth) throws IOException { + private static void consume(JsonReader reader) throws IOException { JsonToken token = reader.peek(); if (token == JsonToken.BEGIN_OBJECT) { - int childDepth = nextDepth(depth); java.util.HashSet names = new java.util.HashSet<>(); reader.beginObject(); while (reader.hasNext()) { @@ -374,14 +342,13 @@ private static void consume(JsonReader reader, int depth) throws IOException { if (!names.add(name)) { throw new JsonParseException("duplicate JSON field: " + name); } - consume(reader, childDepth); + consume(reader); } reader.endObject(); } else if (token == JsonToken.BEGIN_ARRAY) { - int childDepth = nextDepth(depth); reader.beginArray(); while (reader.hasNext()) { - consume(reader, childDepth); + consume(reader); } reader.endArray(); } else if (token == JsonToken.STRING) { @@ -396,11 +363,4 @@ private static void consume(JsonReader reader, int depth) throws IOException { throw new JsonParseException("invalid JSON token: " + token); } } - - private static int nextDepth(int depth) { - if (depth >= MAX_JSON_DEPTH) { - throw new JsonParseException("retrieval JSON exceeds maximum nesting depth"); - } - return depth + 1; - } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java index c2fb3fbda..3cdb32a14 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java @@ -56,10 +56,10 @@ public class RetrievalProperties { @Init public void validateConfiguration() { - if (!"v1".equals(configVersion) || blank(readyGraphName)) { + if (blank(configVersion) || blank(readyGraphName)) { throw new RetrievalException( org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.INVALID_REQUEST, - "retrieval config-version must be v1 and ready-graph-name is required"); + "retrieval config-version and ready-graph-name are required"); } if (!RetrievalMode.KEYWORD.name().equals(defaultMode)) { throw invalid("default-mode must be KEYWORD"); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java index a139fe736..546f02712 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java @@ -19,7 +19,6 @@ package org.apache.geaflow.ai.retrieval.service; import java.util.List; -import java.util.Objects; import org.apache.geaflow.ai.retrieval.api.model.ExecutionMode; import org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget; import org.apache.geaflow.ai.retrieval.api.model.RetrievalCommand; @@ -37,8 +36,7 @@ public class RetrievalRequestValidator { private final RetrievalProperties properties; public RetrievalRequestValidator(RetrievalProperties properties) { - this.properties = Objects.requireNonNull(properties, "properties"); - this.properties.validateConfiguration(); + this.properties = properties; } public RetrievalCommand validate(RetrievalRequest request) { @@ -57,10 +55,23 @@ public RetrievalCommand validate(RetrievalRequest request) { } RetrievalBudget input = request.getBudget(); - final RetrievalBudget effectiveBudget = RetrievalBudgetValidator.validate(input, - properties, false); + final int topK = valueOrDefault(input == null ? null : input.getTopK(), + properties.getDefaultTopK()); + final int timeoutMs = valueOrDefault(input == null ? null : input.getTimeoutMs(), + properties.getDefaultTimeoutMs()); + final int maxCandidates = valueOrDefault(input == null ? null : input.getMaxCandidates(), + properties.getDefaultMaxCandidates()); + final int tokenBudget = valueOrDefault(input == null ? null : input.getTokenBudget(), + properties.getDefaultTokenBudget()); + range("topK", topK, properties.getMaxTopK()); + range("timeoutMs", timeoutMs, properties.getMaxTimeoutMs()); + range("maxCandidates", maxCandidates, properties.getMaxCandidates()); + range("tokenBudget", tokenBudget, properties.getMaxTokenBudget()); + if (topK > maxCandidates) { + throw invalid("topK must not exceed maxCandidates"); + } return new RetrievalCommand(graphName, query, mode, executionMode, - effectiveBudget); + new RetrievalBudget(topK, timeoutMs, maxCandidates, tokenBudget)); } private static String trimmed(String value, String name, int maxLength) { @@ -96,6 +107,16 @@ private static ExecutionMode parseExecutionMode(String value) { } } + private static int valueOrDefault(Integer value, int defaultValue) { + return value == null ? defaultValue : value; + } + + private static void range(String name, int value, int max) { + if (value < 1 || value > max) { + throw invalid(name + " must be between 1 and " + max); + } + } + private static RetrievalException invalid(String message) { return new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, message); } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java index 67d2d9783..85cccd8ae 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java @@ -19,7 +19,6 @@ package org.apache.geaflow.ai.retrieval.api; import com.google.gson.JsonObject; -import com.google.gson.JsonParseException; import com.google.gson.JsonParser; import java.io.IOException; import java.io.InputStream; @@ -33,18 +32,14 @@ import org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget; import org.apache.geaflow.ai.retrieval.api.model.RetrievalError; import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; -import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; -import org.apache.geaflow.ai.retrieval.api.model.RetrievalMode; import org.apache.geaflow.ai.retrieval.api.model.RetrievalRequest; import org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse; import org.apache.geaflow.ai.retrieval.api.model.RetrievalTrace; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalMode; import org.apache.geaflow.ai.retrieval.api.model.TraceStage; import org.apache.geaflow.ai.retrieval.codec.RetrievalApiJson; -import org.apache.geaflow.ai.retrieval.config.RetrievalProperties; -import org.apache.geaflow.ai.retrieval.model.document.SourceRef; import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; import org.apache.geaflow.ai.retrieval.model.evidence.EvidenceKind; -import org.apache.geaflow.ai.retrieval.model.graph.GraphPathRef; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -116,75 +111,6 @@ public void rejectsWrongTypesMalformedJsonAndDuplicateFields() { () -> RetrievalApiJson.parseRequest(unicodeOversized.toString())); } - @Test - public void appliesSemanticValidationAtRequestJsonBoundary() { - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.parseRequest( - "{\"graphName\":\"g\",\"query\":\"q\",\"mode\":\"HYBRID\"}")); - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.parseRequest( - "{\"graphName\":\"g\",\"query\":\"q\"," - + "\"executionMode\":\"PARALLEL\"}")); - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.parseRequest( - "{\"graphName\":\"g\",\"query\":\"q\"," - + "\"budget\":{\"topK\":0}}")); - - StringBuilder longQuery = new StringBuilder("{\"graphName\":\"g\",\"query\":\""); - for (int i = 0; i < 4097; i++) { - longQuery.append('q'); - } - longQuery.append("\"}"); - Assertions.assertThrows(RetrievalException.class, - () -> RetrievalApiJson.parseRequest(longQuery.toString())); - - RetrievalProperties properties = new RetrievalProperties(); - properties.setMaxTopK(5); - properties.setDefaultTopK(5); - properties.validateConfiguration(); - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.parseRequest( - "{\"graphName\":\"g\",\"query\":\"q\"," - + "\"budget\":{\"topK\":6}}", properties)); - } - - @Test - public void rejectsSemanticallyInvalidRequestSerialization() { - RetrievalRequest request = new RetrievalRequest(); - request.setGraphName("graph"); - request.setQuery("query"); - request.setMode("HYBRID"); - - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.toJson(request)); - } - - @Test - public void rejectsExcessiveJsonNestingBeforeParsing() { - StringBuilder json = new StringBuilder("{\"graphName\":\"g\",\"query\":\"q\",\"nested\":"); - for (int i = 0; i < 257; i++) { - json.append('['); - } - json.append('0'); - for (int i = 0; i < 257; i++) { - json.append(']'); - } - json.append('}'); - - Assertions.assertThrows(JsonParseException.class, - () -> RetrievalApiJson.parseRequest(json.toString())); - } - - @Test - public void parsesOnlyExactIntegerBudgetNumbers() { - Assertions.assertThrows(JsonParseException.class, () -> RetrievalApiJson.parseRequest( - "{\"graphName\":\"g\",\"query\":\"q\"," - + "\"budget\":{\"topK\":1.0000000000000001}}")); - Assertions.assertThrows(JsonParseException.class, () -> RetrievalApiJson.parseRequest( - "{\"graphName\":\"g\",\"query\":\"q\"," - + "\"budget\":{\"topK\":2147483648}}")); - - RetrievalRequest request = RetrievalApiJson.parseRequest( - "{\"graphName\":\"g\",\"query\":\"q\"," - + "\"budget\":{\"topK\":1e1}}"); - Assertions.assertEquals(Integer.valueOf(10), request.getBudget().getTopK()); - } - @Test public void responseKeepsRequiredCollectionsAndRoundTrips() { RetrievalResponse response = new RetrievalResponse(); @@ -254,76 +180,6 @@ public void rejectsResponseTraceWithoutStages() { Assertions.assertThrows(RuntimeException.class, () -> RetrievalApiJson.parseResponse(json)); } - @Test - public void rejectsUnsupportedTraceModesAndUnboundedEffectiveBudget() { - String parallel = responseJson("PARALLEL", "SEQUENTIAL", 10, 3000, 100, 4096); - Assertions.assertThrows(RuntimeException.class, - () -> RetrievalApiJson.parseResponse(parallel)); - - String overLimit = responseJson("KEYWORD", "SEQUENTIAL", 101, 3000, 100, 4096); - Assertions.assertThrows(RuntimeException.class, - () -> RetrievalApiJson.parseResponse(overLimit)); - } - - @Test - public void rejectsBlankRequiredResponseAndErrorFields() { - String blankGraph = responseJson("KEYWORD", "SEQUENTIAL", 10, 3000, 100, 4096) - .replace("\"graphName\":\"graph\"", "\"graphName\":\" \""); - Assertions.assertThrows(RuntimeException.class, - () -> RetrievalApiJson.parseResponse(blankGraph)); - Assertions.assertThrows(RuntimeException.class, () -> RetrievalApiJson.parseError( - "{\"requestId\":\" \",\"code\":\"INTERNAL_ERROR\"," - + "\"message\":\"x\",\"retriable\":false}")); - } - - @Test - public void rejectsIncompleteResponseSerialization() { - RetrievalResponse response = new RetrievalResponse(); - Assertions.assertThrows(RuntimeException.class, () -> RetrievalApiJson.toJson(response)); - } - - @Test - public void rejectsPartialEffectiveBudgetSerialization() { - RetrievalResponse response = validResponse(); - response.setEffectiveBudget(new RetrievalBudget(null, 3000, 100, 4096)); - - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.toJson(response)); - } - - @Test - public void rejectsNullResponseCollectionElements() { - RetrievalResponse evidence = validResponse(); - evidence.setEvidence(Collections.singletonList(null)); - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.toJson(evidence)); - - RetrievalResponse paths = validResponse(); - paths.setPaths(Collections.singletonList(null)); - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.toJson(paths)); - - RetrievalResponse sources = validResponse(); - sources.setSources(Collections.singletonList(null)); - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.toJson(sources)); - - RetrievalResponse channels = validResponse(); - channels.setDegradedChannels(Collections.singletonList(null)); - Assertions.assertThrows(RetrievalException.class, () -> RetrievalApiJson.toJson(channels)); - } - - @Test - public void reportsMissingTraceModesAsInvalidRequest() { - RetrievalResponse missingMode = validResponse(); - missingMode.getTrace().setSelectedMode(null); - RetrievalException modeException = Assertions.assertThrows(RetrievalException.class, - () -> RetrievalApiJson.toJson(missingMode)); - Assertions.assertEquals(RetrievalErrorCode.INVALID_REQUEST, modeException.getCode()); - - RetrievalResponse missingExecutionMode = validResponse(); - missingExecutionMode.getTrace().setExecutionMode(null); - RetrievalException executionException = Assertions.assertThrows(RetrievalException.class, - () -> RetrievalApiJson.toJson(missingExecutionMode)); - Assertions.assertEquals(RetrievalErrorCode.INVALID_REQUEST, executionException.getCode()); - } - @Test public void errorCodeAndRetriableFlagAreStable() { RetrievalError error = new RetrievalError("req-3", RetrievalErrorCode.INDEX_NOT_READY, @@ -375,25 +231,6 @@ public void outputUsesDocumentedCamelCaseNames() { Assertions.assertFalse(json.has("graph_name")); } - private static RetrievalResponse validResponse() { - RetrievalResponse response = new RetrievalResponse(); - response.setRequestId("req-1"); - response.setGraphName("graph"); - response.setGraphVersion("v1"); - response.setEvidence(Collections.singletonList(new Evidence("e-1", EvidenceKind.CHUNK, - "text", null, null, null, null, null, null, 1))); - RetrievalTrace trace = new RetrievalTrace(); - trace.setTraceVersion("v1"); - trace.setOriginalQuery("confucius"); - trace.setSelectedMode(RetrievalMode.KEYWORD); - trace.setExecutionMode(ExecutionMode.SEQUENTIAL); - trace.setStages(Collections.singletonList(new TraceStage("keyword", "COMPLETED", null))); - trace.setStopReason("COMPLETED"); - response.setTrace(trace); - response.setEffectiveBudget(new RetrievalBudget(10, 3000, 100, 4096)); - return response; - } - private static String read(String resource) throws IOException { InputStream stream = RetrievalApiJsonTest.class.getClassLoader() .getResourceAsStream(resource); @@ -408,17 +245,4 @@ private static String read(String resource) throws IOException { return result.toString(); } } - - private static String responseJson(String mode, String executionMode, int topK, - int timeoutMs, int maxCandidates, int tokenBudget) { - return "{\"requestId\":\"r\",\"graphName\":\"graph\"," - + "\"graphVersion\":\"v1\",\"evidence\":[],\"paths\":[]," - + "\"sources\":[],\"degradedChannels\":[],\"trace\":{" - + "\"traceVersion\":\"v1\",\"originalQuery\":\"q\"," - + "\"selectedMode\":\"" + mode + "\",\"executionMode\":\"" - + executionMode + "\",\"stages\":[]},\"effectiveBudget\":{" - + "\"topK\":" + topK + ",\"timeoutMs\":" + timeoutMs - + ",\"maxCandidates\":" + maxCandidates + ",\"tokenBudget\":" - + tokenBudget + "}}"; - } } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java index 0a8a639dc..11181c4ea 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java @@ -83,10 +83,6 @@ public void rejectsBlankIdentityAndInvertedDefaults() { properties = new RetrievalProperties(); properties.setMaxTopK(0); assertInvalid(properties); - - properties = new RetrievalProperties(); - properties.setConfigVersion("v2"); - assertInvalid(properties); } private static void assertInvalid(RetrievalProperties properties) { From c928d6692211bafc3c2e23f5746e07735ace7a2f Mon Sep 17 00:00:00 2001 From: aotenjou Date: Thu, 10 Sep 2026 21:47:48 +0800 Subject: [PATCH 02/30] fix(ai): harden retrieval api validation --- .../ai/retrieval/codec/RetrievalApiJson.java | 29 +++++++++++-- .../retrieval/config/RetrievalProperties.java | 4 +- .../service/RetrievalBudgetValidator.java | 4 -- .../service/RetrievalRequestValidator.java | 33 +++------------ .../service/RetrievalResponseValidator.java | 12 +----- .../retrieval/api/RetrievalApiJsonTest.java | 41 +++++++++++++++++++ .../api/RetrievalPropertiesTest.java | 4 ++ 7 files changed, 80 insertions(+), 47 deletions(-) diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java index dfeeb650c..1a51985a4 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java @@ -38,9 +38,12 @@ import org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse; import org.apache.geaflow.ai.retrieval.api.model.RetrievalTrace; import org.apache.geaflow.ai.retrieval.api.model.TraceStage; +import org.apache.geaflow.ai.retrieval.config.RetrievalProperties; import org.apache.geaflow.ai.retrieval.model.document.SourceRef; import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; import org.apache.geaflow.ai.retrieval.model.graph.GraphPathRef; +import org.apache.geaflow.ai.retrieval.service.RetrievalErrorValidator; +import org.apache.geaflow.ai.retrieval.service.RetrievalResponseValidator; /** Strict, dependency-light JSON boundary for the v1 retrieval API. */ public final class RetrievalApiJson { @@ -58,6 +61,10 @@ public static RetrievalRequest parseRequest(String json) { } public static RetrievalResponse parseResponse(String json) { + return parseResponse(json, defaultProperties()); + } + + public static RetrievalResponse parseResponse(String json, RetrievalProperties properties) { JsonObject object = object(json); requireString(object, "requestId"); requireString(object, "graphName"); @@ -81,7 +88,7 @@ public static RetrievalResponse parseResponse(String json) { response.setTrace(parseTrace(object.getAsJsonObject("trace"))); response.setEffectiveBudget(parseBudget(object.getAsJsonObject("effectiveBudget"), true)); response.setDegradedChannels(strings(object, "degradedChannels")); - return response; + return RetrievalResponseValidator.validate(response, properties); } public static RetrievalError parseError(String json) { @@ -107,13 +114,28 @@ public static RetrievalError parseError(String json) { RetrievalError error = new RetrievalError(object.get("requestId").getAsString(), errorCode, object.get("message").getAsString()); error.setRetriable(retriable); - return error; + return RetrievalErrorValidator.validate(error); } public static String toJson(Object value) { + return toJson(value, defaultProperties()); + } + + public static String toJson(Object value, RetrievalProperties properties) { + if (value instanceof RetrievalResponse) { + RetrievalResponseValidator.validate((RetrievalResponse) value, properties); + } else if (value instanceof RetrievalError) { + RetrievalErrorValidator.validate((RetrievalError) value); + } return GSON.toJson(value); } + private static RetrievalProperties defaultProperties() { + RetrievalProperties properties = new RetrievalProperties(); + properties.validateConfiguration(); + return properties; + } + private static JsonObject object(String json) { if (json == null || json.trim().isEmpty()) { throw new JsonParseException("JSON object is required"); @@ -167,8 +189,7 @@ private static RetrievalBudget parseBudget(JsonObject object, boolean requiredVa Integer maxCandidates = optionalInt(object, "maxCandidates"); Integer tokenBudget = optionalInt(object, "tokenBudget"); if (requiredValues && (topK == null || timeoutMs == null || maxCandidates == null - || tokenBudget == null || topK < 1 || timeoutMs < 1 || maxCandidates < 1 - || tokenBudget < 1 || topK > maxCandidates)) { + || tokenBudget == null)) { throw new JsonParseException("effectiveBudget contains invalid values"); } return new RetrievalBudget(topK, timeoutMs, maxCandidates, tokenBudget); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java index 3cdb32a14..c2fb3fbda 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java @@ -56,10 +56,10 @@ public class RetrievalProperties { @Init public void validateConfiguration() { - if (blank(configVersion) || blank(readyGraphName)) { + if (!"v1".equals(configVersion) || blank(readyGraphName)) { throw new RetrievalException( org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.INVALID_REQUEST, - "retrieval config-version and ready-graph-name are required"); + "retrieval config-version must be v1 and ready-graph-name is required"); } if (!RetrievalMode.KEYWORD.name().equals(defaultMode)) { throw invalid("default-mode must be KEYWORD"); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalBudgetValidator.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalBudgetValidator.java index db4dbf324..b2fa47526 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalBudgetValidator.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalBudgetValidator.java @@ -38,10 +38,6 @@ public static RetrievalBudget validate(RetrievalBudget budget, RetrievalProperti } return defaults(properties); } - if (requireAllValues && (budget.getTopK() == null || budget.getTimeoutMs() == null - || budget.getMaxCandidates() == null || budget.getTokenBudget() == null)) { - throw invalid("all budget values are required"); - } Integer topK = valueOrDefault(budget.getTopK(), properties.getDefaultTopK()); Integer timeoutMs = valueOrDefault(budget.getTimeoutMs(), properties.getDefaultTimeoutMs()); Integer maxCandidates = valueOrDefault(budget.getMaxCandidates(), diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java index 546f02712..a139fe736 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java @@ -19,6 +19,7 @@ package org.apache.geaflow.ai.retrieval.service; import java.util.List; +import java.util.Objects; import org.apache.geaflow.ai.retrieval.api.model.ExecutionMode; import org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget; import org.apache.geaflow.ai.retrieval.api.model.RetrievalCommand; @@ -36,7 +37,8 @@ public class RetrievalRequestValidator { private final RetrievalProperties properties; public RetrievalRequestValidator(RetrievalProperties properties) { - this.properties = properties; + this.properties = Objects.requireNonNull(properties, "properties"); + this.properties.validateConfiguration(); } public RetrievalCommand validate(RetrievalRequest request) { @@ -55,23 +57,10 @@ public RetrievalCommand validate(RetrievalRequest request) { } RetrievalBudget input = request.getBudget(); - final int topK = valueOrDefault(input == null ? null : input.getTopK(), - properties.getDefaultTopK()); - final int timeoutMs = valueOrDefault(input == null ? null : input.getTimeoutMs(), - properties.getDefaultTimeoutMs()); - final int maxCandidates = valueOrDefault(input == null ? null : input.getMaxCandidates(), - properties.getDefaultMaxCandidates()); - final int tokenBudget = valueOrDefault(input == null ? null : input.getTokenBudget(), - properties.getDefaultTokenBudget()); - range("topK", topK, properties.getMaxTopK()); - range("timeoutMs", timeoutMs, properties.getMaxTimeoutMs()); - range("maxCandidates", maxCandidates, properties.getMaxCandidates()); - range("tokenBudget", tokenBudget, properties.getMaxTokenBudget()); - if (topK > maxCandidates) { - throw invalid("topK must not exceed maxCandidates"); - } + final RetrievalBudget effectiveBudget = RetrievalBudgetValidator.validate(input, + properties, false); return new RetrievalCommand(graphName, query, mode, executionMode, - new RetrievalBudget(topK, timeoutMs, maxCandidates, tokenBudget)); + effectiveBudget); } private static String trimmed(String value, String name, int maxLength) { @@ -107,16 +96,6 @@ private static ExecutionMode parseExecutionMode(String value) { } } - private static int valueOrDefault(Integer value, int defaultValue) { - return value == null ? defaultValue : value; - } - - private static void range(String name, int value, int max) { - if (value < 1 || value > max) { - throw invalid(name + " must be between 1 and " + max); - } - } - private static RetrievalException invalid(String message) { return new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, message); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalResponseValidator.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalResponseValidator.java index 80c4bfc1d..02a846496 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalResponseValidator.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalResponseValidator.java @@ -51,15 +51,9 @@ public static RetrievalResponse validate(RetrievalResponse response, } required(trace.getTraceVersion(), "traceVersion"); required(trace.getOriginalQuery(), "originalQuery"); - if (trace.getSelectedMode() == null) { - throw invalid("selectedMode is required"); - } if (trace.getSelectedMode() != RetrievalMode.KEYWORD) { throw unsupported("unsupported retrieval mode in trace"); } - if (trace.getExecutionMode() == null) { - throw invalid("executionMode is required"); - } if (trace.getExecutionMode() != ExecutionMode.SEQUENTIAL) { throw unsupported("unsupported execution mode in trace"); } @@ -75,10 +69,8 @@ public static RetrievalResponse validate(RetrievalResponse response, required(stage.getStatus(), "stage.status"); } if (response.getEvidence() == null || response.getPaths() == null - || response.getSources() == null || response.getDegradedChannels() == null - || response.getEvidence().contains(null) || response.getPaths().contains(null) - || response.getSources().contains(null) || response.getDegradedChannels().contains(null)) { - throw invalid("response collections are required and must not contain null"); + || response.getSources() == null || response.getDegradedChannels() == null) { + throw invalid("response collections are required"); } RetrievalBudget effectiveBudget = response.getEffectiveBudget(); RetrievalBudgetValidator.validate(effectiveBudget, properties, true); diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java index 85cccd8ae..5a9e68f07 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java @@ -180,6 +180,34 @@ public void rejectsResponseTraceWithoutStages() { Assertions.assertThrows(RuntimeException.class, () -> RetrievalApiJson.parseResponse(json)); } + @Test + public void rejectsUnsupportedTraceModesAndUnboundedEffectiveBudget() { + String parallel = responseJson("PARALLEL", "SEQUENTIAL", 10, 3000, 100, 4096); + Assertions.assertThrows(RuntimeException.class, + () -> RetrievalApiJson.parseResponse(parallel)); + + String overLimit = responseJson("KEYWORD", "SEQUENTIAL", 101, 3000, 100, 4096); + Assertions.assertThrows(RuntimeException.class, + () -> RetrievalApiJson.parseResponse(overLimit)); + } + + @Test + public void rejectsBlankRequiredResponseAndErrorFields() { + String blankGraph = responseJson("KEYWORD", "SEQUENTIAL", 10, 3000, 100, 4096) + .replace("\"graphName\":\"graph\"", "\"graphName\":\" \""); + Assertions.assertThrows(RuntimeException.class, + () -> RetrievalApiJson.parseResponse(blankGraph)); + Assertions.assertThrows(RuntimeException.class, () -> RetrievalApiJson.parseError( + "{\"requestId\":\" \",\"code\":\"INTERNAL_ERROR\"," + + "\"message\":\"x\",\"retriable\":false}")); + } + + @Test + public void rejectsIncompleteResponseSerialization() { + RetrievalResponse response = new RetrievalResponse(); + Assertions.assertThrows(RuntimeException.class, () -> RetrievalApiJson.toJson(response)); + } + @Test public void errorCodeAndRetriableFlagAreStable() { RetrievalError error = new RetrievalError("req-3", RetrievalErrorCode.INDEX_NOT_READY, @@ -245,4 +273,17 @@ private static String read(String resource) throws IOException { return result.toString(); } } + + private static String responseJson(String mode, String executionMode, int topK, + int timeoutMs, int maxCandidates, int tokenBudget) { + return "{\"requestId\":\"r\",\"graphName\":\"graph\"," + + "\"graphVersion\":\"v1\",\"evidence\":[],\"paths\":[]," + + "\"sources\":[],\"degradedChannels\":[],\"trace\":{" + + "\"traceVersion\":\"v1\",\"originalQuery\":\"q\"," + + "\"selectedMode\":\"" + mode + "\",\"executionMode\":\"" + + executionMode + "\",\"stages\":[]},\"effectiveBudget\":{" + + "\"topK\":" + topK + ",\"timeoutMs\":" + timeoutMs + + ",\"maxCandidates\":" + maxCandidates + ",\"tokenBudget\":" + + tokenBudget + "}}"; + } } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java index 11181c4ea..0a8a639dc 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java @@ -83,6 +83,10 @@ public void rejectsBlankIdentityAndInvertedDefaults() { properties = new RetrievalProperties(); properties.setMaxTopK(0); assertInvalid(properties); + + properties = new RetrievalProperties(); + properties.setConfigVersion("v2"); + assertInvalid(properties); } private static void assertInvalid(RetrievalProperties properties) { From 9818d52f6b3199d37981de890fe0821537e9cba6 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Fri, 11 Sep 2026 09:38:32 +0800 Subject: [PATCH 03/30] feat(ai): add versioned ingestion metadata and import state model --- .../retrieval/ingest/ImportStateMachine.java | 55 ++++ .../metadata/ChunkingConfiguration.java | 82 ++++++ .../retrieval/metadata/DatasetManifest.java | 202 +++++++++++++++ .../metadata/GraphBuildMetadata.java | 157 +++++++++++ .../ai/retrieval/metadata/ImportMetadata.java | 223 ++++++++++++++++ .../ai/retrieval/metadata/ImportState.java | 28 ++ .../metadata/InMemoryMetadataStore.java | 245 ++++++++++++++++++ .../metadata/IndexBuildMetadata.java | 119 +++++++++ .../retrieval/metadata/MetadataException.java | 48 ++++ .../ai/retrieval/metadata/MetadataJson.java | 204 +++++++++++++++ .../ai/retrieval/metadata/MetadataStore.java | 89 +++++++ .../metadata/MetadataValidation.java | 54 ++++ .../retrieval/metadata/QualityCounters.java | 124 +++++++++ .../ingest/ImportStateMachineTest.java | 58 +++++ .../metadata/InMemoryMetadataStoreTest.java | 125 +++++++++ .../retrieval/metadata/MetadataModelTest.java | 114 ++++++++ 16 files changed, 1927 insertions(+) create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ImportStateMachine.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ChunkingConfiguration.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/DatasetManifest.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/GraphBuildMetadata.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ImportMetadata.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ImportState.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/InMemoryMetadataStore.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/IndexBuildMetadata.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataException.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataJson.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataStore.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataValidation.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/QualityCounters.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/ImportStateMachineTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/metadata/InMemoryMetadataStoreTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/metadata/MetadataModelTest.java diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ImportStateMachine.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ImportStateMachine.java new file mode 100644 index 000000000..a7aad805a --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ImportStateMachine.java @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import org.apache.geaflow.ai.retrieval.metadata.ImportState; +import org.apache.geaflow.ai.retrieval.metadata.MetadataException; + +/** Monotonic transitions within one attempt; retries create a new graph version. */ +public final class ImportStateMachine { + + private ImportStateMachine() { + } + + public static boolean canTransition(ImportState current, ImportState target) { + return current == ImportState.IMPORTING + && (target == ImportState.INDEXING || target == ImportState.FAILED) + || current == ImportState.INDEXING + && (target == ImportState.READY || target == ImportState.FAILED); + } + + public static void validate(ImportState current, ImportState target) { + if (!canTransition(current, target)) { + throw new MetadataException(MetadataException.Code.INVALID_TRANSITION, + "illegal import transition: " + current + " -> " + target); + } + } + + /** Validates and returns the next state for callers driving an import explicitly. */ + public static ImportState transition(ImportState current, ImportState target) { + validate(current, target); + return target; + } + + /** Alias used by adapters that name the operation as a transition validation. */ + public static void validateTransition(ImportState current, ImportState target) { + validate(current, target); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ChunkingConfiguration.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ChunkingConfiguration.java new file mode 100644 index 000000000..5f9f09b84 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ChunkingConfiguration.java @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.util.Objects; + +/** Reproducible character-based chunking policy. */ +public final class ChunkingConfiguration { + + private final String policyVersion; + private final int chunkSize; + private final int overlap; + + public ChunkingConfiguration( + String policyVersion, + int chunkSize, + int overlap) { + MetadataValidation.required(policyVersion, "policyVersion"); + if (chunkSize <= 0 || overlap < 0 || overlap >= chunkSize) { + throw MetadataValidation.invalid("chunkSize must be positive and 0 <= overlap < chunkSize"); + } + this.policyVersion = policyVersion; + this.chunkSize = chunkSize; + this.overlap = overlap; + } + + public String getPolicyVersion() { + return policyVersion; + } + + public String getStrategyVersion() { + return policyVersion; + } + + public int getChunkSize() { + return chunkSize; + } + + public int getOverlap() { + return overlap; + } + + public int getOverlapSize() { + return overlap; + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof ChunkingConfiguration)) { + return false; + } + ChunkingConfiguration that = (ChunkingConfiguration) object; + return chunkSize == that.chunkSize + && overlap == that.overlap + && Objects.equals(policyVersion, that.policyVersion); + } + + @Override + public int hashCode() { + return Objects.hash(policyVersion, chunkSize, overlap); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/DatasetManifest.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/DatasetManifest.java new file mode 100644 index 000000000..31a245e19 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/DatasetManifest.java @@ -0,0 +1,202 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.util.Objects; + +/** Immutable inputs needed to reproduce a dataset build. */ +public final class DatasetManifest { + + public static final String CURRENT_VERSION = "v1"; + + private final String manifestVersion; + private final String dataset; + private final String datasetRelease; + private final String split; + private final String sourceUri; + private final String cachePath; + private final String sha256; + private final String preprocessingVersion; + private final ChunkingConfiguration chunking; + private final String graphSchemaVersion; + private final String vectorSource; + private final String vectorVersion; + private final long randomSeed; + + public DatasetManifest( + String manifestVersion, + String dataset, + String datasetRelease, + String split, + String sourceUri, + String cachePath, + String sha256, + String preprocessingVersion, + ChunkingConfiguration chunking, + String graphSchemaVersion, + String vectorSource, + String vectorVersion, + long randomSeed) { + if (!CURRENT_VERSION.equals(manifestVersion)) { + throw MetadataValidation.invalid("unsupported manifestVersion"); + } + MetadataValidation.required(dataset, "dataset"); + MetadataValidation.required(datasetRelease, "datasetRelease"); + MetadataValidation.required(split, "split"); + if ((sourceUri == null || sourceUri.trim().isEmpty()) && (cachePath == null || cachePath.trim().isEmpty())) { + throw MetadataValidation.invalid("sourceUri or cachePath is required"); + } + sha256 = MetadataValidation.checksum(sha256); + MetadataValidation.required(preprocessingVersion, "preprocessingVersion"); + if (chunking == null) { + throw MetadataValidation.invalid("chunking is required"); + } + MetadataValidation.required(graphSchemaVersion, "graphSchemaVersion"); + if ((vectorSource == null) != (vectorVersion == null)) { + throw MetadataValidation.invalid("vectorSource and vectorVersion must be supplied together"); + } + if (vectorSource != null) { + MetadataValidation.required(vectorSource, "vectorSource"); + MetadataValidation.required(vectorVersion, "vectorVersion"); + } + this.manifestVersion = manifestVersion; + this.dataset = dataset; + this.datasetRelease = datasetRelease; + this.split = split; + this.sourceUri = sourceUri; + this.cachePath = cachePath; + this.sha256 = sha256; + this.preprocessingVersion = preprocessingVersion; + this.chunking = chunking; + this.graphSchemaVersion = graphSchemaVersion; + this.vectorSource = vectorSource; + this.vectorVersion = vectorVersion; + this.randomSeed = randomSeed; + } + + public String getManifestVersion() { + return manifestVersion; + } + + public String getDataset() { + return dataset; + } + + public String getDatasetRelease() { + return datasetRelease; + } + + public String getDatasetVersion() { + return datasetRelease; + } + + public String getRelease() { + return datasetRelease; + } + + public String getSplit() { + return split; + } + + public String getSourceUri() { + return sourceUri; + } + + public String getSourceUrl() { + return sourceUri; + } + + public String getCachePath() { + return cachePath; + } + + public String getSha256() { + return sha256; + } + + public String getChecksum() { + return sha256; + } + + public String getPreprocessingVersion() { + return preprocessingVersion; + } + + public ChunkingConfiguration getChunking() { + return chunking; + } + + public String getGraphSchemaVersion() { + return graphSchemaVersion; + } + + public String getVectorSource() { + return vectorSource; + } + + public String getVectorVersion() { + return vectorVersion; + } + + public long getRandomSeed() { + return randomSeed; + } + + /** Returns the dataset identifier using the terminology used by ingestion clients. */ + public String getDatasetId() { + return dataset; + } + + /** Returns the dataset name alias retained for manifest consumers. */ + public String getDatasetName() { + return dataset; + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof DatasetManifest)) { + return false; + } + DatasetManifest that = (DatasetManifest) object; + return randomSeed == that.randomSeed + && Objects.equals(manifestVersion, that.manifestVersion) + && Objects.equals(dataset, that.dataset) + && Objects.equals(datasetRelease, that.datasetRelease) + && Objects.equals(split, that.split) + && Objects.equals(sourceUri, that.sourceUri) + && Objects.equals(cachePath, that.cachePath) + && Objects.equals(sha256, that.sha256) + && Objects.equals(preprocessingVersion, that.preprocessingVersion) + && Objects.equals(chunking, that.chunking) + && Objects.equals(graphSchemaVersion, that.graphSchemaVersion) + && Objects.equals(vectorSource, that.vectorSource) + && Objects.equals(vectorVersion, that.vectorVersion); + } + + @Override + public int hashCode() { + return Objects.hash(manifestVersion, dataset, datasetRelease, split, sourceUri, cachePath, + sha256, preprocessingVersion, chunking, graphSchemaVersion, vectorSource, vectorVersion, + randomSeed); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/GraphBuildMetadata.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/GraphBuildMetadata.java new file mode 100644 index 000000000..38c7f40f4 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/GraphBuildMetadata.java @@ -0,0 +1,157 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.util.List; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.apache.geaflow.ai.retrieval.validation.ModelValidation; + +/** Graph artifact and the exact required index names for its publication. */ +public final class GraphBuildMetadata { + + private final GraphVersion graphVersion; + private final String artifactUri; + private final String graphType; + private final String builderVersion; + private final boolean ready; + private final List requiredIndexes; + + /** Creates graph metadata with the artifact details used by the current build. */ + public GraphBuildMetadata( + GraphVersion graphVersion, + String artifactUri, + String graphType, + String builderVersion, + boolean ready, + List requiredIndexes) { + this.graphVersion = Objects.requireNonNull(graphVersion, "graphVersion"); + this.artifactUri = artifactUri; + this.graphType = MetadataValidation.required(graphType, "graphType"); + this.builderVersion = MetadataValidation.required(builderVersion, "builderVersion"); + if (ready) { + MetadataValidation.required(artifactUri, "artifactUri"); + } + List names = ModelValidation.sortedStrings(requiredIndexes, "requiredIndexes"); + if (names.stream().distinct().count() != names.size()) { + throw MetadataValidation.invalid("requiredIndexes must be unique"); + } + this.ready = ready; + this.requiredIndexes = names; + } + + /** Creates graph metadata while retaining the original four-argument API. */ + public GraphBuildMetadata( + GraphVersion graphVersion, + String artifactUri, + boolean ready, + List requiredIndexes) { + this(graphVersion, artifactUri, "graph", "unknown", ready, requiredIndexes); + } + + /** Creates graph metadata without required indexes for a graph-only build. */ + public GraphBuildMetadata( + GraphVersion graphVersion, + String artifactUri, + String graphType, + String builderVersion, + boolean ready) { + this(graphVersion, artifactUri, graphType, builderVersion, ready, java.util.Collections.emptyList()); + } + + /** Creates graph metadata with the list before the readiness flag. */ + public GraphBuildMetadata( + GraphVersion graphVersion, + String artifactUri, + String graphType, + String builderVersion, + List requiredIndexes, + boolean ready) { + this(graphVersion, artifactUri, graphType, builderVersion, ready, requiredIndexes); + } + + public GraphVersion getGraphVersion() { + return graphVersion; + } + + public GraphVersion getVersion() { + return graphVersion; + } + + public String getArtifactUri() { + return artifactUri; + } + + public String getGraphType() { + return graphType; + } + + /** Alias for consumers that use a generic artifact type name. */ + public String getType() { + return graphType; + } + + public String getBuilderVersion() { + return builderVersion; + } + + public String getBuildVersion() { + return builderVersion; + } + + public boolean isReady() { + return ready; + } + + public boolean getReady() { + return ready; + } + + public List getRequiredIndexes() { + return requiredIndexes; + } + + /** Alias that makes the publication role explicit at call sites. */ + public List getRequiredIndexNames() { + return requiredIndexes; + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof GraphBuildMetadata)) { + return false; + } + GraphBuildMetadata that = (GraphBuildMetadata) object; + return ready == that.ready + && Objects.equals(graphVersion, that.graphVersion) + && Objects.equals(artifactUri, that.artifactUri) + && Objects.equals(graphType, that.graphType) + && Objects.equals(builderVersion, that.builderVersion) + && Objects.equals(requiredIndexes, that.requiredIndexes); + } + + @Override + public int hashCode() { + return Objects.hash(graphVersion, artifactUri, graphType, builderVersion, ready, requiredIndexes); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ImportMetadata.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ImportMetadata.java new file mode 100644 index 000000000..5ae5325d1 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ImportMetadata.java @@ -0,0 +1,223 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.validation.ModelValidation; + +/** Immutable attempt snapshot; timestamps are UTC epoch milliseconds. */ +public final class ImportMetadata { + + private final DatasetManifest manifest; + private final String importerVersion; + private final GraphBuildMetadata graph; + private final List indexes; + private final QualityCounters counters; + private final ImportState state; + private final String verifiedSha256; + private final MetadataException.Code failureCode; + private final String failureReason; + private final long startedAt; + private final long updatedAt; + + public ImportMetadata(DatasetManifest manifest, String importerVersion, GraphBuildMetadata graph, + List indexes, QualityCounters counters, ImportState state, + String verifiedSha256, MetadataException.Code failureCode, String failureReason, + long startedAt, long updatedAt) { + this.manifest = Objects.requireNonNull(manifest, "manifest"); + this.importerVersion = MetadataValidation.required(importerVersion, "importerVersion"); + this.graph = Objects.requireNonNull(graph, "graph"); + this.indexes = ModelValidation.immutableList(indexes, "indexes"); + this.counters = Objects.requireNonNull(counters, "counters"); + this.state = Objects.requireNonNull(state, "state"); + this.verifiedSha256 = verifiedSha256 == null ? null : MetadataValidation.checksum(verifiedSha256); + this.failureCode = failureCode; + this.failureReason = failureReason; + this.startedAt = MetadataValidation.nonNegative(startedAt, "startedAt"); + this.updatedAt = MetadataValidation.nonNegative(updatedAt, "updatedAt"); + if (updatedAt < startedAt) { + throw MetadataValidation.invalid("updatedAt precedes startedAt"); + } + if (this.verifiedSha256 != null && !manifest.getSha256().equals(this.verifiedSha256)) { + throw MetadataValidation.invalid("verifiedSha256 does not match manifest sha256"); + } + if (state == ImportState.FAILED) { + Objects.requireNonNull(failureCode, "failureCode"); + MetadataValidation.required(failureReason, "failureReason"); + } else if (failureCode != null || failureReason != null) { + throw MetadataValidation.invalid("failure is only allowed in FAILED"); + } + if (this.indexes.stream().map(index -> index.getIndexVersion().getIndexName()).distinct().count() + != this.indexes.size()) { + throw MetadataValidation.invalid("duplicate index names"); + } + for (IndexBuildMetadata index : this.indexes) { + if (!graph.getGraphVersion().equals(index.getGraphVersion())) { + throw MetadataValidation.invalid("index belongs to another graph"); + } + } + if (state == ImportState.IMPORTING && graph.isReady()) { + throw MetadataValidation.invalid("IMPORTING graph must not be ready"); + } + if (state == ImportState.INDEXING || state == ImportState.READY) { + if (!graph.isReady() || !manifest.getSha256().equals(this.verifiedSha256)) { + throw new MetadataException(MetadataException.Code.NOT_READY, + "verified source and completed graph required"); + } + } + if (state == ImportState.READY && !readinessReasons().isEmpty()) { + throw new MetadataException(MetadataException.Code.NOT_READY, readinessReasons().toString()); + } + } + + public List readinessReasons() { + List reasons = new ArrayList<>(); + if (state == ImportState.FAILED) { + reasons.add("FAILED: " + failureReason); + } + if (state == ImportState.IMPORTING) { + reasons.add("STATE_NOT_READY: IMPORTING"); + } + if (!manifest.getSha256().equals(verifiedSha256)) { + reasons.add("CHECKSUM_NOT_VERIFIED"); + } + if (!graph.isReady()) { + reasons.add("GRAPH_NOT_READY"); + } + for (String name : graph.getRequiredIndexes()) { + if (indexes.stream().noneMatch(index -> name.equals(index.getIndexVersion().getIndexName()) && index.isReady())) { + reasons.add("INDEX_NOT_READY: " + name); + } + } + return Collections.unmodifiableList(reasons); + } + + public DatasetManifest getManifest() { + return manifest; + } + + public String getImporterVersion() { + return importerVersion; + } + + public GraphBuildMetadata getGraph() { + return graph; + } + + public List getIndexes() { + return indexes; + } + + public QualityCounters getCounters() { + return counters; + } + + public ImportState getState() { + return state; + } + + public String getVerifiedSha256() { + return verifiedSha256; + } + + public MetadataException.Code getFailureCode() { + return failureCode; + } + + public String getFailureReason() { + return failureReason; + } + + public long getStartedAt() { + return startedAt; + } + + public long getUpdatedAt() { + return updatedAt; + } + + public GraphBuildMetadata getGraphMetadata() { + return graph; + } + + public GraphBuildMetadata getGraphBuildMetadata() { + return graph; + } + + public org.apache.geaflow.ai.retrieval.model.version.GraphVersion getGraphVersion() { + return graph.getGraphVersion(); + } + + public List getReadinessReasons() { + return readinessReasons(); + } + + public String getReadinessReason() { + List reasons = readinessReasons(); + return reasons.isEmpty() ? null : String.join(", ", reasons); + } + + public String getChecksum() { + return verifiedSha256; + } + + public long getCreatedAt() { + return startedAt; + } + + public long getLastUpdatedAt() { + return updatedAt; + } + + public boolean isReady() { + return state == ImportState.READY && readinessReasons().isEmpty(); + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof ImportMetadata)) { + return false; + } + ImportMetadata that = (ImportMetadata) object; + return startedAt == that.startedAt + && updatedAt == that.updatedAt + && Objects.equals(manifest, that.manifest) + && Objects.equals(importerVersion, that.importerVersion) + && Objects.equals(graph, that.graph) + && Objects.equals(indexes, that.indexes) + && Objects.equals(counters, that.counters) + && state == that.state + && Objects.equals(verifiedSha256, that.verifiedSha256) + && failureCode == that.failureCode + && Objects.equals(failureReason, that.failureReason); + } + + @Override + public int hashCode() { + return Objects.hash(manifest, importerVersion, graph, indexes, counters, state, + verifiedSha256, failureCode, failureReason, startedAt, updatedAt); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ImportState.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ImportState.java new file mode 100644 index 000000000..c91d2e080 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/ImportState.java @@ -0,0 +1,28 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +/** Lifecycle of a single import attempt. */ +public enum ImportState { + IMPORTING, + INDEXING, + READY, + FAILED +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/InMemoryMetadataStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/InMemoryMetadataStore.java new file mode 100644 index 000000000..58492ebaa --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/InMemoryMetadataStore.java @@ -0,0 +1,245 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.io.IOException; +import java.io.InputStream; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.time.Clock; +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.stream.Collectors; +import org.apache.geaflow.ai.retrieval.ingest.ImportStateMachine; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; + +/** Single-process reference implementation with atomic snapshot publication. */ +public final class InMemoryMetadataStore implements MetadataStore { + + private final Map attempts = new LinkedHashMap<>(); + private final Map published = new LinkedHashMap<>(); + private final Clock clock; + + public InMemoryMetadataStore() { + this(Clock.systemUTC()); + } + + public InMemoryMetadataStore(Clock clock) { + this.clock = Objects.requireNonNull(clock, "clock"); + } + + @Override + public synchronized ImportMetadata begin(DatasetManifest manifest, String importerVersion, GraphBuildMetadata graph) { + Objects.requireNonNull(manifest, "manifest"); + MetadataValidation.required(importerVersion, "importerVersion"); + Objects.requireNonNull(graph, "graph"); + GraphVersion version = graph.getGraphVersion(); + if (attempts.containsKey(version)) { + throw new MetadataException(MetadataException.Code.VERSION_CONFLICT, "attempt already exists"); + } + if (graph.isReady()) { + throw MetadataValidation.invalid("new attempt cannot have a ready graph"); + } + long now = clock.millis(); + ImportMetadata value = new ImportMetadata(manifest, importerVersion, graph, Collections.emptyList(), + new QualityCounters(0, 0, 0, 0, 0, 0), ImportState.IMPORTING, null, null, null, now, now); + attempts.put(version, value); + return value; + } + + @Override + public synchronized ImportMetadata retry(GraphVersion failedVersion, GraphBuildMetadata newGraph) { + Objects.requireNonNull(failedVersion, "failedVersion"); + Objects.requireNonNull(newGraph, "newGraph"); + ImportMetadata old = require(failedVersion); + if (old.getState() != ImportState.FAILED + || !failedVersion.getGraphName().equals(newGraph.getGraphVersion().getGraphName())) { + throw new MetadataException(MetadataException.Code.INVALID_TRANSITION, + "retry requires a failed attempt of the same graph"); + } + return begin(old.getManifest(), old.getImporterVersion(), newGraph); + } + + @Override + public synchronized ImportMetadata verifySource(GraphVersion version, InputStream source) throws IOException { + Objects.requireNonNull(source, "source"); + final ImportMetadata old = requireState(version, ImportState.IMPORTING); + String hash = digest(source); + if (!old.getManifest().getSha256().equals(hash)) { + fail(version, MetadataException.Code.CHECKSUM_MISMATCH, "source SHA-256 differs from manifest"); + throw new MetadataException(MetadataException.Code.CHECKSUM_MISMATCH, "source SHA-256 differs from manifest"); + } + return save(old, old.getGraph(), old.getIndexes(), old.getCounters(), old.getState(), hash, null, null); + } + + @Override + public synchronized ImportMetadata finishImport(GraphVersion version, GraphBuildMetadata graph, QualityCounters counters) { + Objects.requireNonNull(graph, "graph"); + Objects.requireNonNull(counters, "counters"); + ImportMetadata old = requireState(version, ImportState.IMPORTING); + if (!version.equals(graph.getGraphVersion()) + || !old.getGraph().getRequiredIndexes().equals(graph.getRequiredIndexes())) { + throw MetadataValidation.invalid("graph identity and required indexes cannot change within an attempt"); + } + if (!graph.isReady() || old.getVerifiedSha256() == null) { + throw new MetadataException(MetadataException.Code.NOT_READY, "verified source and completed graph required"); + } + ImportStateMachine.validate(old.getState(), ImportState.INDEXING); + return save(old, graph, old.getIndexes(), counters, ImportState.INDEXING, old.getVerifiedSha256(), null, null); + } + + @Override + public synchronized ImportMetadata putIndex(GraphVersion version, IndexBuildMetadata index) { + Objects.requireNonNull(index, "index"); + ImportMetadata old = requireState(version, ImportState.INDEXING); + if (!version.equals(index.getGraphVersion())) { + throw MetadataValidation.invalid("index belongs to another graph version"); + } + List indexes = new ArrayList<>(old.getIndexes()); + for (IndexBuildMetadata existing : indexes) { + if (existing.getIndexVersion().getIndexName().equals(index.getIndexVersion().getIndexName()) + && existing.isReady()) { + throw new MetadataException(MetadataException.Code.VERSION_CONFLICT, + "completed index is immutable"); + } + } + indexes.removeIf(existing -> existing.getIndexVersion().getIndexName() + .equals(index.getIndexVersion().getIndexName())); + indexes.add(index); + return save(old, old.getGraph(), indexes, old.getCounters(), old.getState(), + old.getVerifiedSha256(), null, null); + } + + @Override + public synchronized ImportMetadata fail(GraphVersion version, MetadataException.Code code, String reason) { + Objects.requireNonNull(code, "code"); + MetadataValidation.required(reason, "failureReason"); + ImportMetadata old = require(version); + ImportStateMachine.validate(old.getState(), ImportState.FAILED); + return save(old, old.getGraph(), old.getIndexes(), old.getCounters(), ImportState.FAILED, + old.getVerifiedSha256(), code, reason); + } + + @Override + public synchronized ImportMetadata publish(GraphVersion version, GraphVersion expectedPublishedVersion) { + ImportMetadata old = requireState(version, ImportState.INDEXING); + ImportMetadata current = published.get(version.getGraphName()); + GraphVersion actual = current == null ? null : current.getGraph().getGraphVersion(); + if (!Objects.equals(expectedPublishedVersion, actual)) { + throw new MetadataException(MetadataException.Code.VERSION_CONFLICT, "published version changed"); + } + if (!old.readinessReasons().isEmpty()) { + throw new MetadataException(MetadataException.Code.NOT_READY, + old.readinessReasons().toString()); + } + ImportStateMachine.validate(old.getState(), ImportState.READY); + ImportMetadata ready = save(old, old.getGraph(), old.getIndexes(), old.getCounters(), + ImportState.READY, old.getVerifiedSha256(), null, null); + published.put(version.getGraphName(), ready); + return ready; + } + + @Override + public synchronized Optional find(GraphVersion version) { + return Optional.ofNullable(attempts.get(version)); + } + + @Override + public synchronized Optional published(String graphName) { + return Optional.ofNullable(published.get(graphName)); + } + + @Override + public synchronized List findDataset(String dataset, String release, String split) { + return Collections.unmodifiableList(attempts.values().stream().filter(value -> + value.getManifest().getDataset().equals(dataset) + && value.getManifest().getDatasetRelease().equals(release) + && value.getManifest().getSplit().equals(split)).collect(Collectors.toList())); + } + + /** Returns the published version pointer without exposing mutable storage state. */ + public synchronized Optional publishedVersion(String graphName) { + ImportMetadata value = published.get(graphName); + return value == null ? Optional.empty() : Optional.of(value.getGraphVersion()); + } + + /** Alias for callers that use a get-style publication query. */ + public synchronized Optional getPublished(String graphName) { + return published(graphName); + } + + /** Alias for callers that use a find-style publication query. */ + public synchronized Optional findPublished(String graphName) { + return published(graphName); + } + + private ImportMetadata require(GraphVersion version) { + ImportMetadata value = attempts.get(version); + if (value == null) { + throw new MetadataException(MetadataException.Code.NOT_FOUND, "unknown graph version"); + } + return value; + } + + private ImportMetadata requireState(GraphVersion version, ImportState state) { + ImportMetadata value = require(version); + if (value.getState() != state) { + throw new MetadataException(MetadataException.Code.INVALID_TRANSITION, + "expected " + state + " but was " + value.getState()); + } + return value; + } + + private ImportMetadata save(ImportMetadata old, GraphBuildMetadata graph, + List indexes, QualityCounters counters, + ImportState state, String checksum, MetadataException.Code code, + String reason) { + ImportMetadata value = new ImportMetadata(old.getManifest(), old.getImporterVersion(), graph, indexes, + counters, state, checksum, code, reason, old.getStartedAt(), + Math.max(old.getUpdatedAt(), clock.millis())); + attempts.put(graph.getGraphVersion(), value); + return value; + } + + private static String digest(InputStream source) throws IOException { + MessageDigest digest; + try { + digest = MessageDigest.getInstance("SHA-256"); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException("SHA-256 unavailable", e); + } + byte[] buffer = new byte[8192]; + int length; + while ((length = source.read(buffer)) != -1) { + digest.update(buffer, 0, length); + } + StringBuilder hash = new StringBuilder(64); + for (byte value : digest.digest()) { + hash.append(Character.forDigit((value & 0xff) >>> 4, 16)); + hash.append(Character.forDigit(value & 0xf, 16)); + } + return hash.toString(); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/IndexBuildMetadata.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/IndexBuildMetadata.java new file mode 100644 index 000000000..feef524a2 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/IndexBuildMetadata.java @@ -0,0 +1,119 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.apache.geaflow.ai.retrieval.model.version.IndexVersion; + +/** Index artifact bound to a complete graph identity. */ +public final class IndexBuildMetadata { + + private final GraphVersion graphVersion; + private final IndexVersion indexVersion; + private final String indexType; + private final String builderVersion; + private final String artifactUri; + private final boolean ready; + + public IndexBuildMetadata( + GraphVersion graphVersion, + IndexVersion indexVersion, + String indexType, + String builderVersion, + String artifactUri, + boolean ready) { + Objects.requireNonNull(graphVersion, "graphVersion"); + Objects.requireNonNull(indexVersion, "indexVersion"); + if (!graphVersion.getVersion().equals(indexVersion.getGraphVersion())) { + throw MetadataValidation.invalid("index source graph version mismatch"); + } + MetadataValidation.required(indexType, "indexType"); + MetadataValidation.required(builderVersion, "builderVersion"); + if (ready) { + MetadataValidation.required(artifactUri, "artifactUri"); + } + this.graphVersion = graphVersion; + this.indexVersion = indexVersion; + this.indexType = indexType; + this.builderVersion = builderVersion; + this.artifactUri = artifactUri; + this.ready = ready; + } + + public GraphVersion getGraphVersion() { + return graphVersion; + } + + public IndexVersion getIndexVersion() { + return indexVersion; + } + + public String getIndexType() { + return indexType; + } + + public String getBuilderVersion() { + return builderVersion; + } + + public String getBuildVersion() { + return builderVersion; + } + + public String getArtifactUri() { + return artifactUri; + } + + /** Alias for callers that use the shorter artifact terminology. */ + public String getType() { + return indexType; + } + + public boolean isReady() { + return ready; + } + + public boolean getReady() { + return ready; + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof IndexBuildMetadata)) { + return false; + } + IndexBuildMetadata that = (IndexBuildMetadata) object; + return ready == that.ready + && Objects.equals(graphVersion, that.graphVersion) + && Objects.equals(indexVersion, that.indexVersion) + && Objects.equals(indexType, that.indexType) + && Objects.equals(builderVersion, that.builderVersion) + && Objects.equals(artifactUri, that.artifactUri); + } + + @Override + public int hashCode() { + return Objects.hash(graphVersion, indexVersion, indexType, builderVersion, artifactUri, ready); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataException.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataException.java new file mode 100644 index 000000000..0c808437e --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataException.java @@ -0,0 +1,48 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.util.Objects; + +/** Storage-neutral failure with a stable machine-readable code. */ +public final class MetadataException extends IllegalArgumentException { + + private static final long serialVersionUID = 1L; + + public enum Code { + INVALID_METADATA, + CHECKSUM_MISMATCH, + INVALID_TRANSITION, + NOT_READY, + VERSION_CONFLICT, + NOT_FOUND + } + + private final Code code; + + public MetadataException(Code code, String message) { + super(message); + this.code = Objects.requireNonNull(code, "code"); + } + + public Code getCode() { + return code; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataJson.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataJson.java new file mode 100644 index 000000000..3c50ad30b --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataJson.java @@ -0,0 +1,204 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import com.google.gson.Gson; +import com.google.gson.JsonElement; +import com.google.gson.JsonObject; +import com.google.gson.JsonParseException; +import com.google.gson.JsonParser; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.List; +import org.apache.geaflow.ai.retrieval.codec.RetrievalModelJson; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.apache.geaflow.ai.retrieval.model.version.IndexVersion; + +/** Constructor-validated JSON boundary; unknown optional fields are ignored. */ +public final class MetadataJson { + + private static final Gson GSON = new Gson(); + + private MetadataJson() { + } + + public static String toJson(Object value) { + return GSON.toJson(value); + } + + /** Alias for serializers that use a serialize-style name. */ + public static String serialize(Object value) { + return toJson(value); + } + + public static T fromJson(String json, Class type) { + if (json == null) { + throw new JsonParseException("metadata JSON is required"); + } + if (type == null) { + throw new JsonParseException("metadata type is required"); + } + try { + return type.cast(decode(new JsonParser().parse(json).getAsJsonObject(), type)); + } catch (RuntimeException e) { + throw new JsonParseException("invalid metadata: " + e.getMessage(), e); + } + } + + /** Alias for callers that use a parse-style name. */ + public static T parse(String json, Class type) { + return fromJson(json, type); + } + + private static Object decode(JsonObject object, Class type) { + if (type == ChunkingConfiguration.class) { + return new ChunkingConfiguration( + string(object, "policyVersion"), + Math.toIntExact(number(object, "chunkSize")), + Math.toIntExact(number(object, "overlap"))); + } + if (type == QualityCounters.class) { + return new QualityCounters( + number(object, "documents"), + number(object, "chunks"), + number(object, "entities"), + number(object, "edges"), + number(object, "skippedRecords"), + number(object, "validationErrors")); + } + if (type == DatasetManifest.class) { + return new DatasetManifest( + string(object, "manifestVersion"), + string(object, "dataset"), + string(object, "datasetRelease"), + string(object, "split"), + string(object, "sourceUri"), + string(object, "cachePath"), + string(object, "sha256"), + string(object, "preprocessingVersion"), + (ChunkingConfiguration) decode(object.getAsJsonObject("chunking"), ChunkingConfiguration.class), + string(object, "graphSchemaVersion"), + string(object, "vectorSource"), + string(object, "vectorVersion"), + number(object, "randomSeed")); + } + if (type == GraphBuildMetadata.class) { + return new GraphBuildMetadata( + RetrievalModelJson.fromJson(object.get("graphVersion").toString(), GraphVersion.class), + string(object, "artifactUri"), + stringOrDefault(object, "graphType", "graph"), + stringOrDefault(object, "builderVersion", "unknown"), + bool(object, "ready"), + strings(object, "requiredIndexes")); + } + if (type == IndexBuildMetadata.class) { + return new IndexBuildMetadata( + RetrievalModelJson.fromJson(object.get("graphVersion").toString(), GraphVersion.class), + RetrievalModelJson.fromJson(object.get("indexVersion").toString(), IndexVersion.class), + string(object, "indexType"), + string(object, "builderVersion"), + string(object, "artifactUri"), + bool(object, "ready")); + } + if (type == ImportMetadata.class) { + return new ImportMetadata( + (DatasetManifest) decode(object.getAsJsonObject("manifest"), DatasetManifest.class), + string(object, "importerVersion"), + (GraphBuildMetadata) decode(object.getAsJsonObject("graph"), GraphBuildMetadata.class), + indexes(object), + (QualityCounters) decode(object.getAsJsonObject("counters"), QualityCounters.class), + ImportState.valueOf(string(object, "state")), + string(object, "verifiedSha256"), + string(object, "failureCode") == null ? null : MetadataException.Code.valueOf(string(object, "failureCode")), + string(object, "failureReason"), + number(object, "startedAt"), + number(object, "updatedAt")); + } + throw new JsonParseException("unsupported metadata type: " + type); + } + + private static String string(JsonObject object, String name) { + JsonElement value = object.get(name); + if (value == null || value.isJsonNull()) { + return null; + } + if (!value.isJsonPrimitive() || !value.getAsJsonPrimitive().isString()) { + throw new JsonParseException(name + " must be a string"); + } + return value.getAsString(); + } + + private static long number(JsonObject object, String name) { + JsonElement value = object.get(name); + if (value == null || !value.isJsonPrimitive() || !value.getAsJsonPrimitive().isNumber()) { + throw new JsonParseException(name + " must be a number"); + } + return new BigDecimal(value.getAsString()).longValueExact(); + } + + private static boolean bool(JsonObject object, String name) { + JsonElement value = object.get(name); + if (value == null || !value.isJsonPrimitive() || !value.getAsJsonPrimitive().isBoolean()) { + throw new JsonParseException(name + " must be boolean"); + } + return value.getAsBoolean(); + } + + private static List strings(JsonObject object, String name) { + List result = new ArrayList<>(); + JsonElement element = object.get(name); + if (element == null || element.isJsonNull()) { + return result; + } + if (!element.isJsonArray()) { + throw new JsonParseException(name + " must be an array"); + } + for (JsonElement value : element.getAsJsonArray()) { + if (!value.isJsonPrimitive() || !value.getAsJsonPrimitive().isString()) { + throw new JsonParseException(name + " must contain strings"); + } + result.add(value.getAsString()); + } + return result; + } + + private static List indexes(JsonObject object) { + List result = new ArrayList<>(); + JsonElement element = object.get("indexes"); + if (element == null || element.isJsonNull()) { + return result; + } + if (!element.isJsonArray()) { + throw new JsonParseException("indexes must be an array"); + } + for (JsonElement value : element.getAsJsonArray()) { + if (!value.isJsonObject()) { + throw new JsonParseException("indexes must contain objects"); + } + result.add((IndexBuildMetadata) decode(value.getAsJsonObject(), IndexBuildMetadata.class)); + } + return result; + } + + private static String stringOrDefault(JsonObject object, String name, String defaultValue) { + String value = string(object, name); + return value == null ? defaultValue : value; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataStore.java new file mode 100644 index 000000000..bb0febf8c --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataStore.java @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.io.IOException; +import java.io.InputStream; +import java.util.List; +import java.util.Optional; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; + +/** + * Storage-neutral build lifecycle. Published snapshots are immutable. + * + *

Implementations must serialize mutations per attempt and atomically compare and replace the + * published graph pointer. Source streams remain owned by the caller. Retry uses a new graph + * version, preserving the failed attempt and any previously published version.

+ */ +public interface MetadataStore { + + ImportMetadata begin(DatasetManifest manifest, String importerVersion, GraphBuildMetadata graph); + + ImportMetadata retry(GraphVersion failedVersion, GraphBuildMetadata newGraph); + + ImportMetadata verifySource(GraphVersion version, InputStream source) throws IOException; + + ImportMetadata finishImport(GraphVersion version, GraphBuildMetadata graph, QualityCounters counters); + + ImportMetadata putIndex(GraphVersion version, IndexBuildMetadata index); + + ImportMetadata fail(GraphVersion version, MetadataException.Code code, String reason); + + ImportMetadata publish(GraphVersion version, GraphVersion expectedPublishedVersion); + + Optional find(GraphVersion version); + + Optional published(String graphName); + + List findDataset(String dataset, String release, String split); + + default List findByDataset(String dataset, String release, String split) { + return findDataset(dataset, release, split); + } + + /** Alias for implementations and callers that use get terminology. */ + default Optional get(GraphVersion version) { + return find(version); + } + + /** Returns the published snapshot for a graph name, if one exists. */ + default Optional getPublished(String graphName) { + return published(graphName); + } + + /** Reads only the published graph version while retaining snapshot immutability. */ + default Optional getPublishedVersion(String graphName) { + Optional snapshot = published(graphName); + return snapshot.map(ImportMetadata::getGraphVersion); + } + + default Optional publishedVersion(String graphName) { + return getPublishedVersion(graphName); + } + + default Optional readPublishedVersion(String graphName) { + return getPublishedVersion(graphName); + } + + /** Alias for stores that expose a find-style publication query. */ + default Optional findPublished(String graphName) { + return published(graphName); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataValidation.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataValidation.java new file mode 100644 index 000000000..c08ce337c --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/MetadataValidation.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.util.Locale; + +/** Validation shared by metadata constructors and their JSON boundary. */ +final class MetadataValidation { + + private MetadataValidation() { + } + + static String required(String value, String field) { + if (value == null || value.trim().isEmpty()) { + throw invalid(field + " is required"); + } + return value; + } + + static String checksum(String value) { + if (value == null || !value.matches("[a-fA-F0-9]{64}")) { + throw invalid("sha256 must contain 64 hexadecimal characters"); + } + return value.toLowerCase(Locale.ROOT); + } + + static long nonNegative(long value, String field) { + if (value < 0) { + throw invalid(field + " must be non-negative"); + } + return value; + } + + static MetadataException invalid(String message) { + return new MetadataException(MetadataException.Code.INVALID_METADATA, message); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/QualityCounters.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/QualityCounters.java new file mode 100644 index 000000000..79e013b7b --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/metadata/QualityCounters.java @@ -0,0 +1,124 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.util.Objects; + +/** Non-negative quality measurements for a build. */ +public final class QualityCounters { + + private final long documents; + private final long chunks; + private final long entities; + private final long edges; + private final long skippedRecords; + private final long validationErrors; + + public QualityCounters( + long documents, + long chunks, + long entities, + long edges, + long skippedRecords, + long validationErrors) { + MetadataValidation.nonNegative(documents, "documents"); + MetadataValidation.nonNegative(chunks, "chunks"); + MetadataValidation.nonNegative(entities, "entities"); + MetadataValidation.nonNegative(edges, "edges"); + MetadataValidation.nonNegative(skippedRecords, "skippedRecords"); + MetadataValidation.nonNegative(validationErrors, "validationErrors"); + this.documents = documents; + this.chunks = chunks; + this.entities = entities; + this.edges = edges; + this.skippedRecords = skippedRecords; + this.validationErrors = validationErrors; + } + + public long getDocuments() { + return documents; + } + + public long getChunks() { + return chunks; + } + + public long getEntities() { + return entities; + } + + public long getEdges() { + return edges; + } + + public long getSkippedRecords() { + return skippedRecords; + } + + public long getValidationErrors() { + return validationErrors; + } + + public long getDocumentCount() { + return documents; + } + + public long getChunkCount() { + return chunks; + } + + public long getEntityCount() { + return entities; + } + + public long getEdgeCount() { + return edges; + } + + public long getSkippedRecordCount() { + return skippedRecords; + } + + public long getValidationErrorCount() { + return validationErrors; + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof QualityCounters)) { + return false; + } + QualityCounters that = (QualityCounters) object; + return documents == that.documents + && chunks == that.chunks + && entities == that.entities + && edges == that.edges + && skippedRecords == that.skippedRecords + && validationErrors == that.validationErrors; + } + + @Override + public int hashCode() { + return Objects.hash(documents, chunks, entities, edges, skippedRecords, validationErrors); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/ImportStateMachineTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/ImportStateMachineTest.java new file mode 100644 index 000000000..7d219e3d7 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/ImportStateMachineTest.java @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import org.apache.geaflow.ai.retrieval.metadata.ImportState; +import org.apache.geaflow.ai.retrieval.metadata.MetadataException; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** State transition coverage for one immutable import attempt. */ +public class ImportStateMachineTest { + + @Test + public void allowsOnlyForwardBuildTransitionsAndFailure() { + Assertions.assertTrue(ImportStateMachine.canTransition(ImportState.IMPORTING, + ImportState.INDEXING)); + Assertions.assertTrue(ImportStateMachine.canTransition(ImportState.IMPORTING, + ImportState.FAILED)); + Assertions.assertTrue(ImportStateMachine.canTransition(ImportState.INDEXING, + ImportState.READY)); + Assertions.assertTrue(ImportStateMachine.canTransition(ImportState.INDEXING, + ImportState.FAILED)); + Assertions.assertEquals(ImportState.INDEXING, ImportStateMachine.transition( + ImportState.IMPORTING, ImportState.INDEXING)); + } + + @Test + public void rejectsTerminalAndBackwardTransitionsWithTypedError() { + ImportState[] states = ImportState.values(); + for (ImportState current : states) { + for (ImportState target : states) { + if (!ImportStateMachine.canTransition(current, target)) { + MetadataException exception = Assertions.assertThrows(MetadataException.class, + () -> ImportStateMachine.validate(current, target)); + Assertions.assertEquals(MetadataException.Code.INVALID_TRANSITION, + exception.getCode()); + } + } + } + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/metadata/InMemoryMetadataStoreTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/metadata/InMemoryMetadataStoreTest.java new file mode 100644 index 000000000..ce8a62e2c --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/metadata/InMemoryMetadataStoreTest.java @@ -0,0 +1,125 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import java.io.ByteArrayInputStream; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.Collections; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.apache.geaflow.ai.retrieval.model.version.IndexVersion; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** Lifecycle, readiness, CAS, and retry tests for the in-memory metadata store. */ +public class InMemoryMetadataStoreTest { + + private static final String SOURCE_SHA256 = "41cf6794ba4200b839c53531555f0f3998df4cbb01a4d5cb0b94e3ca5e23947d"; + + @Test + public void mismatchFailsAttemptAndRetryKeepsTheFailureRecord() throws Exception { + InMemoryMetadataStore store = new InMemoryMetadataStore(); + GraphVersion failedVersion = new GraphVersion("graph", "g1"); + store.begin(manifest(), "importer-v1", graph(failedVersion, false, Collections.emptyList())); + + MetadataException mismatch = Assertions.assertThrows(MetadataException.class, + () -> store.verifySource(failedVersion, input("wrong"))); + Assertions.assertEquals(MetadataException.Code.CHECKSUM_MISMATCH, mismatch.getCode()); + Assertions.assertEquals(ImportState.FAILED, store.find(failedVersion).get().getState()); + + GraphVersion retryVersion = new GraphVersion("graph", "g2"); + ImportMetadata retry = store.retry(failedVersion, + graph(retryVersion, false, Collections.emptyList())); + Assertions.assertEquals(ImportState.IMPORTING, retry.getState()); + Assertions.assertEquals(ImportState.FAILED, store.find(failedVersion).get().getState()); + Assertions.assertEquals(2, store.findDataset("dataset", "release", "dev").size()); + } + + @Test + public void publicationRequiresEveryRequiredIndexAndUsesCompareAndSet() throws Exception { + InMemoryMetadataStore store = new InMemoryMetadataStore(); + GraphVersion version = new GraphVersion("graph", "g1"); + store.begin(manifest(), "importer-v1", graph(version, false, + Collections.singletonList("keyword"))); + store.verifySource(version, input("source")); + store.finishImport(version, graph(version, true, Collections.singletonList("keyword")), + new QualityCounters(1, 1, 1, 0, 0, 0)); + + MetadataException notReady = Assertions.assertThrows(MetadataException.class, + () -> store.publish(version, null)); + Assertions.assertEquals(MetadataException.Code.NOT_READY, notReady.getCode()); + Assertions.assertFalse(store.published("graph").isPresent()); + + IndexVersion indexVersion = new IndexVersion("keyword", "i1", "g1"); + store.putIndex(version, new IndexBuildMetadata(version, indexVersion, "bm25", + "builder-v1", null, false)); + Assertions.assertThrows(MetadataException.class, () -> store.publish(version, null)); + store.putIndex(version, new IndexBuildMetadata(version, indexVersion, "bm25", + "builder-v1", "memory:i1", true)); + ImportMetadata published = store.publish(version, null); + Assertions.assertEquals(ImportState.READY, published.getState()); + Assertions.assertEquals(version, store.getPublishedVersion("graph").get()); + + GraphVersion nextVersion = new GraphVersion("graph", "g2"); + store.begin(manifest(), "importer-v1", graph(nextVersion, false, Collections.emptyList())); + store.verifySource(nextVersion, input("source")); + store.finishImport(nextVersion, graph(nextVersion, true, Collections.emptyList()), + new QualityCounters(1, 1, 1, 0, 0, 0)); + MetadataException conflict = Assertions.assertThrows(MetadataException.class, + () -> store.publish(nextVersion, new GraphVersion("graph", "stale"))); + Assertions.assertEquals(MetadataException.Code.VERSION_CONFLICT, conflict.getCode()); + Assertions.assertEquals(version, store.getPublishedVersion("graph").get()); + } + + @Test + public void completedIndexCannotBeReplaced() throws Exception { + InMemoryMetadataStore store = new InMemoryMetadataStore(); + GraphVersion version = new GraphVersion("graph", "g1"); + store.begin(manifest(), "importer-v1", graph(version, false, + Collections.singletonList("keyword"))); + store.verifySource(version, input("source")); + store.finishImport(version, graph(version, true, Collections.singletonList("keyword")), + new QualityCounters(1, 1, 1, 0, 0, 0)); + IndexVersion indexVersion = new IndexVersion("keyword", "i1", "g1"); + store.putIndex(version, new IndexBuildMetadata(version, indexVersion, "bm25", + "builder-v1", "memory:i1", true)); + + MetadataException conflict = Assertions.assertThrows(MetadataException.class, + () -> store.putIndex(version, new IndexBuildMetadata(version, + new IndexVersion("keyword", "i2", "g1"), "bm25", "builder-v2", "memory:i2", true))); + Assertions.assertEquals(MetadataException.Code.VERSION_CONFLICT, conflict.getCode()); + } + + private static DatasetManifest manifest() { + return new DatasetManifest("v1", "dataset", "release", "dev", "uri", null, + SOURCE_SHA256, "preprocess-v1", new ChunkingConfiguration("chunk-v1", 100, 10), + "schema-v1", null, null, 7); + } + + private static GraphBuildMetadata graph(GraphVersion version, boolean ready, + java.util.List requiredIndexes) { + return new GraphBuildMetadata(version, ready ? "memory:" + version.getVersion() : null, + "memory", "builder-v1", ready, requiredIndexes); + } + + private static ByteArrayInputStream input(String value) { + return new ByteArrayInputStream(value.getBytes(StandardCharsets.UTF_8)); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/metadata/MetadataModelTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/metadata/MetadataModelTest.java new file mode 100644 index 000000000..b371b98d6 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/metadata/MetadataModelTest.java @@ -0,0 +1,114 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.metadata; + +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.apache.geaflow.ai.retrieval.model.version.IndexVersion; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** Contract tests for immutable import metadata models and their JSON boundary. */ +public class MetadataModelTest { + + private static final String SHA256 = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + + @Test + public void manifestValidatesRequiredFieldsAndPreservesLongSeed() { + ChunkingConfiguration chunking = new ChunkingConfiguration("chunk-v1", 512, 64); + DatasetManifest manifest = new DatasetManifest("v1", "hotpotqa", "2025", "dev", + "https://example.test/data", null, SHA256.toUpperCase(), "clean-v1", chunking, + "schema-v1", "offline-v1", "embeddings-v2", Long.MIN_VALUE); + + Assertions.assertEquals("hotpotqa", manifest.getDataset()); + Assertions.assertEquals(-9223372036854775808L, manifest.getRandomSeed()); + Assertions.assertEquals(SHA256, manifest.getSha256()); + Assertions.assertEquals(manifest, + MetadataJson.fromJson(MetadataJson.toJson(manifest), DatasetManifest.class)); + } + + @Test + public void manifestRejectsInvalidRequiredFieldsAndPairs() { + ChunkingConfiguration chunking = new ChunkingConfiguration("chunk-v1", 10, 2); + Assertions.assertThrows(MetadataException.class, () -> new DatasetManifest("v2", "d", "r", + "dev", "uri", null, SHA256, "p", chunking, "schema", null, null, 1)); + Assertions.assertThrows(MetadataException.class, () -> new DatasetManifest("v1", "d", "r", + "dev", null, " ", SHA256, "p", chunking, "schema", null, null, 1)); + Assertions.assertThrows(MetadataException.class, () -> new DatasetManifest("v1", "d", "r", + "dev", "uri", null, "bad", "p", chunking, "schema", null, null, 1)); + Assertions.assertThrows(MetadataException.class, () -> new DatasetManifest("v1", "d", "r", + "dev", "uri", null, SHA256, "p", chunking, "schema", "vectors", null, 1)); + Assertions.assertThrows(MetadataException.class, () -> new ChunkingConfiguration("v1", 2, 2)); + } + + @Test + public void graphAndIndexMetadataValidateBindingAndGraphOnlyBuilds() { + GraphVersion graphVersion = new GraphVersion("graph", "g1"); + GraphBuildMetadata graph = new GraphBuildMetadata(graphVersion, "memory:g1", "memory", + "builder-v1", true, Collections.emptyList()); + Assertions.assertTrue(graph.getRequiredIndexes().isEmpty()); + Assertions.assertEquals(graph, MetadataJson.fromJson(MetadataJson.toJson(graph), + GraphBuildMetadata.class)); + + IndexVersion indexVersion = new IndexVersion("keyword", "i1", "g1"); + IndexBuildMetadata index = new IndexBuildMetadata(graphVersion, indexVersion, "bm25", + "builder-v1", "memory:i1", true); + Assertions.assertEquals(index, MetadataJson.fromJson(MetadataJson.toJson(index), + IndexBuildMetadata.class)); + Assertions.assertThrows(MetadataException.class, () -> new IndexBuildMetadata(graphVersion, + new IndexVersion("keyword", "i1", "other"), "bm25", "builder-v1", "uri", true)); + Assertions.assertThrows(MetadataException.class, () -> new GraphBuildMetadata(graphVersion, + "uri", "memory", "builder-v1", false, Arrays.asList("keyword", "keyword"))); + } + + @Test + public void importJsonIgnoresUnknownFieldsAndDefaultsOldIndexes() { + DatasetManifest manifest = manifest(); + GraphBuildMetadata graph = new GraphBuildMetadata(new GraphVersion("graph", "g1"), null, + false, Collections.emptyList()); + ImportMetadata metadata = new ImportMetadata(manifest, "importer-v1", graph, + Collections.emptyList(), new QualityCounters(1, 2, 3, 4, 5, 6), ImportState.IMPORTING, + null, null, null, 10, 10); + JsonObject json = new JsonParser().parse(MetadataJson.toJson(metadata)).getAsJsonObject(); + json.remove("indexes"); + json.addProperty("futureField", "ignored"); + + ImportMetadata restored = MetadataJson.fromJson(json.toString(), ImportMetadata.class); + Assertions.assertTrue(restored.getIndexes().isEmpty()); + Assertions.assertEquals(metadata.getCounters(), restored.getCounters()); + Assertions.assertEquals(ImportState.IMPORTING, restored.getState()); + } + + @Test + public void countersRejectNegativeValues() { + Assertions.assertThrows(MetadataException.class, + () -> new QualityCounters(0, 0, -1, 0, 0, 0)); + } + + private static DatasetManifest manifest() { + return new DatasetManifest("v1", "dataset", "release", "dev", "uri", null, SHA256, + "preprocess-v1", new ChunkingConfiguration("chunk-v1", 100, 10), "schema-v1", + null, null, 7); + } +} From 79a8451d3453426b2eecb660a7e4aa5ca6257482 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Fri, 11 Sep 2026 09:49:53 +0800 Subject: [PATCH 04/30] test(ai): add deterministic retrieval fixture and HTTP contract baseline --- .../geaflow/ai/GeaFlowMemoryServer.java | 2 +- .../apache/geaflow/ai/MemoryServerTest.java | 267 +++++------------- .../ai/retrieval/support/HttpTestClient.java | 92 ++++++ .../retrieval/support/HttpTestResponse.java | 56 ++++ .../support/RetrievalTestFixture.java | 64 +++++ 5 files changed, 287 insertions(+), 194 deletions(-) create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestClient.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestResponse.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/RetrievalTestFixture.java diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java index 5121ae7c9..a824e5ac5 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java @@ -57,7 +57,7 @@ public static void main(String[] args) { ctx.output("GeaFlow AI Server is running..."); }); app.get("/health", ctx -> { - ctx.output("{\"status\":\"UP\",\"service\":\"" + SERVER_NAME + "\"}"); + ctx.outputAsJson("{\"status\":\"UP\",\"service\":\"" + SERVER_NAME + "\"}"); }); }); } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java index 73db0a1af..22c9b1a69 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java @@ -19,212 +19,93 @@ package org.apache.geaflow.ai; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNotSame; +import static org.junit.jupiter.api.Assertions.assertTrue; + import com.google.gson.Gson; -import java.io.IOException; -import java.util.*; -import java.util.concurrent.TimeUnit; -import java.util.stream.Collectors; -import java.util.stream.IntStream; -import okhttp3.*; -import org.apache.geaflow.ai.common.config.Constants; -import org.apache.geaflow.ai.graph.io.*; -import org.junit.jupiter.api.*; +import com.google.gson.JsonObject; +import com.google.gson.JsonParser; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import org.apache.geaflow.ai.graph.io.Vertex; +import org.apache.geaflow.ai.retrieval.support.HttpTestClient; +import org.apache.geaflow.ai.retrieval.support.HttpTestResponse; +import org.apache.geaflow.ai.retrieval.support.RetrievalTestFixture; +import org.junit.jupiter.api.Test; import org.noear.solon.test.SolonTest; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; @SolonTest(GeaFlowMemoryServer.class) public class MemoryServerTest { - private static final Logger LOGGER = LoggerFactory.getLogger(MemoryServerTest.class); - private static final String BASE_URL = "http://localhost:8080"; - private static final String GRAPH_NAME = "Confucius"; - private static OkHttpClient client; - - @BeforeEach - void setUp() { - LOGGER.info("Setting up test environment..."); - if (client == null) { - OkHttpClient.Builder builder = new OkHttpClient.Builder(); - builder.callTimeout(Constants.HTTP_CALL_TIMEOUT_SECONDS, TimeUnit.SECONDS); - builder.connectTimeout(Constants.HTTP_CONNECT_TIMEOUT_SECONDS, TimeUnit.SECONDS); - builder.readTimeout(Constants.HTTP_READ_TIMEOUT_SECONDS, TimeUnit.SECONDS); - builder.writeTimeout(Constants.HTTP_WRITE_TIMEOUT_SECONDS, TimeUnit.SECONDS); - client = builder.build(); - } - LOGGER.info("Test HTTP client initialized, base URL: {}", BASE_URL); - } - - @AfterEach - void tearDown() { - LOGGER.info("Cleaning up test environment..."); - } - - private String get(String useApi) { - String url = BASE_URL + useApi; - Request request = new Request.Builder().url(url).get().build(); - try (okhttp3.Response response = client.newCall(request).execute()) { - return response.body().string(); - } catch (IOException e) { - throw new RuntimeException(e); - } - } - - private String post(String useApi, String bodyJson) { - return post(useApi, bodyJson, Collections.emptyMap()); - } - - private String post(String useApi, String bodyJson, Map queryParams) { - RequestBody requestBody = RequestBody.create( - MediaType.parse("application/json; charset=utf-8"), bodyJson); - String url = BASE_URL + useApi; - HttpUrl.Builder urlBuilder = HttpUrl.parse(url).newBuilder(); - if (queryParams != null) { - for (Map.Entry entry : queryParams.entrySet()) { - urlBuilder.addQueryParameter(entry.getKey(), entry.getValue()); - } - } - Request request = new Request.Builder().url(urlBuilder.build()).post(requestBody).build(); - try (okhttp3.Response response = client.newCall(request).execute()) { - return response.body().string(); - } catch (IOException e) { - throw new RuntimeException(e); - } - } - @Test - void testMain() throws Exception { - testServerHealth(); - testCreateGraph(); - testAddEntities(); - testQueries(); - } - - void testServerHealth() throws Exception { - LOGGER.info("Testing server health endpoint..."); - String api = "/"; - String response = get(api); - LOGGER.info("API: {} Response: {}", api, response); - api = "/health"; - response = get(api); - LOGGER.info("API: {} Response: {}", api, response); - api = "/api/test"; - response = get(api); - LOGGER.info("API: {} Response: {}", api, response); - } - - void testCreateGraph() throws Exception { - LOGGER.info("Testing server create graph..."); + void legacyWorkflow() { + RetrievalTestFixture fixture = new RetrievalTestFixture(); Gson gson = new Gson(); - String api = "/graph/create"; - GraphSchema testGraph = new GraphSchema(); - String graphName = GRAPH_NAME; - testGraph.setName(graphName); - String response = post(api, gson.toJson(testGraph)); - LOGGER.info("API: {} Response: {}", api, response); - - api = "/graph/getGraphSchema"; - Map queryParams = new HashMap<>(); - queryParams.put("graphName", graphName); - response = post(api, "", queryParams); - LOGGER.info("API: {} Response: {}", api, response); - Assertions.assertEquals(gson.toJson(testGraph), response); - - VertexSchema vertexSchema = new VertexSchema("chunk", "id", - Collections.singletonList("text")); - EdgeSchema edgeSchema = new EdgeSchema("relation", "srcId", "dstId", - Collections.singletonList("rel")); - testGraph.addVertex(vertexSchema); - testGraph.addEdge(edgeSchema); - - api = "/graph/addEntitySchema"; - queryParams = new HashMap<>(); - queryParams.put("graphName", graphName); - response = post(api, gson.toJson(vertexSchema), queryParams); - LOGGER.info("API: {} Response: {}", api, response); - - api = "/graph/addEntitySchema"; - queryParams = new HashMap<>(); - queryParams.put("graphName", graphName); - response = post(api, gson.toJson(edgeSchema), queryParams); - LOGGER.info("API: {} Response: {}", api, response); + Map graph = Collections.singletonMap("graphName", RetrievalTestFixture.GRAPH_NAME); + try (HttpTestClient client = new HttpTestClient("http://localhost:8080")) { + HttpTestResponse health = client.get("/health"); + success(health); + assertTrue(health.getContentType().startsWith("application/json")); + JsonObject body = new JsonParser().parse(health.getBody()).getAsJsonObject(); + assertEquals("UP", body.get("status").getAsString()); + assertEquals("geaflow-memory-server", body.get("service").getAsString()); + + HttpTestResponse created = client.post("/graph/create", gson.toJson(fixture.graphSchema()), Collections.emptyMap()); + success(created); + assertTrue(created.getBody().contains(RetrievalTestFixture.GRAPH_NAME)); + HttpTestResponse schema = client.post("/graph/addEntitySchema", gson.toJson(fixture.vertexSchema()), graph); + success(schema); + assertTrue(schema.getBody().contains("chunk")); + HttpTestResponse graphSchema = client.post("/graph/getGraphSchema", "", graph); + success(graphSchema); + assertTrue(graphSchema.getBody().contains("week1-keyword-graph")); + for (Vertex vertex : fixture.vertices()) { + HttpTestResponse inserted = client.post("/graph/insertEntity", gson.toJson(vertex), graph); + success(inserted); + assertEquals("Success to add entities, num: 1", inserted.getBody()); + } - api = "/graph/getGraphSchema"; - queryParams = new HashMap<>(); - queryParams.put("graphName", graphName); - response = post(api, "", queryParams); - LOGGER.info("API: {} Response: {}", api, response); - Assertions.assertEquals(gson.toJson(testGraph), response); + HttpTestResponse context = client.post("/query/context", "", graph); + success(context); + assertFalse(context.getBody().trim().isEmpty()); + Map session = Collections.singletonMap("sessionId", context.getBody()); + HttpTestResponse query = client.post("/query/exec", "Confucius", session); + success(query); + assertTrue(query.getBody().contains("Confucius"), query.getBody()); + HttpTestResponse result = client.post("/query/result", "", session); + success(result); + assertTrue(result.getBody().contains("confucius-1"), result.getBody()); + assertTrue(result.getBody().contains("Confucius taught ethics."), result.getBody()); + } } - void testAddEntities() throws Exception { - LOGGER.info("Testing server add entities..."); + @Test + void fixtureIsDeterministicAndIndependent() { Gson gson = new Gson(); - String graphName = GRAPH_NAME; - - String api = "/graph/getGraphSchema"; - Map queryParams = new HashMap<>(); - queryParams.put("graphName", graphName); - String response = post(api, "", queryParams); - LOGGER.info("API: {} Response: {}", api, response); - - TextFileReader textFileReader = new TextFileReader(10000); - textFileReader.readFile("text/Confucius"); - List chunks = IntStream.range(0, textFileReader.getRowCount()) - .mapToObj(textFileReader::getRow) - .map(String::trim).collect(Collectors.toList()); - for(String chunk : chunks) { - String vid = UUID.randomUUID().toString().replace("-", ""); - Vertex chunkVertex = new Vertex("chunk", vid, Collections.singletonList(chunk)); - api = "/graph/insertEntity"; - queryParams = new HashMap<>(); - queryParams.put("graphName", graphName); - response = post(api, gson.toJson(chunkVertex), queryParams); - LOGGER.info("API: {} Response: {}", api, response); - } + RetrievalTestFixture first = new RetrievalTestFixture(); + RetrievalTestFixture second = new RetrievalTestFixture(); + assertEquals(gson.toJson(first.graphSchema()), gson.toJson(second.graphSchema())); + assertEquals(gson.toJson(first.vertexSchema()), gson.toJson(second.vertexSchema())); + List firstVertices = first.vertices(); + List secondVertices = second.vertices(); + assertEquals(gson.toJson(firstVertices), gson.toJson(secondVertices)); + assertEquals(RetrievalTestFixture.CONFUCIUS_ID, firstVertices.get(0).getId()); + assertEquals(RetrievalTestFixture.ASTRONOMY_ID, firstVertices.get(1).getId()); + assertNotSame(first.graphSchema(), second.graphSchema()); + assertNotSame(first.vertexSchema(), second.vertexSchema()); + assertNotSame(firstVertices, secondVertices); + assertNotSame(firstVertices.get(0), secondVertices.get(0)); + first.graphSchema().setName("changed"); + assertEquals(RetrievalTestFixture.GRAPH_NAME, second.graphSchema().getName()); } - void testQueries() throws Exception { - LOGGER.info("Testing server queries..."); - String graphName = GRAPH_NAME; - String sessionId = null; - String api = "/query/context"; - Map queryParams = new HashMap<>(); - queryParams.put("graphName", graphName); - String response = post(api, "", queryParams); - LOGGER.info("API: {} Response: {}", api, response); - Assertions.assertNotNull(response); - sessionId = response; - - api = "/query/exec"; - queryParams = new HashMap<>(); - queryParams.put("sessionId", sessionId); - queryParams.put("query", "Who is Confucius?"); - response = post(api, "", queryParams); - LOGGER.info("API: {} Response: {}", api, response); - Assertions.assertNotNull(response); - - api = "/query/result"; - queryParams = new HashMap<>(); - queryParams.put("sessionId", sessionId); - response = post(api, "", queryParams); - LOGGER.info("API: {} Response: {}", api, response); - Assertions.assertNotNull(response); - - api = "/query/exec"; - queryParams = new HashMap<>(); - queryParams.put("sessionId", sessionId); - queryParams.put("query", "What did he say?"); - response = post(api, "", queryParams); - LOGGER.info("API: {} Response: {}", api, response); - Assertions.assertNotNull(response); - - api = "/query/result"; - queryParams = new HashMap<>(); - queryParams.put("sessionId", sessionId); - response = post(api, "", queryParams); - LOGGER.info("API: {} Response: {}", api, response); - Assertions.assertNotNull(response); + private static void success(HttpTestResponse response) { + assertEquals(200, response.getStatus(), response.getBody()); + assertNotNull(response.getHeaders().get("Content-Type")); + assertNotNull(response.getBody()); } - } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestClient.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestClient.java new file mode 100644 index 000000000..707bc6b01 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestClient.java @@ -0,0 +1,92 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.support; + +import java.io.IOException; +import java.util.Collections; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.TimeUnit; +import okhttp3.HttpUrl; +import okhttp3.MediaType; +import okhttp3.OkHttpClient; +import okhttp3.Request; +import okhttp3.RequestBody; +import okhttp3.Response; + +/** Loopback HTTP client which never retains live response resources. */ +public final class HttpTestClient implements AutoCloseable { + + private final String baseUrl; + private final OkHttpClient client = new OkHttpClient.Builder() + .callTimeout(10, TimeUnit.SECONDS) + .build(); + + public HttpTestClient(String baseUrl) { + Objects.requireNonNull(baseUrl, "baseUrl"); + if (baseUrl.trim().isEmpty()) { + throw new IllegalArgumentException("baseUrl must not be blank"); + } + this.baseUrl = baseUrl.endsWith("/") + ? baseUrl.substring(0, baseUrl.length() - 1) : baseUrl; + } + + public HttpTestResponse get(String path) { + return get(path, Collections.emptyMap()); + } + + public HttpTestResponse get(String path, Map parameters) { + return execute(new Request.Builder().url(url(path, parameters)).get().build()); + } + + public HttpTestResponse post(String path, String body, Map parameters) { + return execute(new Request.Builder().url(url(path, parameters)).post( + RequestBody.create(MediaType.parse("application/json; charset=utf-8"), + body == null ? "" : body)).build()); + } + + public HttpTestResponse post(String path, String body) { + return post(path, body, Collections.emptyMap()); + } + + private HttpUrl url(String path, Map parameters) { + Objects.requireNonNull(path, "path"); + String normalizedPath = path.startsWith("/") ? path : "/" + path; + HttpUrl.Builder url = HttpUrl.get(baseUrl + normalizedPath).newBuilder(); + (parameters == null ? Collections.emptyMap() : parameters) + .forEach(url::addQueryParameter); + return url.build(); + } + + private HttpTestResponse execute(Request request) { + try (Response response = client.newCall(request).execute()) { + return new HttpTestResponse(response.code(), response.headers().newBuilder().build(), + response.body() == null ? "" : response.body().string()); + } catch (IOException e) { + throw new IllegalStateException(request.method() + " " + request.url(), e); + } + } + + @Override + public void close() { + client.connectionPool().evictAll(); + client.dispatcher().executorService().shutdown(); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestResponse.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestResponse.java new file mode 100644 index 000000000..a08435427 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestResponse.java @@ -0,0 +1,56 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.support; + +import okhttp3.Headers; + +/** Detached HTTP result, safe to inspect after the response has closed. */ +public final class HttpTestResponse { + + private final int status; + private final Headers headers; + private final String body; + + public HttpTestResponse(int status, Headers headers, String body) { + this.status = status; + this.headers = headers; + this.body = body; + } + + public int getStatus() { + return status; + } + + public int getStatusCode() { + return status; + } + + public Headers getHeaders() { + return headers; + } + + public String getBody() { + return body; + } + + public String getContentType() { + return headers.get("Content-Type"); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/RetrievalTestFixture.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/RetrievalTestFixture.java new file mode 100644 index 000000000..3090bda33 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/RetrievalTestFixture.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.support; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.geaflow.ai.graph.io.GraphSchema; +import org.apache.geaflow.ai.graph.io.Vertex; +import org.apache.geaflow.ai.graph.io.VertexSchema; + +/** Two independent facts with fixed identity and insertion order. */ +public final class RetrievalTestFixture { + + public static final String GRAPH_NAME = "week1-keyword-graph"; + public static final String CHUNK_LABEL = "chunk"; + public static final String CONFUCIUS_ID = "confucius-1"; + public static final String ASTRONOMY_ID = "astronomy-1"; + + public GraphSchema graphSchema() { + GraphSchema schema = new GraphSchema(); + schema.setName(GRAPH_NAME); + return schema; + } + + public GraphSchema getGraphSchema() { + return graphSchema(); + } + + public VertexSchema vertexSchema() { + return new VertexSchema(CHUNK_LABEL, "id", Collections.singletonList("text")); + } + + public VertexSchema getVertexSchema() { + return vertexSchema(); + } + + public List vertices() { + return Arrays.asList( + new Vertex(CHUNK_LABEL, CONFUCIUS_ID, Collections.singletonList("Confucius taught ethics.")), + new Vertex(CHUNK_LABEL, ASTRONOMY_ID, Collections.singletonList("Astronomy studies stars."))); + } + + public List getVertices() { + return vertices(); + } +} From 07bbd25c0150af183009660470f4a91a4f0514e1 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 15:43:39 +0800 Subject: [PATCH 05/30] chore(ai): preserve week1 retrieval baseline --- geaflow-ai/README-week1-retrieval.md | 41 +++ .../geaflow/ai/GeaFlowMemoryServer.java | 188 ++++++++++++-- .../apache/geaflow/ai/GraphMemoryServer.java | 52 ++++ .../ai/index/EntityAttributeIndexStore.java | 9 +- .../geaflow/ai/operator/GraphSearchStore.java | 100 +++++++- .../geaflow/ai/operator/SearchStore.java | 23 +- .../geaflow/ai/operator/SessionOperator.java | 17 +- .../retrieval/codec/RetrievalModelJson.java | 3 +- .../ai/retrieval/index/Bm25IndexBuilder.java | 31 +++ .../ai/retrieval/index/IndexArtifact.java | 29 +++ .../retrieval/index/VectorIndexBuilder.java | 31 +++ .../ai/retrieval/index/package-info.java | 26 ++ .../retrieval/ingest/DocumentNormalizer.java | 32 +++ .../ai/retrieval/ingest/EntityExtractor.java | 30 +++ .../ai/retrieval/ingest/ExtractionResult.java | 54 ++++ .../ai/retrieval/ingest/GraphArtifact.java | 29 +++ .../ai/retrieval/ingest/GraphWriter.java | 32 +++ .../ai/retrieval/ingest/IngestionContext.java | 68 +++++ .../retrieval/ingest/IngestionPipeline.java | 133 ++++++++++ .../retrieval/ingest/MetadataPublisher.java | 36 +++ .../ai/retrieval/ingest/SourceLoader.java | 30 +++ .../ai/retrieval/ingest/package-info.java | 28 +++ .../ai/retrieval/model/evidence/Evidence.java | 24 +- .../ai/retrieval/service/EvidenceMapper.java | 76 ++++++ .../retrieval/service/RetrievalMetrics.java | 99 ++++++++ .../retrieval/service/RetrievalService.java | 155 ++++++++++++ .../geaflow/ai/service/ServerMemoryCache.java | 104 +++++++- geaflow-ai/src/main/resources/application.yml | 2 +- .../geaflow/ai/GeaFlowMemoryServerTest.java | 45 ++++ .../ai/GraphMemoryKeywordSearchTest.java | 86 +++++++ .../apache/geaflow/ai/MemoryServerTest.java | 42 +++- .../geaflow/ai/operator/SearchStoreTest.java | 46 ++++ .../ingest/IngestionPipelineTest.java | 233 ++++++++++++++++++ .../service/RetrievalServiceTest.java | 143 +++++++++++ .../ai/retrieval/support/HttpTestClient.java | 9 + 35 files changed, 2037 insertions(+), 49 deletions(-) create mode 100644 geaflow-ai/README-week1-retrieval.md create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/Bm25IndexBuilder.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/IndexArtifact.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/VectorIndexBuilder.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/package-info.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/DocumentNormalizer.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/EntityExtractor.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ExtractionResult.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/GraphArtifact.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/GraphWriter.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionContext.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataPublisher.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/SourceLoader.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/package-info.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/EvidenceMapper.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalMetrics.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalService.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/GeaFlowMemoryServerTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/GraphMemoryKeywordSearchTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/SearchStoreTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/RetrievalServiceTest.java diff --git a/geaflow-ai/README-week1-retrieval.md b/geaflow-ai/README-week1-retrieval.md new file mode 100644 index 000000000..590868824 --- /dev/null +++ b/geaflow-ai/README-week1-retrieval.md @@ -0,0 +1,41 @@ +# Week 1 Retrieval API + +Build the memory server; run its `GeaFlowMemoryServer` main class with the project runtime classpath: + +```bash +GEAFLOW_SERVER_PORT=8080 mvn -pl geaflow-ai -DskipTests package +``` + +Create a small keyword graph and load one vertex: + +```bash +curl -sS -X POST http://localhost:8080/graph/create \ + -H 'Content-Type: application/json' \ + -d '{"graphName":"Confucius","vertexSchemaList":[],"edgeSchemaList":[]}' + +curl -sS -X POST 'http://localhost:8080/graph/addEntitySchema?graphName=Confucius' \ + -H 'Content-Type: application/json' \ + -d '{"label":"chunk","idField":"id","fields":["text"]}' + +curl -sS -X POST 'http://localhost:8080/graph/insertEntity?graphName=Confucius' \ + -H 'Content-Type: application/json' \ + -d '{"label":"chunk","id":"confucius-1","values":["Confucius taught ethics."]}' +``` + +Run bounded keyword retrieval without creating a session: + +```bash +curl -sS -X POST http://localhost:8080/api/v1/retrievals \ + -H 'Content-Type: application/json' \ + -H 'X-Request-Id: curl-example-1' \ + -d '{"graphName":"Confucius","query":"Confucius","mode":"KEYWORD",\ +"budget":{"topK":10,"timeoutMs":3000,"maxCandidates":100,"tokenBudget":4096}}' +``` + +Health and readiness are separate checks: + +```bash +curl -i http://localhost:8080/health +curl -i http://localhost:8080/ready +curl -sS http://localhost:8080/metrics/retrieval +``` diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java index a824e5ac5..d9967240b 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java @@ -23,11 +23,18 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.regex.Pattern; import org.apache.geaflow.ai.common.util.SeDeUtil; import org.apache.geaflow.ai.graph.*; import org.apache.geaflow.ai.graph.io.*; import org.apache.geaflow.ai.index.EntityAttributeIndexStore; import org.apache.geaflow.ai.index.vector.KeywordVector; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalError; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.codec.RetrievalApiJson; +import org.apache.geaflow.ai.retrieval.config.RetrievalProperties; +import org.apache.geaflow.ai.retrieval.service.RetrievalMetrics; +import org.apache.geaflow.ai.retrieval.service.RetrievalService; import org.apache.geaflow.ai.search.VectorSearch; import org.apache.geaflow.ai.service.ServerMemoryCache; import org.apache.geaflow.ai.verbalization.Context; @@ -44,36 +51,174 @@ public class GeaFlowMemoryServer { private static final String SERVER_NAME = "geaflow-memory-server"; private static final int DEFAULT_PORT = 8080; + private static final Pattern SAFE_REQUEST_ID = Pattern.compile("[A-Za-z0-9._:-]{1,128}"); - private static final ServerMemoryCache CACHE = new ServerMemoryCache(); + private static final ServerMemoryCache FALLBACK_CACHE = new ServerMemoryCache(); + private static final RetrievalMetrics FALLBACK_METRICS = new RetrievalMetrics(); + + @Inject + private ServerMemoryCache cache; + + @Inject + private RetrievalMetrics metrics; + + @Inject + private RetrievalProperties retrievalProperties; + + @Inject + private RetrievalService retrievalService; + + private static final RetrievalService FALLBACK_SERVICE = new RetrievalService( + FALLBACK_CACHE, new RetrievalProperties(), FALLBACK_METRICS); public static void main(String[] args) { System.setProperty("solon.app.name", SERVER_NAME); + int port = configuredPort(); + System.setProperty("server.port", String.valueOf(port)); Solon.start(GeaFlowMemoryServer.class, args, app -> { app.cfg().loadAdd("application.yml"); - int port = app.cfg().getInt("server.port", DEFAULT_PORT); + app.cfg().put("server.port", port); LOGGER.info("Starting {} on port {}", SERVER_NAME, port); app.get("/", ctx -> { ctx.output("GeaFlow AI Server is running..."); }); - app.get("/health", ctx -> { - ctx.outputAsJson("{\"status\":\"UP\",\"service\":\"" + SERVER_NAME + "\"}"); - }); }); } + private static int configuredPort() { + String environmentPort = System.getenv("GEAFLOW_SERVER_PORT"); + if (environmentPort == null || environmentPort.trim().isEmpty()) { + return DEFAULT_PORT; + } + try { + return Integer.parseInt(environmentPort.trim()); + } catch (NumberFormatException ignored) { + LOGGER.warn("Ignoring invalid GEAFLOW_SERVER_PORT: {}", environmentPort); + return DEFAULT_PORT; + } + } + @Get @Mapping("/api/test") public String test() { return "GeaFlow Memory Server is working!"; } + @Post + @Mapping("/api/v1/retrievals") + public String retrieve(org.noear.solon.core.handle.Context ctx, @Body String input) { + String requestId = resolveRequestId(ctx.header("X-Request-Id")); + long startedAt = System.nanoTime(); + ctx.headerSet("X-Request-Id", requestId); + ctx.contentType("application/json; charset=utf-8"); + try { + String body = RetrievalApiJson.toJson( + runtimeRetrievalService().retrieve(RetrievalApiJson.parseRequest(input), requestId)); + ctx.status(200); + return body; + } catch (RetrievalException exception) { + ctx.status(exception.getCode().getHttpStatus()); + RetrievalError error = new RetrievalError(requestId, exception.getCode(), + exception.getMessage()); + return RetrievalApiJson.toJson(error); + } catch (com.google.gson.JsonParseException exception) { + ctx.status(400); + runtimeMetrics().recordFailure( + org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.INVALID_REQUEST, + elapsedMs(startedAt)); + RetrievalError error = new RetrievalError(requestId, + org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.INVALID_REQUEST, + "Malformed JSON request"); + return RetrievalApiJson.toJson(error); + } catch (RuntimeException exception) { + ctx.status(500); + runtimeMetrics().recordFailure( + org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.INTERNAL_ERROR, + elapsedMs(startedAt)); + RetrievalError error = new RetrievalError(requestId, + org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.INTERNAL_ERROR, + "internal retrieval error"); + return RetrievalApiJson.toJson(error); + } + } + + static String resolveRequestId(String requestedId) { + if (requestedId == null) { + return java.util.UUID.randomUUID().toString(); + } + String trimmed = requestedId.trim(); + if (!SAFE_REQUEST_ID.matcher(trimmed).matches()) { + return java.util.UUID.randomUUID().toString(); + } + return trimmed; + } + + private static long elapsedMs(long startedAt) { + return (System.nanoTime() - startedAt) / 1_000_000L; + } + + @Get + @Mapping("/health") + public String health(org.noear.solon.core.handle.Context ctx) { + ctx.status(200); + ctx.contentType("application/json; charset=utf-8"); + return "{\"status\":\"UP\",\"service\":\"" + SERVER_NAME + "\"}"; + } + + @Get + @Mapping("/ready") + public String ready(org.noear.solon.core.handle.Context ctx) { + RetrievalProperties properties = properties(); + ServerMemoryCache.ReadinessStatus status; + try { + properties.validateConfiguration(); + status = runtimeCache().keywordReadiness(properties.getReadyGraphName()); + } catch (RetrievalException exception) { + status = new ServerMemoryCache.ReadinessStatus(false, null, null, + "INVALID_CONFIGURATION"); + } + ctx.status(status.isReady() ? 200 : 503); + ctx.contentType("application/json; charset=utf-8"); + return "{\"status\":\"" + (status.isReady() ? "READY" : "NOT_READY") + + "\",\"graphName\":\"" + String.valueOf(status.getGraphName()) + + "\",\"graphVersion\":\"" + String.valueOf(status.getGraphVersion()) + + "\",\"reason\":\"" + status.getReason() + "\"}"; + } + + @Get + @Mapping("/metrics/retrieval") + public String metrics(org.noear.solon.core.handle.Context ctx) { + RetrievalMetrics.Snapshot snapshot = runtimeMetrics().snapshot(); + ctx.status(200); + ctx.contentType("application/json; charset=utf-8"); + return "{\"total\":" + snapshot.getTotal() + ",\"success\":" + + snapshot.getSuccess() + ",\"timeout\":" + snapshot.getTimeout() + + ",\"failure\":" + snapshot.getFailure() + ",\"totalElapsedMs\":" + + snapshot.getTotalElapsedMs() + "}"; + } + + private RetrievalProperties properties() { + return retrievalProperties == null ? new RetrievalProperties() : retrievalProperties; + } + + private ServerMemoryCache runtimeCache() { + return cache == null ? FALLBACK_CACHE : cache; + } + + private RetrievalMetrics runtimeMetrics() { + return metrics == null ? FALLBACK_METRICS : metrics; + } + + private RetrievalService runtimeRetrievalService() { + return retrievalService == null ? FALLBACK_SERVICE : retrievalService; + } + @Post @Mapping("/graph/create") public String createGraph(@Body String input) { GraphSchema graphSchema = SeDeUtil.deserializeGraphSchema(input); String graphName = graphSchema.getName(); - if (graphName == null || CACHE.getGraphByName(graphName) != null) { + if (graphName == null || runtimeCache().getGraphByName(graphName) != null) { throw new RuntimeException("Cannot create graph name: " + graphName); } Map entities = new HashMap<>(); @@ -84,7 +229,6 @@ public String createGraph(@Body String input) { entities.put(edgeSchema.getName(), new EdgeGroup(edgeSchema, new ArrayList<>())); } MemoryGraph graph = new MemoryGraph(graphSchema, entities); - CACHE.putGraph(graph); LocalMemoryGraphAccessor graphAccessor = new LocalMemoryGraphAccessor(graph); LOGGER.info("Success to init empty graph."); @@ -96,7 +240,8 @@ public String createGraph(@Body String input) { server.addGraphAccessor(graphAccessor); server.addIndexStore(indexStore); LOGGER.info("Success to init GraphMemoryServer."); - CACHE.putServer(server); + runtimeCache().putGraph(graph); + runtimeCache().putServer(server); LOGGER.info("Success to init graph. SCHEMA: {}", graphSchema); return "createGraph has been called, graphName: " + graphName; @@ -106,7 +251,7 @@ public String createGraph(@Body String input) { @Mapping("/graph/addEntitySchema") public String addSchema(@Param("graphName") String graphName, @Body String input) { - Graph graph = CACHE.getGraphByName(graphName); + Graph graph = runtimeCache().getGraphByName(graphName); if (graph == null) { throw new RuntimeException("Graph not exist."); } @@ -123,13 +268,14 @@ public String addSchema(@Param("graphName") String graphName, } else { throw new RuntimeException("Cannot add schema: " + input); } + runtimeCache().markGraphUpdated(graphName); return "addSchema has been called, schemaName: " + schemaName; } @Post @Mapping("/graph/getGraphSchema") public String getSchema(@Param("graphName") String graphName) { - Graph graph = CACHE.getGraphByName(graphName); + Graph graph = runtimeCache().getGraphByName(graphName); if (graph == null) { throw new RuntimeException("Graph not exist."); } @@ -143,7 +289,7 @@ public String getSchema(@Param("graphName") String graphName) { @Mapping("/graph/insertEntity") public String addEntity(@Param("graphName") String graphName, @Body String input) { - Graph graph = CACHE.getGraphByName(graphName); + Graph graph = runtimeCache().getGraphByName(graphName); if (graph == null) { throw new RuntimeException("Graph not exist."); } @@ -160,12 +306,13 @@ public String addEntity(@Param("graphName") String graphName, memoryMutableGraph.addEdge(((GraphEdge) entity).getEdge()); } } - GraphMemoryServer insertServer = CACHE.getServerByName(graphName); + GraphMemoryServer insertServer = runtimeCache().getServerByName(graphName); if (insertServer == null || insertServer.getGraphAccessors().isEmpty()) { throw new RuntimeException("Server or graph accessor not available for graph: " + graphName); } - CACHE.getConsolidateServer().executeConsolidateTask( + runtimeCache().getConsolidateServer().executeConsolidateTask( insertServer.getGraphAccessors().get(0), memoryMutableGraph); + runtimeCache().markGraphUpdated(graphName); return "Success to add entities, num: " + graphEntities.size(); } @@ -173,7 +320,7 @@ public String addEntity(@Param("graphName") String graphName, @Mapping("/graph/delEntity") public String deleteEntity(@Param("graphName") String graphName, @Body String input) { - Graph graph = CACHE.getGraphByName(graphName); + Graph graph = runtimeCache().getGraphByName(graphName); if (graph == null) { throw new RuntimeException("Graph not exist."); } @@ -190,18 +337,19 @@ public String deleteEntity(@Param("graphName") String graphName, memoryMutableGraph.removeEdge(((GraphEdge) entity).getEdge()); } } + runtimeCache().markGraphUpdated(graphName); return "Success to remove entities, num: " + graphEntities.size(); } @Post @Mapping("/query/context") public String createContext(@Param("graphName") String graphName) { - GraphMemoryServer server = CACHE.getServerByName(graphName); + GraphMemoryServer server = runtimeCache().getServerByName(graphName); if (server == null) { throw new RuntimeException("Server not exist."); } String sessionId = server.createSession(); - CACHE.putSession(server, sessionId); + runtimeCache().putSession(server, sessionId); return sessionId; } @@ -209,11 +357,11 @@ public String createContext(@Param("graphName") String graphName) { @Mapping("/query/exec") public String execQuery(@Param("sessionId") String sessionId, @Body String query) { - String graphName = CACHE.getGraphNameBySession(sessionId); + String graphName = runtimeCache().getGraphNameBySession(sessionId); if (graphName == null) { throw new RuntimeException("Graph not exist."); } - GraphMemoryServer server = CACHE.getServerByName(graphName); + GraphMemoryServer server = runtimeCache().getServerByName(graphName); VectorSearch search = new VectorSearch(null, sessionId); search.addVector(new KeywordVector(query)); server.search(search); @@ -228,11 +376,11 @@ public String execQuery(@Param("sessionId") String sessionId, @Post @Mapping("/query/result") public String getResult(@Param("sessionId") String sessionId) { - String graphName = CACHE.getGraphNameBySession(sessionId); + String graphName = runtimeCache().getGraphNameBySession(sessionId); if (graphName == null) { throw new RuntimeException("Graph not exist."); } - GraphMemoryServer server = CACHE.getServerByName(graphName); + GraphMemoryServer server = runtimeCache().getServerByName(graphName); List result = server.getSessionEntities(sessionId); return result.toString(); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java index 57a8d37a2..02d29087f 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java @@ -26,10 +26,13 @@ import java.util.stream.Collectors; import org.apache.geaflow.ai.graph.GraphAccessor; import org.apache.geaflow.ai.graph.GraphEntity; +import org.apache.geaflow.ai.graph.GraphVertex; import org.apache.geaflow.ai.index.EmbeddingIndexStore; import org.apache.geaflow.ai.index.EntityAttributeIndexStore; import org.apache.geaflow.ai.index.IndexStore; import org.apache.geaflow.ai.operator.EmbeddingOperator; +import org.apache.geaflow.ai.operator.GraphSearchStore; +import org.apache.geaflow.ai.operator.GraphSearchStore.ScoredGraphEntity; import org.apache.geaflow.ai.operator.SearchOperator; import org.apache.geaflow.ai.operator.SessionOperator; import org.apache.geaflow.ai.search.VectorSearch; @@ -64,6 +67,55 @@ public List getIndexStores() { return indexStores; } + /** Executes bounded keyword retrieval without creating a session. */ + public List searchKeyword(String query, int topK, int maxCandidates, + long deadlineNanos) { + if (graphAccessors.isEmpty()) { + throw new IllegalStateException("No graph accessor available"); + } + for (IndexStore indexStore : indexStores) { + if (indexStore instanceof EntityAttributeIndexStore) { + EntityAttributeIndexStore keywordIndex = (EntityAttributeIndexStore) indexStore; + if (!keywordIndex.isInitialized()) { + throw new IllegalStateException("Keyword index is not initialized"); + } + GraphAccessor accessor = graphAccessors.get(0); + GraphSearchStore store = new GraphSearchStore(); + try { + java.util.Iterator vertices = accessor.scanVertex(); + int candidates = 0; + while (vertices.hasNext()) { + if (candidates >= maxCandidates || System.nanoTime() > deadlineNanos) { + if (System.nanoTime() > deadlineNanos) { + throw new org.apache.geaflow.ai.retrieval.api.model.RetrievalException( + org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.RETRIEVAL_TIMEOUT, + "retrieval deadline exceeded"); + } + break; + } + candidates++; + GraphVertex vertex = vertices.next(); + List vectors = + keywordIndex.getEntityIndex(vertex); + if (vectors != null && !vectors.isEmpty()) { + store.indexVertex(vertex, vectors); + } + } + if (System.nanoTime() > deadlineNanos) { + throw new org.apache.geaflow.ai.retrieval.api.model.RetrievalException( + org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.RETRIEVAL_TIMEOUT, + "retrieval deadline exceeded"); + } + store.finishWriting(); + return store.searchScored(query, accessor, topK, maxCandidates, deadlineNanos); + } finally { + store.close(); + } + } + } + throw new IllegalStateException("Keyword index is not available"); + } + public String createSession() { String sessionId = sessionManagement.createSession(); if (sessionId == null) { diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EntityAttributeIndexStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EntityAttributeIndexStore.java index aa9823876..d8a827cc0 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EntityAttributeIndexStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/index/EntityAttributeIndexStore.java @@ -31,7 +31,11 @@ public class EntityAttributeIndexStore implements IndexStore { - private VerbalizationFunction verbFunc; + private volatile VerbalizationFunction verbFunc; + + public boolean isInitialized() { + return verbFunc != null; + } public void initStore(VerbalizationFunction func) { if (func != null) { @@ -41,6 +45,9 @@ public void initStore(VerbalizationFunction func) { @Override public List getEntityIndex(GraphEntity entity) { + if (!isInitialized()) { + throw new IllegalStateException("Entity attribute index is not initialized"); + } if (entity instanceof GraphVertex) { String verbalization = verbFunc.verbalize(new SubGraph().addVertex((GraphVertex) entity)); List sentences = new ArrayList<>(); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java index 20beb48ad..43a4da513 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java @@ -19,8 +19,14 @@ package org.apache.geaflow.ai.operator; -import java.util.*; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; import java.util.stream.Collectors; +import org.apache.geaflow.ai.common.config.Constants; import org.apache.geaflow.ai.graph.GraphAccessor; import org.apache.geaflow.ai.graph.GraphEdge; import org.apache.geaflow.ai.graph.GraphEntity; @@ -40,6 +46,30 @@ public class GraphSearchStore { + public static final class ScoredGraphEntity { + private final GraphEntity entity; + private final float score; + private final int rank; + + public ScoredGraphEntity(GraphEntity entity, float score, int rank) { + this.entity = entity; + this.score = score; + this.rank = rank; + } + + public GraphEntity getEntity() { + return entity; + } + + public float getScore() { + return score; + } + + public int getRank() { + return rank; + } + } + private SearchStore store; private long entityNum = 0L; @@ -90,16 +120,36 @@ public boolean indexEdge(GraphEdge graphEdge, List indexVectors) { } public List search(String key1, GraphAccessor graphAccessor) { + List scored = searchScored(key1, graphAccessor, + Constants.GRAPH_SEARCH_STORE_DEFAULT_TOPN, Constants.GRAPH_SEARCH_STORE_DEFAULT_TOPN, + Long.MAX_VALUE); + List result = new ArrayList<>(scored.size()); + for (ScoredGraphEntity hit : scored) { + result.add(hit.getEntity()); + } + return result; + } + + public List searchScored(String key1, GraphAccessor graphAccessor, + int topK, int maxCandidates, + long deadlineNanos) { try { String query = SearchUtils.formatQuery(key1); - TopDocs docs = store.searchDoc(SearchConstants.CONTENT, query); + TopDocs docs = store.searchDoc(SearchConstants.CONTENT, query, + Math.max(topK, maxCandidates)); ScoreDoc[] scoreDocArray = docs.scoreDocs; Set vertexLabels = graphAccessor.getGraphSchema().getVertexSchemaList() .stream().map(VertexSchema::getLabel).collect(Collectors.toSet()); Set edgeLabels = graphAccessor.getGraphSchema().getEdgeSchemaList() .stream().map(EdgeSchema::getLabel).collect(Collectors.toSet()); - List result = new ArrayList<>(); + Map result = new LinkedHashMap<>(); + int candidateCount = 0; for (ScoreDoc scoreDoc : scoreDocArray) { + if (candidateCount >= maxCandidates || result.size() >= topK + || System.nanoTime() > deadlineNanos) { + break; + } + candidateCount++; int docId = scoreDoc.doc; Document document = store.getDoc(docId); String label = document.get(SearchConstants.LABEL); @@ -107,25 +157,53 @@ public List search(String key1, GraphAccessor graphAccessor) { String id = document.get(SearchConstants.ID); GraphVertex graphVertex = graphAccessor.getVertex(label, id); if (graphVertex != null) { - result.add(graphVertex); + putHighest(result, graphVertex, scoreDoc.score); } } else if (edgeLabels.contains(label)) { String src = document.get(SearchConstants.SRC); String dst = document.get(SearchConstants.DST); List graphEdge = graphAccessor.getEdge(label, src, dst); if (graphEdge != null) { - result.addAll(graphEdge); + for (GraphEdge edge : graphEdge) { + putHighest(result, edge, scoreDoc.score); + } } } } - return result; + List ordered = new ArrayList<>(result.values()); + ordered.sort((left, right) -> Float.compare(right.getScore(), left.getScore())); + List limited = new ArrayList<>(); + for (ScoredGraphEntity hit : ordered) { + if (limited.size() >= topK) { + break; + } + limited.add(new ScoredGraphEntity(hit.getEntity(), hit.getScore(), limited.size() + 1)); + } + return limited; } catch (IndexNotFoundException notFoundException) { return new ArrayList<>(); - } catch (Throwable e) { + } catch (Exception e) { throw new RuntimeException("Cannot read search store", e); } } + private static void putHighest(Map result, GraphEntity entity, + float score) { + String key; + if (entity instanceof GraphVertex) { + GraphVertex vertex = (GraphVertex) entity; + key = "V:" + vertex.getVertex().getLabel() + ":" + vertex.getVertex().getId(); + } else { + GraphEdge edge = (GraphEdge) entity; + key = "E:" + edge.getEdge().getLabel() + ":" + edge.getEdge().getSrcId() + + ":" + edge.getEdge().getDstId(); + } + ScoredGraphEntity previous = result.get(key); + if (previous == null || score > previous.getScore()) { + result.put(key, new ScoredGraphEntity(entity, score, 0)); + } + } + private void addItem() { entityNum++; } @@ -138,6 +216,14 @@ public void close() { } } + public void finishWriting() { + try { + store.finishWriting(); + } catch (Throwable e) { + throw new RuntimeException("Cannot finish search store writes", e); + } + } + public Directory getDirectory() { return store.getDirectory(); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java index 81183ae47..6749246d1 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java @@ -62,13 +62,20 @@ public void addDoc(Map kv) throws IOException { } public TopDocs searchDoc(String field, String content) throws ParseException, IOException { + return searchDoc(field, content, Constants.GRAPH_SEARCH_STORE_DEFAULT_TOPN); + } + + public TopDocs searchDoc(String field, String content, int topN) throws ParseException, IOException { + if (topN < 1) { + throw new IllegalArgumentException("topN must be positive"); + } if (!readStats) { reader = DirectoryReader.open(directory); searcher = new IndexSearcher(reader); readStats = true; } QueryParser parser = new QueryParser(field, analyzer); - return searcher.search(parser.parse(content), Constants.GRAPH_SEARCH_STORE_DEFAULT_TOPN); + return searcher.search(parser.parse(content), topN); } public Document getDoc(int docId) { @@ -93,14 +100,20 @@ public void initWriter() throws IOException { } public void close() throws IOException { - if (writeStats) { - writer.close(); - writeStats = false; - } + finishWriting(); if (readStats) { reader.close(); readStats = false; } + directory.close(); + analyzer.close(); + } + + public void finishWriting() throws IOException { + if (writeStats) { + writer.close(); + writeStats = false; + } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SessionOperator.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SessionOperator.java index 4800c9341..5ed203ccd 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SessionOperator.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SessionOperator.java @@ -81,8 +81,13 @@ public List apply(List subGraphList, VectorSearch search) { } //recall compute GraphSearchStore searchStore = initSearchStore(extendEntityIndexMap); - searchStore.close(); - List matchEntities = searchStore.search(query, graphAccessor); + searchStore.finishWriting(); + List matchEntities; + try { + matchEntities = searchStore.search(query, graphAccessor); + } finally { + searchStore.close(); + } Set matchEntitiesSet = new HashSet<>(matchEntities); //Apply to subgraph @@ -123,8 +128,12 @@ private List searchWithGlobalGraph(String query) { } //recall compute GraphSearchStore searchStore = initSearchStore(entityIndexMap); - searchStore.close(); - return searchStore.search(query, graphAccessor); + searchStore.finishWriting(); + try { + return searchStore.search(query, graphAccessor); + } finally { + searchStore.close(); + } } private GraphSearchStore initSearchStore(Map> entityIndexMap) { diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java index dea798bf6..2c73715ba 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java @@ -130,7 +130,8 @@ private static Object parse(JsonObject object, Class type) { optionalString(object, "text"), models(object, "chunks", TextChunk.class), models(object, "entities", EntityRef.class), models(object, "paths", GraphPathRef.class), models(object, "sources", SourceRef.class), scores(object), - optionalDouble(object, "fusedScore"), optionalInt(object, "rank")); + optionalDouble(object, "fusedScore"), optionalDouble(object, "finalScore"), + optionalInt(object, "rank")); } throw new JsonParseException("unsupported retrieval model: " + type.getName()); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/Bm25IndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/Bm25IndexBuilder.java new file mode 100644 index 000000000..4cfcb2f61 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/Bm25IndexBuilder.java @@ -0,0 +1,31 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index; + +import java.io.IOException; +import java.util.List; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; + +/** Builds a versioned lexical index from canonical text chunks. */ +public interface Bm25IndexBuilder { + + IndexArtifact build(IngestionContext context, List chunks) throws IOException; +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/IndexArtifact.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/IndexArtifact.java new file mode 100644 index 000000000..b4f04b5b5 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/IndexArtifact.java @@ -0,0 +1,29 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index; + +import java.io.Closeable; +import org.apache.geaflow.ai.retrieval.metadata.IndexBuildMetadata; + +/** Versioned index output. The pipeline closes temporary resources after publication or failure. */ +public interface IndexArtifact extends Closeable { + + IndexBuildMetadata getMetadata(); +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/VectorIndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/VectorIndexBuilder.java new file mode 100644 index 000000000..62e522e49 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/VectorIndexBuilder.java @@ -0,0 +1,31 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index; + +import java.io.IOException; +import java.util.List; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; + +/** Builds a replaceable vector index from canonical text chunks. */ +public interface VectorIndexBuilder { + + IndexArtifact build(IngestionContext context, List chunks) throws IOException; +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/package-info.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/package-info.java new file mode 100644 index 000000000..5c303da8c --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/package-info.java @@ -0,0 +1,26 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +/** + * Dataset-neutral index builder contracts. + * + *

Builders receive versioned canonical chunks and return closeable artifacts. A builder must + * not publish a ready version directly; publication is coordinated by the ingestion pipeline.

+ */ +package org.apache.geaflow.ai.retrieval.index; diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/DocumentNormalizer.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/DocumentNormalizer.java new file mode 100644 index 000000000..f271e79d7 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/DocumentNormalizer.java @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.io.IOException; +import java.util.List; +import org.apache.geaflow.ai.retrieval.model.document.SourceDocument; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; + +/** Normalizes canonical documents and applies the manifest's chunking policy. */ +public interface DocumentNormalizer { + + List normalize(List documents, IngestionContext context) + throws IOException; +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/EntityExtractor.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/EntityExtractor.java new file mode 100644 index 000000000..5517be5f1 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/EntityExtractor.java @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.io.IOException; +import java.util.List; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; + +/** Extracts canonical entities and relationships from normalized chunks. */ +public interface EntityExtractor { + + ExtractionResult extract(List chunks, IngestionContext context) throws IOException; +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ExtractionResult.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ExtractionResult.java new file mode 100644 index 000000000..df10544ed --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ExtractionResult.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; + +/** Dataset-neutral extracted graph records. */ +public final class ExtractionResult { + + private final List entities; + private final List edges; + + public ExtractionResult(List entities, List edges) { + this.entities = immutable(entities, "entities"); + this.edges = immutable(edges, "edges"); + } + + private static List immutable(List values, String name) { + Objects.requireNonNull(values, name); + if (values.stream().anyMatch(Objects::isNull)) { + throw new IllegalArgumentException(name + " must not contain null"); + } + return Collections.unmodifiableList(new java.util.ArrayList<>(values)); + } + + public List getEntities() { + return entities; + } + + public List getEdges() { + return edges; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/GraphArtifact.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/GraphArtifact.java new file mode 100644 index 000000000..0feb520b5 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/GraphArtifact.java @@ -0,0 +1,29 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.io.Closeable; +import org.apache.geaflow.ai.retrieval.metadata.GraphBuildMetadata; + +/** Temporary or durable graph build output owned by the pipeline until publication. */ +public interface GraphArtifact extends Closeable { + + GraphBuildMetadata getMetadata(); +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/GraphWriter.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/GraphWriter.java new file mode 100644 index 000000000..2cf396115 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/GraphWriter.java @@ -0,0 +1,32 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.io.IOException; +import java.util.List; +import org.apache.geaflow.ai.retrieval.model.document.SourceDocument; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; + +/** Writes canonical documents, chunks, and extracted graph records to a versioned graph. */ +public interface GraphWriter { + + GraphArtifact write(IngestionContext context, List documents, + List chunks, ExtractionResult extraction) throws IOException; +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionContext.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionContext.java new file mode 100644 index 000000000..b653da519 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionContext.java @@ -0,0 +1,68 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; +import org.apache.geaflow.ai.retrieval.metadata.QualityCounters; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; + +/** Immutable build inputs plus the counters collected by the pipeline. */ +public final class IngestionContext { + + private final DatasetManifest manifest; + private final GraphVersion graphVersion; + private final String importerVersion; + private volatile QualityCounters qualityCounters = new QualityCounters(0, 0, 0, 0, 0, 0); + + public IngestionContext(DatasetManifest manifest, GraphVersion graphVersion, + String importerVersion) { + this.manifest = Objects.requireNonNull(manifest, "manifest"); + this.graphVersion = Objects.requireNonNull(graphVersion, "graphVersion"); + this.importerVersion = Objects.requireNonNull(importerVersion, "importerVersion"); + if (importerVersion.trim().isEmpty()) { + throw new IllegalArgumentException("importerVersion must not be blank"); + } + } + + public DatasetManifest getManifest() { + return manifest; + } + + public GraphVersion getGraphVersion() { + return graphVersion; + } + + public String getImporterVersion() { + return importerVersion; + } + + public long getRandomSeed() { + return manifest.getRandomSeed(); + } + + public QualityCounters getQualityCounters() { + return qualityCounters; + } + + void setQualityCounters(QualityCounters qualityCounters) { + this.qualityCounters = Objects.requireNonNull(qualityCounters, "qualityCounters"); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java new file mode 100644 index 000000000..d3ca791ec --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java @@ -0,0 +1,133 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.index.Bm25IndexBuilder; +import org.apache.geaflow.ai.retrieval.index.IndexArtifact; +import org.apache.geaflow.ai.retrieval.index.VectorIndexBuilder; +import org.apache.geaflow.ai.retrieval.metadata.ImportMetadata; +import org.apache.geaflow.ai.retrieval.metadata.QualityCounters; +import org.apache.geaflow.ai.retrieval.model.document.SourceDocument; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; + +/** Deterministic orchestration boundary for source, graph, and index builders. */ +public final class IngestionPipeline { + + private final SourceLoader sourceLoader; + private final DocumentNormalizer normalizer; + private final EntityExtractor entityExtractor; + private final GraphWriter graphWriter; + private final Bm25IndexBuilder bm25Builder; + private final VectorIndexBuilder vectorBuilder; + private final MetadataPublisher metadataPublisher; + + public IngestionPipeline(SourceLoader sourceLoader, DocumentNormalizer normalizer, + EntityExtractor entityExtractor, GraphWriter graphWriter, + Bm25IndexBuilder bm25Builder, VectorIndexBuilder vectorBuilder, + MetadataPublisher metadataPublisher) { + this.sourceLoader = Objects.requireNonNull(sourceLoader, "sourceLoader"); + this.normalizer = Objects.requireNonNull(normalizer, "normalizer"); + this.entityExtractor = Objects.requireNonNull(entityExtractor, "entityExtractor"); + this.graphWriter = Objects.requireNonNull(graphWriter, "graphWriter"); + this.bm25Builder = Objects.requireNonNull(bm25Builder, "bm25Builder"); + this.vectorBuilder = Objects.requireNonNull(vectorBuilder, "vectorBuilder"); + this.metadataPublisher = Objects.requireNonNull(metadataPublisher, "metadataPublisher"); + } + + /** Runs one attempt and publishes only after every artifact has been built successfully. */ + public ImportMetadata run(IngestionContext context) throws IOException { + Objects.requireNonNull(context, "context"); + List graphArtifacts = new ArrayList<>(); + List indexArtifacts = new ArrayList<>(); + try { + metadataPublisher.begin(context); + List documents = requireList(sourceLoader.load(context), "documents"); + List chunks = requireList(normalizer.normalize(documents, context), "chunks"); + ExtractionResult extraction = Objects.requireNonNull( + entityExtractor.extract(chunks, context), "extraction"); + GraphArtifact graph = Objects.requireNonNull( + graphWriter.write(context, documents, chunks, extraction), "graph artifact"); + graphArtifacts.add(graph); + IndexArtifact bm25 = Objects.requireNonNull(bm25Builder.build(context, chunks), "bm25 artifact"); + indexArtifacts.add(bm25); + IndexArtifact vector = Objects.requireNonNull(vectorBuilder.build(context, chunks), "vector artifact"); + indexArtifacts.add(vector); + context.setQualityCounters(new QualityCounters(documents.size(), chunks.size(), + extraction.getEntities().size(), extraction.getEdges().size(), 0, 0)); + return metadataPublisher.publish(context, graph, + Collections.unmodifiableList(indexArtifacts)); + } catch (Exception failure) { + closeAll(indexArtifacts, failure); + indexArtifacts.clear(); + closeAll(graphArtifacts, failure); + graphArtifacts.clear(); + try { + metadataPublisher.fail(context, failure); + } catch (IOException publicationFailure) { + failure.addSuppressed(publicationFailure); + } + if (failure instanceof IOException) { + throw (IOException) failure; + } + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } + throw new IOException("ingestion failed", failure); + } finally { + closeAll(indexArtifacts, null); + closeAll(graphArtifacts, null); + } + } + + private static List requireList(List values, String name) { + Objects.requireNonNull(values, name); + if (values.stream().anyMatch(Objects::isNull)) { + throw new IllegalArgumentException(name + " must not contain null"); + } + return values; + } + + private static void closeAll(List resources, Exception failure) + throws IOException { + IOException closeFailure = null; + for (AutoCloseable resource : resources) { + try { + resource.close(); + } catch (Exception closeError) { + if (closeFailure == null) { + closeFailure = new IOException("failed to close ingestion artifact", closeError); + } else { + closeFailure.addSuppressed(closeError); + } + } + } + if (closeFailure != null && failure == null) { + throw closeFailure; + } + if (closeFailure != null && failure != null) { + failure.addSuppressed(closeFailure); + } + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataPublisher.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataPublisher.java new file mode 100644 index 000000000..0cd58af14 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataPublisher.java @@ -0,0 +1,36 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.io.IOException; +import java.util.List; +import org.apache.geaflow.ai.retrieval.index.IndexArtifact; +import org.apache.geaflow.ai.retrieval.metadata.ImportMetadata; + +/** Owns import state changes and atomic publication of completed artifacts. */ +public interface MetadataPublisher { + + void begin(IngestionContext context) throws IOException; + + ImportMetadata publish(IngestionContext context, GraphArtifact graph, + List indexes) throws IOException; + + void fail(IngestionContext context, Exception failure) throws IOException; +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/SourceLoader.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/SourceLoader.java new file mode 100644 index 000000000..a20b680a9 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/SourceLoader.java @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.io.IOException; +import java.util.List; +import org.apache.geaflow.ai.retrieval.model.document.SourceDocument; + +/** Loads canonical source documents without exposing dataset-specific record classes. */ +public interface SourceLoader { + + List load(IngestionContext context) throws IOException; +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/package-info.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/package-info.java new file mode 100644 index 000000000..e0e028f20 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/package-info.java @@ -0,0 +1,28 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +/** + * Storage-neutral ingestion extension points. + * + *

Loaders own the source records they return. The pipeline owns graph and index artifacts + * created during an attempt and closes them after publication or failure. Implementations must + * write versioned temporary output and expose it only through their artifact metadata; the + * metadata publisher is the sole component allowed to publish a ready version.

+ */ +package org.apache.geaflow.ai.retrieval.ingest; diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/Evidence.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/Evidence.java index ef2dd250c..a2e7a94fb 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/Evidence.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/Evidence.java @@ -60,6 +60,8 @@ public final class Evidence { private final Map stageScores; @SerializedName("fusedScore") private final Double fusedScore; + @SerializedName("finalScore") + private final Double finalScore; @SerializedName("rank") private final Integer rank; @@ -73,12 +75,21 @@ private Evidence() { this.sources = Collections.emptyList(); this.stageScores = Collections.emptyMap(); this.fusedScore = null; + this.finalScore = null; this.rank = null; } public Evidence(String evidenceId, EvidenceKind kind, String text, List chunks, List entities, List paths, List sources, Map stageScores, Double fusedScore, Integer rank) { + this(evidenceId, kind, text, chunks, entities, paths, sources, stageScores, fusedScore, + null, rank); + } + + public Evidence(String evidenceId, EvidenceKind kind, String text, List chunks, + List entities, List paths, List sources, + Map stageScores, Double fusedScore, Double finalScore, + Integer rank) { this.evidenceId = ModelValidation.optionalNonBlank(evidenceId, "evidenceId"); this.kind = Objects.requireNonNull(kind, "kind"); this.text = ModelValidation.optional(text); @@ -96,6 +107,7 @@ public Evidence(String evidenceId, EvidenceKind kind, String text, List getEntities() { return Collections.unmodifiableList(entities == null ? Collections.emptyList() : entities); } + /** Returns the primary entity identifier for entity evidence, when present. */ + public String getEntityId() { + return entities == null || entities.isEmpty() ? null : entities.get(0).getEntityId(); + } + public List getPaths() { return Collections.unmodifiableList(paths == null ? Collections.emptyList() : paths); } @@ -136,6 +153,10 @@ public Double getFusedScore() { return fusedScore; } + public Double getFinalScore() { + return finalScore; + } + public Integer getRank() { return rank; } @@ -201,12 +222,13 @@ public boolean equals(Object object) { && Objects.equals(sources, that.sources) && Objects.equals(stageScores, that.stageScores) && Objects.equals(fusedScore, that.fusedScore) + && Objects.equals(finalScore, that.finalScore) && Objects.equals(rank, that.rank); } @Override public int hashCode() { return Objects.hash(evidenceId, kind, text, chunks, entities, paths, sources, - stageScores, fusedScore, rank); + stageScores, fusedScore, finalScore, rank); } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/EvidenceMapper.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/EvidenceMapper.java new file mode 100644 index 000000000..178009238 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/EvidenceMapper.java @@ -0,0 +1,76 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.service; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.apache.geaflow.ai.graph.GraphEdge; +import org.apache.geaflow.ai.graph.GraphEntity; +import org.apache.geaflow.ai.graph.GraphVertex; +import org.apache.geaflow.ai.operator.GraphSearchStore.ScoredGraphEntity; +import org.apache.geaflow.ai.retrieval.model.evidence.ChannelScore; +import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; +import org.apache.geaflow.ai.retrieval.model.evidence.EvidenceKind; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; + +/** Maps graph search hits to stable API evidence records. */ +public final class EvidenceMapper { + + public List map(List hits) { + Map unique = new LinkedHashMap<>(); + for (ScoredGraphEntity hit : hits == null ? Collections.emptyList() : hits) { + GraphEntity entity = hit.getEntity(); + String identity = identity(entity); + Map scores = new LinkedHashMap<>(); + scores.put("keyword", new ChannelScore("keyword", hit.getScore(), + null, hit.getRank())); + Evidence evidence = new Evidence(identity, EvidenceKind.ENTITY, entity.toString(), + Collections.emptyList(), Collections.singletonList(toRef(entity)), + Collections.emptyList(), Collections.emptyList(), scores, + null, (double) hit.getScore(), hit.getRank()); + unique.putIfAbsent(identity, evidence); + } + return new ArrayList<>(unique.values()); + } + + private static String identity(GraphEntity entity) { + if (entity instanceof GraphVertex) { + GraphVertex vertex = (GraphVertex) entity; + return "VERTEX:" + vertex.getVertex().getLabel() + ":" + vertex.getVertex().getId(); + } + GraphEdge edge = (GraphEdge) entity; + return "EDGE:" + edge.getEdge().getLabel() + ":" + edge.getEdge().getSrcId() + + ":" + edge.getEdge().getDstId(); + } + + private static EntityRef toRef(GraphEntity entity) { + if (entity instanceof GraphVertex) { + GraphVertex vertex = (GraphVertex) entity; + String id = vertex.getVertex().getId(); + return new EntityRef(id, id, vertex.getVertex().getLabel()); + } + GraphEdge edge = (GraphEdge) entity; + String id = edge.getEdge().getSrcId() + "->" + edge.getEdge().getDstId(); + return new EntityRef(id, id, edge.getEdge().getLabel()); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalMetrics.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalMetrics.java new file mode 100644 index 000000000..279236108 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalMetrics.java @@ -0,0 +1,99 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.service; + +import java.util.concurrent.atomic.AtomicLong; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.noear.solon.annotation.Component; + +/** Process-local counters for the Week 1 retrieval endpoint. */ +@Component +public final class RetrievalMetrics { + + private final AtomicLong success = new AtomicLong(); + private final AtomicLong failure = new AtomicLong(); + private final AtomicLong timeout = new AtomicLong(); + private final AtomicLong total = new AtomicLong(); + private final AtomicLong totalElapsedMs = new AtomicLong(); + + public void recordSuccess() { + recordSuccess(0L); + } + + public void recordSuccess(long elapsedMs) { + total.incrementAndGet(); + success.incrementAndGet(); + totalElapsedMs.addAndGet(Math.max(0L, elapsedMs)); + } + + public void recordFailure(RetrievalErrorCode code) { + recordFailure(code, 0L); + } + + public void recordFailure(RetrievalErrorCode code, long elapsedMs) { + total.incrementAndGet(); + failure.incrementAndGet(); + totalElapsedMs.addAndGet(Math.max(0L, elapsedMs)); + if (code == RetrievalErrorCode.RETRIEVAL_TIMEOUT) { + timeout.incrementAndGet(); + } + } + + public Snapshot snapshot() { + return new Snapshot(total.get(), success.get(), failure.get(), timeout.get(), + totalElapsedMs.get()); + } + + public static final class Snapshot { + private final long total; + private final long success; + private final long failure; + private final long timeout; + private final long totalElapsedMs; + + private Snapshot(long total, long success, long failure, long timeout, long totalElapsedMs) { + this.total = total; + this.success = success; + this.failure = failure; + this.timeout = timeout; + this.totalElapsedMs = totalElapsedMs; + } + + public long getTotal() { + return total; + } + + public long getSuccess() { + return success; + } + + public long getFailure() { + return failure; + } + + public long getTimeout() { + return timeout; + } + + public long getTotalElapsedMs() { + return totalElapsedMs; + } + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalService.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalService.java new file mode 100644 index 000000000..5bfd0e624 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalService.java @@ -0,0 +1,155 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.service; + +import java.util.Collections; +import java.util.List; +import java.util.UUID; +import org.apache.geaflow.ai.GraphMemoryServer; +import org.apache.geaflow.ai.operator.GraphSearchStore.ScoredGraphEntity; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalCommand; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalRequest; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalTrace; +import org.apache.geaflow.ai.retrieval.api.model.TraceStage; +import org.apache.geaflow.ai.retrieval.config.RetrievalProperties; +import org.apache.geaflow.ai.service.ServerMemoryCache; +import org.noear.solon.annotation.Component; +import org.noear.solon.annotation.Init; +import org.noear.solon.annotation.Inject; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Executes the bounded keyword retrieval flow without creating a session. */ +@Component +public class RetrievalService { + + private static final Logger LOGGER = LoggerFactory.getLogger(RetrievalService.class); + + @Inject + private ServerMemoryCache cache; + @Inject + private RetrievalProperties properties; + @Inject + private RetrievalMetrics metrics; + private RetrievalRequestValidator validator; + private final EvidenceMapper evidenceMapper = new EvidenceMapper(); + + public RetrievalService() { + this(new ServerMemoryCache(), new RetrievalProperties(), new RetrievalMetrics()); + } + + public RetrievalService(ServerMemoryCache cache, RetrievalProperties properties) { + this(cache, properties, new RetrievalMetrics()); + } + + public RetrievalService(ServerMemoryCache cache, RetrievalProperties properties, + RetrievalMetrics metrics) { + this.cache = cache; + this.properties = properties; + this.metrics = metrics; + this.properties.validateConfiguration(); + this.validator = new RetrievalRequestValidator(properties); + } + + @Init + public void initialize() { + properties.validateConfiguration(); + validator = new RetrievalRequestValidator(properties); + } + + public RetrievalResponse retrieve(RetrievalRequest request) { + return retrieve(request, UUID.randomUUID().toString()); + } + + public RetrievalResponse retrieve(RetrievalRequest request, String requestId) { + long startedAt = System.nanoTime(); + try { + RetrievalCommand command = validator.validate(request); + ServerMemoryCache.ReadyStatus readiness = cache.keywordReadiness(command.getGraphName()); + if ("GRAPH_NOT_LOADED".equals(readiness.getReason())) { + throw new RetrievalException(RetrievalErrorCode.GRAPH_NOT_FOUND, + "graph not found: " + command.getGraphName()); + } + if (!readiness.isReady()) { + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, + "keyword retrieval is not ready: " + readiness.getReason()); + } + GraphMemoryServer server = cache.getServerByName(command.getGraphName()); + if (server == null) { + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, + "keyword retrieval is not ready: SERVER_NOT_LOADED"); + } + RetrievalBudget budget = command.getBudget(); + long deadline = System.nanoTime() + budget.getTimeoutMs() * 1_000_000L; + final List hits = server.searchKeyword(command.getQuery(), + budget.getTopK(), budget.getMaxCandidates(), deadline); + if (System.nanoTime() > deadline) { + throw new RetrievalException(RetrievalErrorCode.RETRIEVAL_TIMEOUT, + "retrieval deadline exceeded"); + } + RetrievalResponse response = new RetrievalResponse(); + response.setRequestId(requestId); + response.setGraphName(command.getGraphName()); + response.setGraphVersion(cache.getGraphVersion(command.getGraphName())); + response.setEvidence(evidenceMapper.map(hits)); + response.setPaths(Collections.emptyList()); + response.setSources(Collections.emptyList()); + response.setDegradedChannels(Collections.emptyList()); + response.setEffectiveBudget(budget); + RetrievalTrace trace = new RetrievalTrace(); + trace.setTraceVersion("v1"); + trace.setOriginalQuery(command.getQuery()); + trace.setSelectedMode(command.getMode()); + trace.setExecutionMode(command.getExecutionMode()); + trace.setStages(Collections.singletonList(new TraceStage("keyword", "SUCCESS", null))); + trace.setStopReason(hits.isEmpty() ? "NO_MATCH" : "TOP_K"); + response.setTrace(trace); + long elapsedMs = elapsedMs(startedAt); + metrics.recordSuccess(elapsedMs); + LOGGER.info("retrieval completed graph={} mode={} topK={} maxCandidates={} " + + "timeoutMs={} candidateBudget={} results={} status=SUCCESS elapsedMs={}", + command.getGraphName(), command.getMode(), budget.getTopK(), + budget.getMaxCandidates(), budget.getTimeoutMs(), budget.getMaxCandidates(), + hits.size(), elapsedMs); + return response; + } catch (RetrievalException exception) { + long elapsedMs = elapsedMs(startedAt); + metrics.recordFailure(exception.getCode(), elapsedMs); + LOGGER.info("retrieval completed status={} code={} elapsedMs={}", + "FAILURE", exception.getCode(), elapsedMs); + throw exception; + } catch (RuntimeException exception) { + long elapsedMs = elapsedMs(startedAt); + metrics.recordFailure(RetrievalErrorCode.INTERNAL_ERROR, elapsedMs); + LOGGER.warn("retrieval failed status=FAILURE code={} elapsedMs={}", + RetrievalErrorCode.INTERNAL_ERROR, elapsedMs, exception); + throw new RetrievalException(RetrievalErrorCode.INTERNAL_ERROR, + "internal retrieval error"); + } + } + + private static long elapsedMs(long startedAt) { + return (System.nanoTime() - startedAt) / 1_000_000L; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/service/ServerMemoryCache.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/service/ServerMemoryCache.java index fd412bde3..bedd7a676 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/service/ServerMemoryCache.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/service/ServerMemoryCache.java @@ -19,30 +19,36 @@ package org.apache.geaflow.ai.service; -import java.util.HashMap; -import java.util.LinkedHashMap; import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicLong; import org.apache.geaflow.ai.GraphMemoryServer; import org.apache.geaflow.ai.consolidate.ConsolidateServer; import org.apache.geaflow.ai.graph.Graph; +import org.noear.solon.annotation.Component; +@Component public class ServerMemoryCache { - private final Map name2Graph = new LinkedHashMap<>(); - private final Map name2Server = new LinkedHashMap<>(); - private final Map session2GraphName = new HashMap<>(); + private final Map name2Graph = new ConcurrentHashMap<>(); + private final Map name2Server = new ConcurrentHashMap<>(); + private final Map session2GraphName = new ConcurrentHashMap<>(); + private final Map graphVersions = new ConcurrentHashMap<>(); private final ConsolidateServer consolidateServer = new ConsolidateServer(); public void putGraph(Graph g) { - name2Graph.put(g.getGraphSchema().getName(), g); + String name = g.getGraphSchema().getName(); + name2Graph.put(name, g); + graphVersions.putIfAbsent(name, new AtomicLong(1L)); } public void putServer(GraphMemoryServer server) { if (server.getGraphAccessors().isEmpty()) { throw new RuntimeException("Cannot register server without graph accessor"); } - name2Server.put(server.getGraphAccessors().get(0) - .getGraphSchema().getName(), server); + String name = server.getGraphAccessors().get(0).getGraphSchema().getName(); + name2Server.put(name, server); + graphVersions.putIfAbsent(name, new AtomicLong(1L)); } public void putSession(GraphMemoryServer server, String sessionId) { @@ -68,4 +74,86 @@ public String getGraphNameBySession(String sessionId) { public ConsolidateServer getConsolidateServer() { return consolidateServer; } + + public String getGraphVersion(String graphName) { + AtomicLong version = graphVersions.get(graphName); + return version == null ? null : "v" + version.get(); + } + + public String markGraphUpdated(String graphName) { + AtomicLong version = graphVersions.get(graphName); + if (version == null) { + throw new IllegalArgumentException("Unknown graph: " + graphName); + } + return "v" + version.incrementAndGet(); + } + + public ReadinessStatus keywordReadiness(String graphName) { + if (name2Graph.get(graphName) == null) { + return new ReadinessStatus(false, graphName, getGraphVersion(graphName), "GRAPH_NOT_LOADED"); + } + if (name2Server.get(graphName) == null) { + return new ReadinessStatus(false, graphName, getGraphVersion(graphName), "SERVER_NOT_LOADED"); + } + GraphMemoryServer server = name2Server.get(graphName); + if (server.getGraphAccessors().isEmpty()) { + return new ReadinessStatus(false, graphName, getGraphVersion(graphName), "GRAPH_ACCESSOR_NOT_READY"); + } + for (org.apache.geaflow.ai.index.IndexStore store : server.getIndexStores()) { + if (store instanceof org.apache.geaflow.ai.index.EntityAttributeIndexStore + && ((org.apache.geaflow.ai.index.EntityAttributeIndexStore) store).isInitialized()) { + return new ReadinessStatus(true, graphName, getGraphVersion(graphName), "READY"); + } + } + return new ReadinessStatus(false, graphName, getGraphVersion(graphName), "KEYWORD_INDEX_NOT_READY"); + } + + public static class ReadyStatus { + private final boolean ready; + private final String graphName; + private final String graphVersion; + private final String reason; + + public ReadyStatus(boolean ready, String reason) { + this(ready, null, null, reason); + } + + public ReadyStatus(boolean ready, String graphName, String graphVersion, String reason) { + this.ready = ready; + this.graphName = graphName; + this.graphVersion = graphVersion; + this.reason = reason; + } + + public boolean isReady() { + return ready; + } + + public String getReason() { + return reason; + } + + public String getGraphName() { + return graphName; + } + + public String getGraphVersion() { + return graphVersion; + } + } + + public static final class ReadinessStatus extends ReadyStatus { + + public static ReadinessStatus ready(String graphName, String graphVersion) { + return new ReadinessStatus(true, graphName, graphVersion, "READY"); + } + + public static ReadinessStatus notReady(String reason) { + return new ReadinessStatus(false, null, null, reason); + } + + public ReadinessStatus(boolean ready, String graphName, String graphVersion, String reason) { + super(ready, graphName, graphVersion, reason); + } + } } diff --git a/geaflow-ai/src/main/resources/application.yml b/geaflow-ai/src/main/resources/application.yml index b21ef36b5..b808dcc56 100644 --- a/geaflow-ai/src/main/resources/application.yml +++ b/geaflow-ai/src/main/resources/application.yml @@ -21,7 +21,7 @@ solon: group: geaflow server: - port: 8080 + port: ${GEAFLOW_SERVER_PORT:8080} host: 0.0.0.0 logging: diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/GeaFlowMemoryServerTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/GeaFlowMemoryServerTest.java new file mode 100644 index 000000000..3e4a1a67c --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/GeaFlowMemoryServerTest.java @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; + +import org.junit.jupiter.api.Test; + +class GeaFlowMemoryServerTest { + + @Test + void preservesValidRequestIdAndRegeneratesInvalidValues() { + assertEquals("trace-1", GeaFlowMemoryServer.resolveRequestId(" trace-1 ")); + assertNotEquals("", GeaFlowMemoryServer.resolveRequestId(" ")); + assertNotEquals("", GeaFlowMemoryServer.resolveRequestId(null)); + assertNotEquals("", GeaFlowMemoryServer.resolveRequestId(repeat('x', 129))); + assertNotEquals("trace/1", GeaFlowMemoryServer.resolveRequestId("trace/1")); + } + + private static String repeat(char value, int count) { + StringBuilder result = new StringBuilder(count); + for (int index = 0; index < count; index++) { + result.append(value); + } + return result.toString(); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/GraphMemoryKeywordSearchTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/GraphMemoryKeywordSearchTest.java new file mode 100644 index 000000000..401840566 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/GraphMemoryKeywordSearchTest.java @@ -0,0 +1,86 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.apache.geaflow.ai.graph.LocalMemoryGraphAccessor; +import org.apache.geaflow.ai.graph.io.EntityGroup; +import org.apache.geaflow.ai.graph.io.GraphSchema; +import org.apache.geaflow.ai.graph.io.MemoryGraph; +import org.apache.geaflow.ai.graph.io.Vertex; +import org.apache.geaflow.ai.graph.io.VertexGroup; +import org.apache.geaflow.ai.graph.io.VertexSchema; +import org.apache.geaflow.ai.index.EntityAttributeIndexStore; +import org.apache.geaflow.ai.operator.GraphSearchStore.ScoredGraphEntity; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.verbalization.SubgraphSemanticPromptFunction; +import org.junit.jupiter.api.Test; + +class GraphMemoryKeywordSearchTest { + + @Test + void maxCandidatesStopsScanningBeforeTailMatch() { + GraphMemoryServer server = createServer(); + + List hits = server.searchKeyword("stars", 10, 1, Long.MAX_VALUE); + + assertEquals(0, hits.size()); + } + + @Test + void expiredDeadlineIsReportedAsTimeout() { + GraphMemoryServer server = createServer(); + + RetrievalException exception = assertThrows(RetrievalException.class, + () -> server.searchKeyword("confucius", 1, 10, System.nanoTime() - 1)); + + assertEquals(RetrievalErrorCode.RETRIEVAL_TIMEOUT, exception.getCode()); + } + + private static GraphMemoryServer createServer() { + VertexSchema vertexSchema = new VertexSchema("chunk", "id", + Collections.singletonList("text")); + GraphSchema schema = new GraphSchema(); + schema.setName("keyword-search-test"); + schema.addVertex(vertexSchema); + List vertices = Arrays.asList( + new Vertex("chunk", "confucius-1", Collections.singletonList("Confucius taught ethics.")), + new Vertex("chunk", "astronomy-1", Collections.singletonList("Astronomy studies stars."))); + Map groups = new LinkedHashMap<>(); + groups.put("chunk", new VertexGroup(vertexSchema, new ArrayList<>(vertices))); + MemoryGraph graph = new MemoryGraph(schema, groups); + LocalMemoryGraphAccessor accessor = new LocalMemoryGraphAccessor(graph); + EntityAttributeIndexStore indexStore = new EntityAttributeIndexStore(); + indexStore.initStore(new SubgraphSemanticPromptFunction(accessor)); + GraphMemoryServer server = new GraphMemoryServer(); + server.addGraphAccessor(accessor); + server.addIndexStore(indexStore); + return server; + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java index 22c9b1a69..60757e216 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java @@ -32,6 +32,8 @@ import java.util.List; import java.util.Map; import org.apache.geaflow.ai.graph.io.Vertex; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse; +import org.apache.geaflow.ai.retrieval.codec.RetrievalApiJson; import org.apache.geaflow.ai.retrieval.support.HttpTestClient; import org.apache.geaflow.ai.retrieval.support.HttpTestResponse; import org.apache.geaflow.ai.retrieval.support.RetrievalTestFixture; @@ -46,7 +48,9 @@ void legacyWorkflow() { RetrievalTestFixture fixture = new RetrievalTestFixture(); Gson gson = new Gson(); Map graph = Collections.singletonMap("graphName", RetrievalTestFixture.GRAPH_NAME); - try (HttpTestClient client = new HttpTestClient("http://localhost:8080")) { + String port = System.getProperty("server.port", + System.getenv().getOrDefault("GEAFLOW_SERVER_PORT", "8080")); + try (HttpTestClient client = new HttpTestClient("http://localhost:" + port)) { HttpTestResponse health = client.get("/health"); success(health); assertTrue(health.getContentType().startsWith("application/json")); @@ -80,6 +84,42 @@ void legacyWorkflow() { success(result); assertTrue(result.getBody().contains("confucius-1"), result.getBody()); assertTrue(result.getBody().contains("Confucius taught ethics."), result.getBody()); + + HttpTestResponse ready = client.get("/ready"); + assertEquals(503, ready.getStatus(), ready.getBody()); + assertTrue(ready.getBody().contains("GRAPH_NOT_LOADED"), ready.getBody()); + + String request = "{\"graphName\":\"" + RetrievalTestFixture.GRAPH_NAME + + "\",\"query\":\"Confucius\",\"budget\":{\"topK\":1," + + "\"timeoutMs\":3000,\"maxCandidates\":10,\"tokenBudget\":4096}}"; + HttpTestResponse retrieval = client.postWithHeaders("/api/v1/retrievals", request, + Collections.singletonMap("X-Request-Id", "request-1")); + assertEquals(200, retrieval.getStatus(), retrieval.getBody()); + assertEquals("request-1", retrieval.getHeaders().get("X-Request-Id")); + RetrievalResponse response = RetrievalApiJson.parseResponse(retrieval.getBody()); + assertEquals("request-1", response.getRequestId()); + assertEquals("v4", response.getGraphVersion()); + assertEquals(1, response.getEvidence().size()); + assertTrue(response.getEvidence().get(0).getFinalScore() > 0); + assertTrue(response.getEvidence().get(0).getStageScores().containsKey("keyword")); + + HttpTestResponse noMatch = client.post("/api/v1/retrievals", + request.replace("Confucius", "not-present")); + assertEquals(200, noMatch.getStatus(), noMatch.getBody()); + assertEquals(0, RetrievalApiJson.parseResponse(noMatch.getBody()).getEvidence().size()); + + HttpTestResponse missing = client.post("/api/v1/retrievals", + request.replace(RetrievalTestFixture.GRAPH_NAME, "missing-graph")); + assertEquals(404, missing.getStatus(), missing.getBody()); + assertTrue(missing.getBody().contains("GRAPH_NOT_FOUND"), missing.getBody()); + + HttpTestResponse malformed = client.post("/api/v1/retrievals", "{"); + assertEquals(400, malformed.getStatus(), malformed.getBody()); + assertTrue(malformed.getBody().contains("Malformed JSON request"), malformed.getBody()); + + HttpTestResponse metrics = client.get("/metrics/retrieval"); + assertEquals(200, metrics.getStatus(), metrics.getBody()); + assertTrue(metrics.getBody().contains("\"success\":"), metrics.getBody()); } } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/SearchStoreTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/SearchStoreTest.java new file mode 100644 index 000000000..211d7aac6 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/SearchStoreTest.java @@ -0,0 +1,46 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.operator; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.Collections; +import org.apache.lucene.search.TopDocs; +import org.junit.jupiter.api.Test; + +class SearchStoreTest { + + @Test + void searchDocHonorsRequestedTopN() throws Exception { + SearchStore store = new SearchStore(); + try { + store.addDoc(Collections.singletonMap(SearchConstants.CONTENT, "confucius teacher")); + store.addDoc(Collections.singletonMap(SearchConstants.CONTENT, "confucius philosopher")); + store.addDoc(Collections.singletonMap(SearchConstants.CONTENT, "confucius scholar")); + store.finishWriting(); + + TopDocs docs = store.searchDoc(SearchConstants.CONTENT, "confucius", 1); + + assertEquals(1, docs.scoreDocs.length); + } finally { + store.close(); + } + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java new file mode 100644 index 000000000..4ef738e2c --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java @@ -0,0 +1,233 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.geaflow.ai.retrieval.index.Bm25IndexBuilder; +import org.apache.geaflow.ai.retrieval.index.IndexArtifact; +import org.apache.geaflow.ai.retrieval.index.VectorIndexBuilder; +import org.apache.geaflow.ai.retrieval.metadata.ChunkingConfiguration; +import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; +import org.apache.geaflow.ai.retrieval.metadata.GraphBuildMetadata; +import org.apache.geaflow.ai.retrieval.metadata.ImportMetadata; +import org.apache.geaflow.ai.retrieval.metadata.MetadataException; +import org.apache.geaflow.ai.retrieval.metadata.QualityCounters; +import org.apache.geaflow.ai.retrieval.metadata.IndexBuildMetadata; +import org.apache.geaflow.ai.retrieval.model.document.SourceDocument; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.junit.jupiter.api.Test; + +/** Contract tests for the storage-neutral ingestion extension points. */ +public class IngestionPipelineTest { + + private static final String SHA256 = "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef"; + + @Test + public void runsStagesInOrderAndReportsCounters() throws Exception { + List calls = new java.util.ArrayList<>(); + AtomicBoolean published = new AtomicBoolean(); + IngestionContext context = context(); + SourceDocument document = document(); + TextChunk chunk = chunk(); + TrackingGraphArtifact graph = new TrackingGraphArtifact(calls, "graph"); + TrackingIndexArtifact bm25 = new TrackingIndexArtifact(calls, "bm25"); + TrackingIndexArtifact vector = new TrackingIndexArtifact(calls, "vector"); + IngestionPipeline pipeline = new IngestionPipeline( + value -> { + calls.add("source"); + return Collections.singletonList(document); + }, + (documents, value) -> { + calls.add("normalize"); + return Collections.singletonList(chunk); + }, + (chunks, value) -> { + calls.add("extract"); + return new ExtractionResult(Collections.singletonList(entity()), + Collections.singletonList(edge())); + }, + (value, documents, chunks, extraction) -> { + calls.add("graph"); + return graph; + }, + (value, chunks) -> { + calls.add("bm25"); + return bm25; + }, + (value, chunks) -> { + calls.add("vector"); + return vector; + }, + new MetadataPublisher() { + @Override + public void begin(IngestionContext value) { + calls.add("begin"); + } + + @Override + public ImportMetadata publish(IngestionContext value, GraphArtifact valueGraph, + List indexes) { + calls.add("publish"); + assertEquals(2, indexes.size()); + published.set(true); + return null; + } + + @Override + public void fail(IngestionContext value, Exception failure) { + throw new AssertionError("unexpected failure", failure); + } + }); + + pipeline.run(context); + + assertEquals(Arrays.asList("begin", "source", "normalize", "extract", "graph", + "bm25", "vector", "publish", "close-bm25", "close-vector", "close-graph"), calls); + assertTrue(published.get()); + assertEquals(new QualityCounters(1, 1, 1, 1, 0, 0), context.getQualityCounters()); + } + + @Test + public void failureClosesArtifactsAndNeverPublishes() throws Exception { + AtomicBoolean published = new AtomicBoolean(); + AtomicBoolean failed = new AtomicBoolean(); + TrackingGraphArtifact graph = new TrackingGraphArtifact(new java.util.ArrayList<>(), "graph"); + TrackingIndexArtifact bm25 = new TrackingIndexArtifact(new java.util.ArrayList<>(), "bm25"); + IngestionPipeline pipeline = new IngestionPipeline( + value -> Collections.singletonList(document()), + (documents, value) -> Collections.singletonList(chunk()), + (chunks, value) -> new ExtractionResult(Collections.emptyList(), Collections.emptyList()), + (value, documents, chunks, extraction) -> graph, + (value, chunks) -> bm25, + (value, chunks) -> { + throw new IllegalStateException("vector build failed"); + }, + new MetadataPublisher() { + @Override + public void begin(IngestionContext value) { + } + + @Override + public ImportMetadata publish(IngestionContext value, GraphArtifact valueGraph, + List indexes) { + published.set(true); + return null; + } + + @Override + public void fail(IngestionContext value, Exception failure) { + failed.set(true); + assertTrue(failure.getMessage().contains("vector build failed")); + } + }); + + try { + pipeline.run(context()); + } catch (IllegalStateException expected) { + assertEquals("vector build failed", expected.getMessage()); + } + + assertFalse(published.get()); + assertTrue(failed.get()); + assertTrue(graph.closed); + assertTrue(bm25.closed); + } + + private static IngestionContext context() { + DatasetManifest manifest = new DatasetManifest("v1", "hotpotqa", "release-1", "train", + null, "cache/hotpotqa", SHA256, "normalizer-1", + new ChunkingConfiguration("chunker-1", 100, 10), "graph-schema-1", "fixture", "v1", 7); + return new IngestionContext(manifest, new GraphVersion("graph-1", "version-1"), "importer-1"); + } + + private static SourceDocument document() { + return new SourceDocument("doc-1", "hotpotqa", "release-1", "train", "uri:doc-1", SHA256); + } + + private static TextChunk chunk() { + return new TextChunk("chunk-1", "doc-1", 0, 0, 10, 3, "A fixed chunk."); + } + + private static EntityRef entity() { + return new EntityRef("entity-1", "Fixed entity", "person"); + } + + private static GraphEdgeRef edge() { + return new GraphEdgeRef("edge-1", "related", "entity-1", "entity-1"); + } + + private static final class TrackingGraphArtifact implements GraphArtifact { + private final List calls; + private final String name; + private boolean closed; + + private TrackingGraphArtifact(List calls, String name) { + this.calls = calls; + this.name = name; + } + + @Override + public GraphBuildMetadata getMetadata() { + return new GraphBuildMetadata(new GraphVersion("graph-1", "version-1"), + "memory:" + name, "fixture", "builder-1", true); + } + + @Override + public void close() { + closed = true; + calls.add("close-" + name); + } + } + + private static final class TrackingIndexArtifact implements IndexArtifact { + private final List calls; + private final String name; + private boolean closed; + + private TrackingIndexArtifact(List calls, String name) { + this.calls = calls; + this.name = name; + } + + @Override + public IndexBuildMetadata getMetadata() { + return new IndexBuildMetadata(new GraphVersion("graph-1", "version-1"), + new org.apache.geaflow.ai.retrieval.model.version.IndexVersion(name, "v1", "version-1"), + name, "builder-1", "memory:" + name, true); + } + + @Override + public void close() { + closed = true; + calls.add("close-" + name); + } + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/RetrievalServiceTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/RetrievalServiceTest.java new file mode 100644 index 000000000..c1de6d2a2 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/RetrievalServiceTest.java @@ -0,0 +1,143 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.service; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import org.apache.geaflow.ai.GraphMemoryServer; +import org.apache.geaflow.ai.graph.LocalMemoryGraphAccessor; +import org.apache.geaflow.ai.graph.io.EntityGroup; +import org.apache.geaflow.ai.graph.io.GraphSchema; +import org.apache.geaflow.ai.graph.io.MemoryGraph; +import org.apache.geaflow.ai.graph.io.Vertex; +import org.apache.geaflow.ai.graph.io.VertexGroup; +import org.apache.geaflow.ai.graph.io.VertexSchema; +import org.apache.geaflow.ai.index.EntityAttributeIndexStore; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalRequest; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse; +import org.apache.geaflow.ai.retrieval.config.RetrievalProperties; +import org.apache.geaflow.ai.service.ServerMemoryCache; +import org.apache.geaflow.ai.verbalization.SubgraphSemanticPromptFunction; +import org.junit.jupiter.api.Test; + +class RetrievalServiceTest { + + @Test + void retrievesScoredEvidenceWithoutCreatingSession() { + ServerMemoryCache cache = new ServerMemoryCache(); + MemoryGraph graph = createGraph(); + GraphMemoryServer server = createServer(graph); + cache.putGraph(graph); + cache.putServer(server); + + RetrievalRequest request = new RetrievalRequest(); + request.setGraphName("week1-service-graph"); + request.setQuery("Confucius"); + request.setBudget(new RetrievalBudget(1, 3000, 10, 4096)); + + RetrievalMetrics metrics = new RetrievalMetrics(); + RetrievalResponse response = new RetrievalService(cache, new RetrievalProperties(), metrics) + .retrieve(request, "request-1"); + + assertEquals("request-1", response.getRequestId()); + assertEquals("v1", response.getGraphVersion()); + assertEquals(1, response.getEvidence().size()); + assertTrue(response.getEvidence().get(0).getFinalScore() > 0); + assertTrue(response.getEvidence().get(0).getStageScores().get("keyword").getRawScore() > 0); + assertFalse(response.getEvidence().get(0).getStageScores().isEmpty()); + assertEquals(1, metrics.snapshot().getTotal()); + assertEquals(1, metrics.snapshot().getSuccess()); + } + + @Test + void preservesUnboundedKeywordScoresWithoutInventingFusion() { + MemoryGraph graph = createGraph(); + org.apache.geaflow.ai.graph.GraphVertex vertex = new LocalMemoryGraphAccessor(graph) + .getVertex("chunk", "chunk-1"); + org.apache.geaflow.ai.retrieval.model.evidence.Evidence evidence = new EvidenceMapper() + .map(Collections.singletonList( + new org.apache.geaflow.ai.operator.GraphSearchStore.ScoredGraphEntity(vertex, 8.5f, 1))) + .get(0); + assertEquals(8.5, evidence.getStageScores().get("keyword").getRawScore()); + assertEquals(8.5, evidence.getFinalScore()); + org.junit.jupiter.api.Assertions.assertNull(evidence.getFusedScore()); + } + + @Test + void distinguishesMissingGraph() { + RetrievalRequest request = new RetrievalRequest(); + request.setGraphName("missing"); + request.setQuery("query"); + + RetrievalException exception = org.junit.jupiter.api.Assertions.assertThrows( + RetrievalException.class, + () -> new RetrievalService(new ServerMemoryCache(), new RetrievalProperties()) + .retrieve(request, "request-2")); + + assertEquals(RetrievalErrorCode.GRAPH_NOT_FOUND, exception.getCode()); + } + + @Test + void distinguishesLoadedGraphWithUnavailableIndex() { + ServerMemoryCache cache = new ServerMemoryCache(); + cache.putGraph(createGraph()); + RetrievalRequest request = new RetrievalRequest(); + request.setGraphName("week1-service-graph"); + request.setQuery("query"); + + RetrievalException exception = org.junit.jupiter.api.Assertions.assertThrows( + RetrievalException.class, + () -> new RetrievalService(cache, new RetrievalProperties()) + .retrieve(request, "request-not-ready")); + + assertEquals(RetrievalErrorCode.INDEX_NOT_READY, exception.getCode()); + } + + private static MemoryGraph createGraph() { + VertexSchema schema = new VertexSchema("chunk", "id", Collections.singletonList("text")); + GraphSchema graphSchema = new GraphSchema(); + graphSchema.setName("week1-service-graph"); + graphSchema.addVertex(schema); + Vertex vertex = new Vertex("chunk", "chunk-1", + Collections.singletonList("Confucius was a philosopher and teacher.")); + Map groups = new HashMap<>(); + groups.put("chunk", new VertexGroup(schema, new ArrayList<>(Collections.singletonList(vertex)))); + return new MemoryGraph(graphSchema, groups); + } + + private static GraphMemoryServer createServer(MemoryGraph graph) { + LocalMemoryGraphAccessor accessor = new LocalMemoryGraphAccessor(graph); + EntityAttributeIndexStore indexStore = new EntityAttributeIndexStore(); + indexStore.initStore(new SubgraphSemanticPromptFunction(accessor)); + GraphMemoryServer server = new GraphMemoryServer(); + server.addGraphAccessor(accessor); + server.addIndexStore(indexStore); + return server; + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestClient.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestClient.java index 707bc6b01..ef5e24743 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestClient.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/support/HttpTestClient.java @@ -66,6 +66,15 @@ public HttpTestResponse post(String path, String body) { return post(path, body, Collections.emptyMap()); } + public HttpTestResponse postWithHeaders(String path, String body, Map headers) { + Request.Builder builder = new Request.Builder().url(url(path, Collections.emptyMap())).post( + RequestBody.create(MediaType.parse("application/json; charset=utf-8"), + body == null ? "" : body)); + (headers == null ? Collections.emptyMap() : headers) + .forEach(builder::header); + return execute(builder.build()); + } + private HttpUrl url(String path, Map parameters) { Objects.requireNonNull(path, "path"); String normalizedPath = path.startsWith("/") ? path : "/" + path; From abbad8b8a416d627bdc4429c9d9bbcb7eca9843f Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 15:49:19 +0800 Subject: [PATCH 06/30] feat(graphrag): add verified dataset source loader --- tools/__init__.py | 1 + tools/__pycache__/__init__.cpython-311.pyc | Bin 0 -> 193 bytes tools/graphrag/__init__.py | 1 + .../__pycache__/__init__.cpython-311.pyc | Bin 0 -> 208 bytes tools/graphrag/ingest/__init__.py | 1 + .../__pycache__/__init__.cpython-311.pyc | Bin 0 -> 219 bytes tools/graphrag/ingest/sources/__init__.py | 5 + .../__pycache__/__init__.cpython-311.pyc | Bin 0 -> 367 bytes .../__pycache__/source_loader.cpython-311.pyc | Bin 0 -> 16286 bytes .../graphrag/ingest/sources/source_loader.py | 197 ++++++++++++++++++ .../manifests/2wikimultihopqa-dev.json | 9 + .../manifests/hotpotqa-distractor.json | 9 + .../test_source_loader.cpython-311.pyc | Bin 0 -> 5817 bytes tools/graphrag/tests/test_source_loader.py | 44 ++++ 14 files changed, 267 insertions(+) create mode 100644 tools/__init__.py create mode 100644 tools/__pycache__/__init__.cpython-311.pyc create mode 100644 tools/graphrag/__init__.py create mode 100644 tools/graphrag/__pycache__/__init__.cpython-311.pyc create mode 100644 tools/graphrag/ingest/__init__.py create mode 100644 tools/graphrag/ingest/__pycache__/__init__.cpython-311.pyc create mode 100644 tools/graphrag/ingest/sources/__init__.py create mode 100644 tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc create mode 100644 tools/graphrag/ingest/sources/__pycache__/source_loader.cpython-311.pyc create mode 100644 tools/graphrag/ingest/sources/source_loader.py create mode 100644 tools/graphrag/manifests/2wikimultihopqa-dev.json create mode 100644 tools/graphrag/manifests/hotpotqa-distractor.json create mode 100644 tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc create mode 100644 tools/graphrag/tests/test_source_loader.py diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 000000000..798f01b5d --- /dev/null +++ b/tools/__init__.py @@ -0,0 +1 @@ +"""Repository tooling packages.""" diff --git a/tools/__pycache__/__init__.cpython-311.pyc b/tools/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4dd5b3811a02d241794564317270bab78979e51a GIT binary patch literal 193 zcmZ3^%ge<81o;))vy_4KV-N=h7@>^MASKfoQW&BbQW%37G?}Vof>I0ei!)2|iz*dL z^7C^t^U@Ux5|gtN(^HG}{4^PFvB$@!&ryk0@&FAkgB{FKt1RJ$TJpjMD0iUonh2WCb_#v2Su L7htG}1t^MASKfoQW&BbQW%37G?}WD+=~(mGJ+i46;cvQ z5)}%HQVSA`5=%1k^At+*^K&xu()Ii_8E>)2$EV~c$H%W^_zW`Ym$80Eer~FMVpUOY zYD#8_L4|&LYGPVWez`tKSFwIN$kd|5bp81F%)HE!_;|g7%3mBdx%nxjIjMF^MASKfoQW&BbQW%37G?}W@{L|8MGV@Xu+=~(m zGJ+i46*BYEQ;SP7^Yavv^K%RG^HTFliuL?78E>)2$EV~c$H%W^_zW`cmz91-er~FM zVpUOYYD#8_L4|&LYGPVWez|^0etu4|emcm+qQrE4h>7~~@tJvb$Lq#k=Q2=wMJy-w$ literal 0 HcmV?d00001 diff --git a/tools/graphrag/ingest/sources/__init__.py b/tools/graphrag/ingest/sources/__init__.py new file mode 100644 index 000000000..d04a2b34d --- /dev/null +++ b/tools/graphrag/ingest/sources/__init__.py @@ -0,0 +1,5 @@ +"""Reproducible offline dataset source readers.""" + +from .source_loader import SourceError, SourceLoader + +__all__ = ["SourceError", "SourceLoader"] diff --git a/tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc b/tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bac663e53f4ee792335470224238b243090bc9e6 GIT binary patch literal 367 zcmZusy-ve05VoC^LPK`Oh%FL|p<+Y~NNk9KZdoi7yJ?UUM`u@*PCNt>yaD0~c#VuG zQza%gq;8#X2rP)Z(|14m?$h0K7>wQ-|B7ssQupkPTBqgPn zdKd{$_&G^^F_@ruYk6=bYtUkz71v5edOB4_B_qKtH_}E%&mof$NG>FpV{OO3sH-49{Adbr9}soN-==u)tu1!4B4$NpQKv0^6UUNKlP?i_TEy9Qn3?m;(A>8M_cGu@^*bGY(DJw<&2f9f$66OWNv&fQ9kLQ+TsV_cBu*l1*k5AwG_6Mb26;YsOK7*^&1rTx5)BLOEWE{zBqI|Yq4@X&9~Oi#$KD9?kzgzVPp3wq z6niTYj&hOM2pbB<;;~357=`YI;ym;u5@UzrlQ9n78VV*tqhaX(H^?PFt=NLGSUeHL z-V1$jZ^2ItMT3F>H_W+UV$^R`^oL`&6s9{8N+{+Ni7+1=iiWePM4VTc(-8rloSjg4 z5Xuh#B(g(Bf9Ham`&5`3q+xpLI2xvt4u(Ll*!ts>d?Mb7 zUY`s`vrkNc02dDh0v}P|VZWdEZXS(~hc^e4{CJp)B(~h#JQ5BLN8?kQ6Y+Rd*gV1q zCr0_;$Y$U*%$v<>7lrIK5RK#f>72O5mqBfWxFI1_1DVPHGEBdiJ9&Ig3SIz~8;ALL zG91(JFuNS+?_yB0z{88MFvGybeUM91`h1EZDFZw;=2NEpQ*%CL!Bm*xD!B67Ov&Og za1IzBr&6wtPi~-asXcx$+d1C?u&d zCU-NeYovJJv@T)G<-q9ZQ}jQ8a7^j`fqswC)|=A3spF_=-5a_odWyQKe~p@=Im5Jm z);LUa#@ptUzMImH%rui?_*PB*CQ`qdv!sjqS7<*~(&mf3c$0q1m^KyOD5SU&(JQp9 zS<R+OxQdG&;IM@1`l5(yeSkApxxi!b6Wc1c7S2A`f zLrD)S*J}CGst;k_e3G9#ATz4_o*O0;!yUUnq7^g9vN-CcWK~v*qtax@vuc-puIpm` z1|DnY*CFvCsY6l^B$QvbaASr)5SFmUA@NK4J~aXyzq)6dE?((x>C(mTKuPsTf!@&? z0opNMT`Fkj*TSw8{f_QV7kneFfbQbGP`9EV2`518CHTk$s54#X7DazOd`r>u!6^YJ z9Lvt|m2iD--4_K#i8}m1vRy08IzjA@2T-KxAs2SdS+?s8ENZTdFxpyJGU+W4$(D9OP>45 zGBzTx;}JoC5$a--x^|Xd5AE=^K$0E#a!mpl6R%H3!1l=MWR&EngbT8q4B9u?OUPZuP#Hs+_AU=E!IeO{*KhXoO=`HC8%}aR&QMrMY-~4DEZd-Dy+a=!) z*|%fGVXm-#O9A;F$qHq*SOwIqf8E+neOi6ER{w>r`mk63h1URSu(ZkaBIX*8^MKz7 z#Bb!wG3`K70Yq^G0!7{iSmHL~XpD!!IpX-ESPmDvY7nEu`w6^v_Rc7_Uf9iScHb@|lw}Cx3o2!_jGZ{VD8HHFRGMy~P&``ZHf4P5mL`BunprsmXU{hvzA2>~+}!vScui9= zlvWBe1dfqnC9Y*ah$O>`9(=M;mM4TUn1(<4yVHa}De(WUlLx$`Sa)+J7WkW>BvMoW z8dZ;!4({^SJ4eCy>ifRI$4OtId3JnKAkHFwPXfGp(8|_J!RTaGZk~Yh{3#%i&j{h@ zFrUxK>T6ZEp!*n!MG}EPvZ|P;ow@7-P>0X~WQNMvs^;mp56m7|cC8m{yCm0M*|k@+ z@6FV;&oZ*vms!_1%iM9y!jDWL#hlY5>Kw;a^SqAu*5FZJL5KPZ`rGI|wX@)UU%h%& zW6D>d)?Uk_3PxP?~@+iZM8WxRBxusun-$hWCW8fT%W4#^HYAXjpPI zq>oFE4%yK$%Vf;{ChGk!4`ogcT{_5(&cByH%+_d{K zE&7f~z9X^^)FET>5%Sil*`B$h8N2frr{6w3U-_hF^HR;`1@mI_;|8gwN3Q9S?7gzR zSG4yoJ6v-+-b&s{&R=*CSaP(Bj`mEoZ_b9ljCb9f;aSF3fqK@Pad~D><$kF2y(F6~ zqvziIk||j&n#F<>mMoT&koa+n1tzT(7L4u~Oo*T1gwg#C1NN#7r^iP$AZC03>I7Gy zFmP8PLn}xa9EPOxg^|H`gyTT6BT`lKq^fzTsyTgAs@ft~ZNU$%C1mzk#^7(5l)uLLoi%ol%UUflh z9V_JU>d=L@91rFCK?D^RGY`>@#K7S#PCcMS96!%)&+89DvF3&k{AV=o+> zfirw+EEW0yo}s&zJ;u;_BS*RjtjXmOuI5JopeBvbr>WQufXx~Ur2!smoi-GY3WsC( zvT-q|3;NF?!%Ry*XKx%_um$Y%swl zckgAxF$A>PoI!(@5P`O^LCK~k1@PcZ5rNAT@1`a1 zrboNwwgZy)pzJ*;xev+iLr>hNmfWW#_ZiuJX0{uY(l1hPr`~$w&Kn^2%GPH*pc5cd zc5Gas43_dot@5V*B$e^h%brc5=b+emaP~O(N>ANB+1((LGd=LZwKV(|2Ji>+(RuJ& zz205^Jp}|hRaX9IGv%uO<)vb@bvVWiGzxEaLctAcjART&5#`Bs~9gwR7 zlJlps^QTXo*O#2vCFi*89G^92%;k5=o|t`0W?#B{VY_JdN#@-n?W7c`mnfErRkpD%+2W}S-PFB@`f`u;SeyPUeZ#S4{a4KfNUKH# z(+3WKMnRwu1Q7=tfG~K7o)=r?1?*fajHh)Lv<8Va1@CJPuEgftLA`BEN0{rW6B0$*dL2DTW9o<_)GuDCvF}cAdg@ z@^-n3RrBq|J^N4kWhVWy6xRa%v=ZTI(f0FA0N+4V+WHh7KNKo(m`no|WJM?lLDsjZ zbn)|ix^#%4I$d9cqACy^We~uEs}cN{aBy6}&&r~p=6%uN_z)L7kld`5B_(EbF%YW0 z5QRzv)aC#&J!<~E?c=tE>(Ztp@}?s*S z@Wm8x8Cdftr}ti+gv=ewRIa;!=I)tvlT_I%SGJ1eL}}_@(L)`iB0|2!D!)fk$XhLi zu-z)9o;~}sk*e^(Bj{@}CGs#+;gu^`k(}w)4?5HETWG}JqJB{i-KpRIEk(Nze2*k^ zvcA4V zXzK-abd5k3hG}EUn9lnOxt8#%Rf_yEu$_y_6|DJG2}iO!AFpWDQNdmZ#pIhdX}-&v zg00GwX;}Z_*QHEE@tEE)6%@~%!0M_IMgQu`QGJp?6gGQ6FCYTwO)7yFX!a;ppN0Nr zH@>w!wh@(NPnllaa$c>+{<6NLzH=lAS}DP!&%ie!K@5h+$^4wYl58Y#GJ+Zz0vl5V z?F+?Yi6F$e1mVO88vY`bPS)kNXVlP|1PjHZ5dc4%xGqR1E3=w{1RrQrgx>+GYjSuv za#LY2@`g`g7j6PU$S{1J-+}k2I2CIIP*aTT3568~x9n6s@ER6qL$VXe9wfy4HWOW- z7*UG=q=O<~F$_m|u+@zK(Sp&g*szlU74yM5p&YTu(aTbipX~M*KL|Oc=wudg!p;g? zfNLta@U(8@gX@njF4b)n>$d)~@?Onc&F>l150#dF&<~!u>Un0KdvsZJ?vR{2Wao|* z%2C!sFIH!~b?@8WwWSaLeffj(IYXwR`o8V1ZT|2(<@d_x$}?{7{L#Dn=JtU)dGK2L zhE&@r*LKeJEc=?$!=G`IZ=3Af2I`~=Hel-f;-+(A|7B^@WvOmJt{V`&1DX0Jv2ykE zPaQiKZajWf?C6y`dgYGZ70OuE1L^^(&C23Ms-m1eiUjy^sr)*Y8zCuG-&uLaS59vUY5Y5n+30wP+!D(~J* zEgY%u-l+d$MfK4Z=8w%Zrdu4{PSYQER(IP?U)l|j{?f@n{+CrYDF5X~rn}knWwRdB zej3x8NP07SY?J;gNAzpn#GW&vwg8RuGPxAyw?Zy z6zS=gX)M?twx^71a(lg&hp@KS((bh-CFv20t9xJpbgRBExuHqe3_gTvVS^Wu4>bV4 zBR(_+z>%U4L^wqUe>Q@@1FDW0=G;xPv9t;ZXOBV@JDm9-J1)6t^-{T7ipwe{bbc^Y z%sc^M`#}Jlk~>#_JzKEQYZm@hpd&m!k+=nZk&yTXeJtW_J*Qa>{(z&f4e^KVkjUW= zSKZps%hEuDtcP2PUL`*SFMdRmAhnOO(R^kB_D&!^G+Z3f$72i0SF~tX>6UpBHkga( zL*t}4qObnASbuz>^`Sm}{2$9cD1)1lQ;owgBw(K>%puCj&4P|WT6&rdNMjsx1Wd6r9ADn@kl9PORa@R|z$dxAU zDvJS*lQCg(Vj|8XA__r{6XCU{W_E6o7hxuoNd;k>??cDvcPooo&ASY=hR{hT6!K%Lb70{{$KTFPOfF^=tu>W%}kJ5zcd$ zesgm1##e&0{gk}@l(^*-W@=BB`JJmXU*!zpdk(oEJXb2Rg8<=KfnmUkal;WeC1LIO zEhX`#N|`F2a`=G_{=EglQDO;@L;T+%Lp}haSo1Sn;BoEvD+sp~hJ&!95-7V6U&^1x z3J9$6Uob^bj(wOy?32HOWIvL1K>THDBul~9`v@PAKHSID$3PTA1pYG#U>3mB$H{b2 zEP1&=W)A7EfZ92SeUK{&e&K%t*~z+>ON)1)j1U762H3cxX1?LQL@u;bXW0vWu~>Il z#h4LYUT6k4{}|l`o_{?O_g=)6_kt5$w)bgu&Ha(PBlF>RV)tTV_1Q=I&rKhj7MPF9 zJ}raWMfbmYe&=~yEA2Te?>Reb%ebrW@4CBdzDsg9%kE}~|Af}tKGT2JpH51St+Hb) zV%>Xx->`UH>N+WRoqW=DZmH{>)OB9&I!_*v&UKe`?)d7R!%u5{cvNz<%8pj@_JaAh z<^NnRZ8;=wIaH|bWJz^>t5^5bRdv7RZp*x1ay81XM$xr5`?Xf7b+6pIch-bo)ADe8 zaf7BpGnY!58C;{8D$lGHgNrUpz7E;f0r-cdhX(irC`1z-o7;}fG(Kb`N0aPmTA{3f zw0zd^i2KO@sej@1k2@cA{(kUl=IdAg-Sj_Sy(sO!B=5iUs58^hB2}<+hPhi9?RxX> zO;z?UM&K&dAD4l8j?13obIcsG?DWldKWvnoEwZy^g)%|?8Snb{&F`Ai9)Ou-uV3^Y zeRSdTfsY3khNX`Ea>sticc&i%!NSl11w zWIra`kBRnUkjP}`PpcY)EV}ljgwS%?#V{YtyexVr!Sg?v>fSB0I1s{Oi>3 zrXDAygXiUg=OIJ9G_c~fVz9JIELI}w`~12d{x}FA2$DVl7ycxYQ%F#hLM{pR^I*mX zB>k)A{&dw`)S3{^!NNrizAmXNp}tjkPFnbD%rzIaOXM>V?6~Jk06DcUC^-H0k}z%_ zV}STqkswe=KyVTpqSyfiPeel5c2TG5poAPOW$Fw)y~`zLBHOP*I75H>B) zd=Zkhg_PNb(c(FhFObV1kH`!lU+MP}w&DMw4o||rGs6F&^0>ps&m#F560)N(3Mr*r zjp6CchVbxl%(UhIF*hvmpvx!>K1z^XX`BWSx^rs+%ma1wa0{6?Bf*VB-ipMAbZHJ`k%4~_}{=Sy7U6Z7=!eo5niVm7!x3N5JMuhvG{yS*=EQu zLzT~vU#5^_nIXRnRVC)m47EYjo*Bw3YR?RHNc=J9ijk(tf0C?9=Qe@#geM?33(w7z z&Nfps*C^5+iT22}XGO=-Ft#s}Z!zzCBrx7OZ@y>`eNu;StWLir`QNYclyLt40#rd> AZ~y=R literal 0 HcmV?d00001 diff --git a/tools/graphrag/ingest/sources/source_loader.py b/tools/graphrag/ingest/sources/source_loader.py new file mode 100644 index 000000000..6406a965d --- /dev/null +++ b/tools/graphrag/ingest/sources/source_loader.py @@ -0,0 +1,197 @@ +"""Manifest-driven readers for the Week 2 public datasets. + +The reader intentionally depends only on the Python standard library. It accepts JSON arrays, +JSONL files, and gzip-compressed variants, while yielding canonical records in bounded batches. +""" + +from __future__ import annotations + +import gzip +import hashlib +import json +import os +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Iterable, Iterator, List, Optional + + +class SourceError(ValueError): + """Raised when a source cannot be verified or parsed into canonical records.""" + + +@dataclass(frozen=True) +class SourceManifest: + dataset: str + dataset_release: str + split: str + source_uri: Optional[str] + cache_path: Optional[str] + sha256: str + parser_version: str + + @classmethod + def from_file(cls, path: Path) -> "SourceManifest": + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise SourceError("cannot read manifest %s: %s" % (path, exc)) from exc + required = ("dataset", "dataset_release", "split", "sha256", "parser_version") + missing = [key for key in required if not str(raw.get(key, "")).strip()] + if missing: + raise SourceError("manifest %s is missing: %s" % (path, ", ".join(missing))) + if not raw.get("source_uri") and not raw.get("cache_path"): + raise SourceError("manifest %s requires source_uri or cache_path" % path) + checksum = str(raw["sha256"]).lower() + if len(checksum) != 64 or any(char not in "0123456789abcdef" for char in checksum): + raise SourceError("manifest %s has invalid sha256" % path) + return cls( + dataset=str(raw["dataset"]), + dataset_release=str(raw["dataset_release"]), + split=str(raw["split"]), + source_uri=raw.get("source_uri"), + cache_path=raw.get("cache_path"), + sha256=checksum, + parser_version=str(raw["parser_version"]), + ) + + +class SourceLoader: + """Loads and validates one manifest without exposing dataset-specific records.""" + + def __init__(self, allow_download: bool = False, batch_size: int = 256): + if batch_size < 1: + raise ValueError("batch_size must be positive") + self.allow_download = allow_download + self.batch_size = batch_size + + def load_batches(self, manifest: SourceManifest) -> Iterator[List[Dict[str, Any]]]: + path = self._resolve_source(manifest) + self._verify_checksum(path, manifest) + batch: List[Dict[str, Any]] = [] + for number, raw in enumerate(self._records(path), start=1): + batch.append(self._canonical_record(manifest, number, raw)) + if len(batch) == self.batch_size: + yield batch + batch = [] + if batch: + yield batch + + def load(self, manifest: SourceManifest) -> Iterator[Dict[str, Any]]: + for batch in self.load_batches(manifest): + yield from batch + + def _resolve_source(self, manifest: SourceManifest) -> Path: + if manifest.cache_path: + path = Path(os.path.expandvars(manifest.cache_path)).expanduser() + if path.is_file(): + return path + if not self.allow_download or not manifest.source_uri: + location = manifest.cache_path or "" + raise SourceError( + "%s/%s cache is missing at %s; enable download explicitly or provide a verified cache" + % (manifest.dataset, manifest.split, location) + ) + target = Path(os.path.expandvars(manifest.cache_path or Path(manifest.source_uri).name)).expanduser() + target.parent.mkdir(parents=True, exist_ok=True) + try: + urllib.request.urlretrieve(manifest.source_uri, target) + except OSError as exc: + raise SourceError("download failed for %s/%s: %s" % (manifest.dataset, manifest.split, exc)) from exc + return target + + @staticmethod + def _verify_checksum(path: Path, manifest: SourceManifest) -> None: + digest = hashlib.sha256() + try: + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + except OSError as exc: + raise SourceError("cannot read %s/%s source %s: %s" % ( + manifest.dataset, manifest.split, path, exc)) from exc + actual = digest.hexdigest() + if actual != manifest.sha256: + raise SourceError( + "%s/%s checksum mismatch for %s: expected %s, got %s" + % (manifest.dataset, manifest.split, path, manifest.sha256, actual) + ) + + @staticmethod + def _records(path: Path) -> Iterable[Any]: + opener = gzip.open if path.suffix == ".gz" else open + try: + with opener(path, "rt", encoding="utf-8") as stream: + first = stream.read(1) + stream.seek(0) + if first == "[": + values = json.load(stream) + if not isinstance(values, list): + raise SourceError("JSON source %s must contain an array" % path) + yield from values + else: + for line_number, line in enumerate(stream, start=1): + if line.strip(): + try: + yield json.loads(line) + except json.JSONDecodeError as exc: + raise SourceError("invalid JSON at %s line %d: %s" % ( + path, line_number, exc)) from exc + except SourceError: + raise + except (OSError, json.JSONDecodeError) as exc: + raise SourceError("cannot parse source %s: %s" % (path, exc)) from exc + + @staticmethod + def _canonical_record(manifest: SourceManifest, number: int, raw: Any) -> Dict[str, Any]: + if not isinstance(raw, dict): + raise SourceError("%s/%s record %d must be an object" % ( + manifest.dataset, manifest.split, number)) + document_id = raw.get("_id") or raw.get("id") or raw.get("question_id") + question = raw.get("question") + if not isinstance(document_id, (str, int)) or not str(document_id).strip(): + raise SourceError("%s/%s record %d field id is required" % ( + manifest.dataset, manifest.split, number)) + if not isinstance(question, str) or not question.strip(): + raise SourceError("%s/%s record %d field question is required" % ( + manifest.dataset, manifest.split, number)) + context = raw.get("context") or raw.get("passages") + if not isinstance(context, list) or not context: + raise SourceError("%s/%s record %d field context must be a non-empty list" % ( + manifest.dataset, manifest.split, number)) + paragraphs = [] + for index, item in enumerate(context): + if isinstance(item, list) and len(item) == 2 and isinstance(item[1], list): + title, sentences = item + text = " ".join(str(sentence) for sentence in sentences) + elif isinstance(item, dict): + title = item.get("title") or item.get("name") or "" + text = item.get("text") or item.get("paragraph") or "" + if isinstance(text, list): + text = " ".join(str(sentence) for sentence in text) + else: + raise SourceError("%s/%s record %d field context[%d] has unsupported shape" % ( + manifest.dataset, manifest.split, number, index)) + if not isinstance(title, str) or not isinstance(text, str) or not text.strip(): + raise SourceError("%s/%s record %d field context[%d] is empty" % ( + manifest.dataset, manifest.split, number, index)) + paragraphs.append({"title": title.strip(), "text": text}) + answers = raw.get("answer") or raw.get("answers") + if isinstance(answers, str): + answers = [answers] + if answers is None: + answers = [] + if not isinstance(answers, list) or any(not isinstance(answer, str) for answer in answers): + raise SourceError("%s/%s record %d field answer has invalid shape" % ( + manifest.dataset, manifest.split, number)) + return { + "dataset": manifest.dataset, + "dataset_release": manifest.dataset_release, + "split": manifest.split, + "record_number": number, + "document_id": str(document_id), + "question": question, + "answers": answers, + "paragraphs": paragraphs, + "supporting_facts": raw.get("supporting_facts") or raw.get("evidence") or [], + } diff --git a/tools/graphrag/manifests/2wikimultihopqa-dev.json b/tools/graphrag/manifests/2wikimultihopqa-dev.json new file mode 100644 index 000000000..180e49332 --- /dev/null +++ b/tools/graphrag/manifests/2wikimultihopqa-dev.json @@ -0,0 +1,9 @@ +{ + "dataset": "2wikimultihopqa", + "dataset_release": "official", + "split": "dev", + "source_uri": null, + "cache_path": "${GRAPH_RAG_CACHE}/2wikimultihopqa-dev.json", + "sha256": "REPLACE_WITH_OFFICIAL_SHA256", + "parser_version": "2wikimultihopqa-parser-v1" +} diff --git a/tools/graphrag/manifests/hotpotqa-distractor.json b/tools/graphrag/manifests/hotpotqa-distractor.json new file mode 100644 index 000000000..415cbba99 --- /dev/null +++ b/tools/graphrag/manifests/hotpotqa-distractor.json @@ -0,0 +1,9 @@ +{ + "dataset": "hotpotqa", + "dataset_release": "official-distractor", + "split": "distractor", + "source_uri": null, + "cache_path": "${GRAPH_RAG_CACHE}/hotpotqa-distractor.json", + "sha256": "REPLACE_WITH_OFFICIAL_SHA256", + "parser_version": "hotpotqa-parser-v1" +} diff --git a/tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc b/tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3213b7edf7e304f99fddfd2d0fb51500befce05b GIT binary patch literal 5817 zcmcgwO>7iL7OtM|`EMH!WEil6*GV?uA&U(dn_#mfYYaHq;4EPCW3sfGskYr_ew^xd zFe5wbNF02~VUbvYBBGrW>j)zGv`eH!@&*nZXRMZNwS*LDFPobY9FR}@Uib9O^mqu_ zBuaJ9>sMW`s$RW%@9S6dkGi@Lfs&c}*Z5ljLjHplr+ACZlh1&;K`5bIk__8VZkV&` zo?(ymbvB@8jE-cNY)jb$ZcTmI28hOaHp(a-2Xkk9f)(L@F!Uxi}PmqRA!PTlp~B~ zm?_$o_cAh7m|1Vvj>~CvMA7s_X+U7OV&zE-FgJ)oaFMV;?ka)hd4LL(2kNE5DH8XZ z4V5Dg;%LqYI1nxo7xqMl<=h~8Mca6ZcQ2}w%y7D6cIZ0I-3G}!C27iK!-ULusE6`Y zpx#k#)HBLaANAi7Xn+Q3hZw8|fxW z^l16rw8uT_v_N0F4Z7SZ=`-WiV`X`LYni9fyPMs;b0tSCOTi9#-BRw_jBlFc>MP|P zp5h(_aQ|owtnFK8?JC83~C6BkU6O3lEBliB#sH9o&ASEqGXq&jQ7Cbnrq?M@&b|`zaI;h*Qqa~|J zmK7y6Eor)ZF{wxwWj!&bXl)0RnS`9w4z-oWLEqGV2jstS&650%6L8=A=u6>lXrZgn z+HJIU+wUsr6Obos-bsCA*PCB(WLvERU(wa~$&3J94=O&rRe{CPGm4Fq;~#>6<*aB8r1+mL-P z>ht=-aihJr^oG3x608k~qo~(1W2L5SH#BFD{X!~QdbQem?Yfk=?_tT-=;0(h3vELN zXzVmvD#kaM0Wjdih?-Q)jf22tn9Q#9so?v1hFxI@=w|3L12aojunOZD6=HWx*2a?R zMboE^$$LAy%n(y#DqXw+K5PcZlqsr?f;TfnXOjoU6wrbUzd5Fe@!+_$8J0CoVfv8| zC*@?^Zwi`{95I8`(Ew(!g2?n2!_f>tBqT0rlPNQhDt=u}e=$HbgxnUpSngoH`r@d{ zcKn1r5Md$>QFxX-s*4(Rt%bUsM%~U;5(=icdG%3r>&^Hz>AIBrOCj1dcl5hhGx_ebcpo^I9~sXpm+}osxL3V=tbXp;YJk*j zS{Ao0iQ8_{zpJ0Ab~HbX!LzXAgaNeqgwcE=|3M-zZYzk?5NTedk7CW&KhBLTZ{NSP zeScy50b~1tLaf_}buY*ImSTN{*ij>PbgutV?Dbq?zGH!382qN=-bY{ceBJXfkUu$8 zh@CcKr=fjwOTM-+U#g~`F zmv3q}m0Th>bn`;d*%ITx6n0%QfQm^&Oy=w2Hqw~K$UOX%CcXQ^BK1$$)?YS}dTwe23D#Yq`g5xXLD^z~L<`a` zNH{8zp6Y}HQ?jargNK9wgKu5YE7AlSa?jieA5uM1?oyffxRat>DX%U|Lz2R~(-L>h z%GPTn7wPcxQj4FEii8KUVr43l4Wu)A3w}=7tDq`-4aw_B@Uz8s05RX*{WA+g+|5QD zTxv^mu|nxdUlMwKasnGNefVqv9#SQ!@> z<~i&wByS@*3`9dwggP7O-mrkq(Lr)_kVa%RsYw@=kqlF$37Ki|ozEVujd(wO`98{Aikd>V8 zMBZ5?T-`p4JZgC9Zy#K{aQ(t^v~?-kI=`h5-D^bm78*K?hK_tg$13kgTH=qfgN2fyv+k^eo1{2}D;*9(U?`tL_MptbR90P>+<2fr$jYP@uy#ABmceWVobdfjLP zSW|J-8801{04smQvCNGvxeCMtV-|}Wvu60pLz{)F0l?LYgIP7MPV?(kk3#*hM%HZaBA{Y`N#&tZL}Tu z78{CN_Oy(sN|LrXXTNx)CNMVPxGy8Yh+zAXyouxhk`ipPeyqBIlDz}(;@)bU>_tfn zEYUd;EQ@Hl-9pQhoCH{e^PFTW3J*wtbq!8o2^kNm+5NSF@`W8S|JSB#V5HBDk#=@i z!pP6*t?~KaF7WeBh2~D9xpUEPA!ffJ_UA>&)Bdk?&by$0b82z(SLeSz|FAVbc%~3L zYsAjZ^}B#|un-wCB12^u{FQ+<3ZvG)#+JFzy{Xdj=o0zzQ4WC zEP;KFqyg{aw-K{9Vxo=>Q0|*{3Y6 z_$>ZqVFww8j(w5~41(xk8P03g$_8*A0TrlP_*&qgZLdS@49Gwj5baAK@KNTt6;eNI zT`MF!Yh5d3!>n~xFkAA(b$>_X{PXWsUo!``T2KBuWyD$^W5d&OYcEFxaW1>WZ_V>t bSA@uH|JCE4ynprmS?gNmiO>iQmYM$tw?bN% literal 0 HcmV?d00001 diff --git a/tools/graphrag/tests/test_source_loader.py b/tools/graphrag/tests/test_source_loader.py new file mode 100644 index 000000000..1a4f16e19 --- /dev/null +++ b/tools/graphrag/tests/test_source_loader.py @@ -0,0 +1,44 @@ +import hashlib +import json +import tempfile +import unittest +from pathlib import Path + +from tools.graphrag.ingest.sources.source_loader import SourceError, SourceLoader, SourceManifest + + +class SourceLoaderTest(unittest.TestCase): + + def test_loads_jsonl_in_stable_batches(self): + with tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "records.jsonl" + records = [ + {"id": "a", "question": "q1", "context": [["T", ["text"]]], "answer": "a1"}, + {"id": "b", "question": "q2", "context": [{"title": "U", "text": "text"}], "answer": ["a2"]}, + ] + source.write_text("\n".join(json.dumps(record) for record in records), encoding="utf-8") + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + manifest = SourceManifest("d", "r", "s", None, str(source), checksum, "p") + batches = list(SourceLoader(batch_size=1).load_batches(manifest)) + self.assertEqual([["a"], ["b"]], [[item["document_id"] for item in batch] for batch in batches]) + + def test_checksum_mismatch_fails_before_parsing(self): + with tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "records.json" + source.write_text("not json", encoding="utf-8") + manifest = SourceManifest("d", "r", "s", None, str(source), "0" * 64, "p") + with self.assertRaisesRegex(SourceError, "checksum mismatch"): + list(SourceLoader().load(manifest)) + + def test_malformed_record_reports_location(self): + with tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "records.jsonl" + source.write_text(json.dumps({"id": "a"}), encoding="utf-8") + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + manifest = SourceManifest("d", "r", "dev", None, str(source), checksum, "p") + with self.assertRaisesRegex(SourceError, "d/dev record 1 field question"): + list(SourceLoader().load(manifest)) + + +if __name__ == "__main__": + unittest.main() From ceb101d19f37bd7cf39fe0fc01f7dd30741e2310 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 15:49:44 +0800 Subject: [PATCH 07/30] chore(graphrag): remove generated Python caches --- tools/__pycache__/__init__.cpython-311.pyc | Bin 193 -> 0 bytes .../__pycache__/__init__.cpython-311.pyc | Bin 208 -> 0 bytes .../ingest/__pycache__/__init__.cpython-311.pyc | Bin 219 -> 0 bytes .../__pycache__/__init__.cpython-311.pyc | Bin 367 -> 0 bytes .../__pycache__/source_loader.cpython-311.pyc | Bin 16286 -> 0 bytes .../test_source_loader.cpython-311.pyc | Bin 5817 -> 0 bytes 6 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 tools/__pycache__/__init__.cpython-311.pyc delete mode 100644 tools/graphrag/__pycache__/__init__.cpython-311.pyc delete mode 100644 tools/graphrag/ingest/__pycache__/__init__.cpython-311.pyc delete mode 100644 tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc delete mode 100644 tools/graphrag/ingest/sources/__pycache__/source_loader.cpython-311.pyc delete mode 100644 tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc diff --git a/tools/__pycache__/__init__.cpython-311.pyc b/tools/__pycache__/__init__.cpython-311.pyc deleted file mode 100644 index 4dd5b3811a02d241794564317270bab78979e51a..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 193 zcmZ3^%ge<81o;))vy_4KV-N=h7@>^MASKfoQW&BbQW%37G?}Vof>I0ei!)2|iz*dL z^7C^t^U@Ux5|gtN(^HG}{4^PFvB$@!&ryk0@&FAkgB{FKt1RJ$TJpjMD0iUonh2WCb_#v2Su L7htG}1t^MASKfoQW&BbQW%37G?}WD+=~(mGJ+i46;cvQ z5)}%HQVSA`5=%1k^At+*^K&xu()Ii_8E>)2$EV~c$H%W^_zW`Ym$80Eer~FMVpUOY zYD#8_L4|&LYGPVWez`tKSFwIN$kd|5bp81F%)HE!_;|g7%3mBdx%nxjIjMF^MASKfoQW&BbQW%37G?}W@{L|8MGV@Xu+=~(m zGJ+i46*BYEQ;SP7^Yavv^K%RG^HTFliuL?78E>)2$EV~c$H%W^_zW`cmz91-er~FM zVpUOYYD#8_L4|&LYGPVWez|^0etu4|emcm+qQrE4h>7~~@tJvb$Lq#k=Q2=wMJy-w$ diff --git a/tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc b/tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc deleted file mode 100644 index bac663e53f4ee792335470224238b243090bc9e6..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 367 zcmZusy-ve05VoC^LPK`Oh%FL|p<+Y~NNk9KZdoi7yJ?UUM`u@*PCNt>yaD0~c#VuG zQza%gq;8#X2rP)Z(|14m?$h0K7>wQ-|B7ssQupkPTBqgPn zdKd{$_&G^^F_@ruYk6=bYtUkz71v5edOB4_B_qKtH_}E%&mof$NG>FpV{OO3sH-49{Adbr9}soN-==u)tu1!4B4$NpQKv0^6UUNKlP?i_TEy9Qn3?m;(A>8M_cGu@^*bGY(DJw<&2f9f$66OWNv&fQ9kLQ+TsV_cBu*l1*k5AwG_6Mb26;YsOK7*^&1rTx5)BLOEWE{zBqI|Yq4@X&9~Oi#$KD9?kzgzVPp3wq z6niTYj&hOM2pbB<;;~357=`YI;ym;u5@UzrlQ9n78VV*tqhaX(H^?PFt=NLGSUeHL z-V1$jZ^2ItMT3F>H_W+UV$^R`^oL`&6s9{8N+{+Ni7+1=iiWePM4VTc(-8rloSjg4 z5Xuh#B(g(Bf9Ham`&5`3q+xpLI2xvt4u(Ll*!ts>d?Mb7 zUY`s`vrkNc02dDh0v}P|VZWdEZXS(~hc^e4{CJp)B(~h#JQ5BLN8?kQ6Y+Rd*gV1q zCr0_;$Y$U*%$v<>7lrIK5RK#f>72O5mqBfWxFI1_1DVPHGEBdiJ9&Ig3SIz~8;ALL zG91(JFuNS+?_yB0z{88MFvGybeUM91`h1EZDFZw;=2NEpQ*%CL!Bm*xD!B67Ov&Og za1IzBr&6wtPi~-asXcx$+d1C?u&d zCU-NeYovJJv@T)G<-q9ZQ}jQ8a7^j`fqswC)|=A3spF_=-5a_odWyQKe~p@=Im5Jm z);LUa#@ptUzMImH%rui?_*PB*CQ`qdv!sjqS7<*~(&mf3c$0q1m^KyOD5SU&(JQp9 zS<R+OxQdG&;IM@1`l5(yeSkApxxi!b6Wc1c7S2A`f zLrD)S*J}CGst;k_e3G9#ATz4_o*O0;!yUUnq7^g9vN-CcWK~v*qtax@vuc-puIpm` z1|DnY*CFvCsY6l^B$QvbaASr)5SFmUA@NK4J~aXyzq)6dE?((x>C(mTKuPsTf!@&? z0opNMT`Fkj*TSw8{f_QV7kneFfbQbGP`9EV2`518CHTk$s54#X7DazOd`r>u!6^YJ z9Lvt|m2iD--4_K#i8}m1vRy08IzjA@2T-KxAs2SdS+?s8ENZTdFxpyJGU+W4$(D9OP>45 zGBzTx;}JoC5$a--x^|Xd5AE=^K$0E#a!mpl6R%H3!1l=MWR&EngbT8q4B9u?OUPZuP#Hs+_AU=E!IeO{*KhXoO=`HC8%}aR&QMrMY-~4DEZd-Dy+a=!) z*|%fGVXm-#O9A;F$qHq*SOwIqf8E+neOi6ER{w>r`mk63h1URSu(ZkaBIX*8^MKz7 z#Bb!wG3`K70Yq^G0!7{iSmHL~XpD!!IpX-ESPmDvY7nEu`w6^v_Rc7_Uf9iScHb@|lw}Cx3o2!_jGZ{VD8HHFRGMy~P&``ZHf4P5mL`BunprsmXU{hvzA2>~+}!vScui9= zlvWBe1dfqnC9Y*ah$O>`9(=M;mM4TUn1(<4yVHa}De(WUlLx$`Sa)+J7WkW>BvMoW z8dZ;!4({^SJ4eCy>ifRI$4OtId3JnKAkHFwPXfGp(8|_J!RTaGZk~Yh{3#%i&j{h@ zFrUxK>T6ZEp!*n!MG}EPvZ|P;ow@7-P>0X~WQNMvs^;mp56m7|cC8m{yCm0M*|k@+ z@6FV;&oZ*vms!_1%iM9y!jDWL#hlY5>Kw;a^SqAu*5FZJL5KPZ`rGI|wX@)UU%h%& zW6D>d)?Uk_3PxP?~@+iZM8WxRBxusun-$hWCW8fT%W4#^HYAXjpPI zq>oFE4%yK$%Vf;{ChGk!4`ogcT{_5(&cByH%+_d{K zE&7f~z9X^^)FET>5%Sil*`B$h8N2frr{6w3U-_hF^HR;`1@mI_;|8gwN3Q9S?7gzR zSG4yoJ6v-+-b&s{&R=*CSaP(Bj`mEoZ_b9ljCb9f;aSF3fqK@Pad~D><$kF2y(F6~ zqvziIk||j&n#F<>mMoT&koa+n1tzT(7L4u~Oo*T1gwg#C1NN#7r^iP$AZC03>I7Gy zFmP8PLn}xa9EPOxg^|H`gyTT6BT`lKq^fzTsyTgAs@ft~ZNU$%C1mzk#^7(5l)uLLoi%ol%UUflh z9V_JU>d=L@91rFCK?D^RGY`>@#K7S#PCcMS96!%)&+89DvF3&k{AV=o+> zfirw+EEW0yo}s&zJ;u;_BS*RjtjXmOuI5JopeBvbr>WQufXx~Ur2!smoi-GY3WsC( zvT-q|3;NF?!%Ry*XKx%_um$Y%swl zckgAxF$A>PoI!(@5P`O^LCK~k1@PcZ5rNAT@1`a1 zrboNwwgZy)pzJ*;xev+iLr>hNmfWW#_ZiuJX0{uY(l1hPr`~$w&Kn^2%GPH*pc5cd zc5Gas43_dot@5V*B$e^h%brc5=b+emaP~O(N>ANB+1((LGd=LZwKV(|2Ji>+(RuJ& zz205^Jp}|hRaX9IGv%uO<)vb@bvVWiGzxEaLctAcjART&5#`Bs~9gwR7 zlJlps^QTXo*O#2vCFi*89G^92%;k5=o|t`0W?#B{VY_JdN#@-n?W7c`mnfErRkpD%+2W}S-PFB@`f`u;SeyPUeZ#S4{a4KfNUKH# z(+3WKMnRwu1Q7=tfG~K7o)=r?1?*fajHh)Lv<8Va1@CJPuEgftLA`BEN0{rW6B0$*dL2DTW9o<_)GuDCvF}cAdg@ z@^-n3RrBq|J^N4kWhVWy6xRa%v=ZTI(f0FA0N+4V+WHh7KNKo(m`no|WJM?lLDsjZ zbn)|ix^#%4I$d9cqACy^We~uEs}cN{aBy6}&&r~p=6%uN_z)L7kld`5B_(EbF%YW0 z5QRzv)aC#&J!<~E?c=tE>(Ztp@}?s*S z@Wm8x8Cdftr}ti+gv=ewRIa;!=I)tvlT_I%SGJ1eL}}_@(L)`iB0|2!D!)fk$XhLi zu-z)9o;~}sk*e^(Bj{@}CGs#+;gu^`k(}w)4?5HETWG}JqJB{i-KpRIEk(Nze2*k^ zvcA4V zXzK-abd5k3hG}EUn9lnOxt8#%Rf_yEu$_y_6|DJG2}iO!AFpWDQNdmZ#pIhdX}-&v zg00GwX;}Z_*QHEE@tEE)6%@~%!0M_IMgQu`QGJp?6gGQ6FCYTwO)7yFX!a;ppN0Nr zH@>w!wh@(NPnllaa$c>+{<6NLzH=lAS}DP!&%ie!K@5h+$^4wYl58Y#GJ+Zz0vl5V z?F+?Yi6F$e1mVO88vY`bPS)kNXVlP|1PjHZ5dc4%xGqR1E3=w{1RrQrgx>+GYjSuv za#LY2@`g`g7j6PU$S{1J-+}k2I2CIIP*aTT3568~x9n6s@ER6qL$VXe9wfy4HWOW- z7*UG=q=O<~F$_m|u+@zK(Sp&g*szlU74yM5p&YTu(aTbipX~M*KL|Oc=wudg!p;g? zfNLta@U(8@gX@njF4b)n>$d)~@?Onc&F>l150#dF&<~!u>Un0KdvsZJ?vR{2Wao|* z%2C!sFIH!~b?@8WwWSaLeffj(IYXwR`o8V1ZT|2(<@d_x$}?{7{L#Dn=JtU)dGK2L zhE&@r*LKeJEc=?$!=G`IZ=3Af2I`~=Hel-f;-+(A|7B^@WvOmJt{V`&1DX0Jv2ykE zPaQiKZajWf?C6y`dgYGZ70OuE1L^^(&C23Ms-m1eiUjy^sr)*Y8zCuG-&uLaS59vUY5Y5n+30wP+!D(~J* zEgY%u-l+d$MfK4Z=8w%Zrdu4{PSYQER(IP?U)l|j{?f@n{+CrYDF5X~rn}knWwRdB zej3x8NP07SY?J;gNAzpn#GW&vwg8RuGPxAyw?Zy z6zS=gX)M?twx^71a(lg&hp@KS((bh-CFv20t9xJpbgRBExuHqe3_gTvVS^Wu4>bV4 zBR(_+z>%U4L^wqUe>Q@@1FDW0=G;xPv9t;ZXOBV@JDm9-J1)6t^-{T7ipwe{bbc^Y z%sc^M`#}Jlk~>#_JzKEQYZm@hpd&m!k+=nZk&yTXeJtW_J*Qa>{(z&f4e^KVkjUW= zSKZps%hEuDtcP2PUL`*SFMdRmAhnOO(R^kB_D&!^G+Z3f$72i0SF~tX>6UpBHkga( zL*t}4qObnASbuz>^`Sm}{2$9cD1)1lQ;owgBw(K>%puCj&4P|WT6&rdNMjsx1Wd6r9ADn@kl9PORa@R|z$dxAU zDvJS*lQCg(Vj|8XA__r{6XCU{W_E6o7hxuoNd;k>??cDvcPooo&ASY=hR{hT6!K%Lb70{{$KTFPOfF^=tu>W%}kJ5zcd$ zesgm1##e&0{gk}@l(^*-W@=BB`JJmXU*!zpdk(oEJXb2Rg8<=KfnmUkal;WeC1LIO zEhX`#N|`F2a`=G_{=EglQDO;@L;T+%Lp}haSo1Sn;BoEvD+sp~hJ&!95-7V6U&^1x z3J9$6Uob^bj(wOy?32HOWIvL1K>THDBul~9`v@PAKHSID$3PTA1pYG#U>3mB$H{b2 zEP1&=W)A7EfZ92SeUK{&e&K%t*~z+>ON)1)j1U762H3cxX1?LQL@u;bXW0vWu~>Il z#h4LYUT6k4{}|l`o_{?O_g=)6_kt5$w)bgu&Ha(PBlF>RV)tTV_1Q=I&rKhj7MPF9 zJ}raWMfbmYe&=~yEA2Te?>Reb%ebrW@4CBdzDsg9%kE}~|Af}tKGT2JpH51St+Hb) zV%>Xx->`UH>N+WRoqW=DZmH{>)OB9&I!_*v&UKe`?)d7R!%u5{cvNz<%8pj@_JaAh z<^NnRZ8;=wIaH|bWJz^>t5^5bRdv7RZp*x1ay81XM$xr5`?Xf7b+6pIch-bo)ADe8 zaf7BpGnY!58C;{8D$lGHgNrUpz7E;f0r-cdhX(irC`1z-o7;}fG(Kb`N0aPmTA{3f zw0zd^i2KO@sej@1k2@cA{(kUl=IdAg-Sj_Sy(sO!B=5iUs58^hB2}<+hPhi9?RxX> zO;z?UM&K&dAD4l8j?13obIcsG?DWldKWvnoEwZy^g)%|?8Snb{&F`Ai9)Ou-uV3^Y zeRSdTfsY3khNX`Ea>sticc&i%!NSl11w zWIra`kBRnUkjP}`PpcY)EV}ljgwS%?#V{YtyexVr!Sg?v>fSB0I1s{Oi>3 zrXDAygXiUg=OIJ9G_c~fVz9JIELI}w`~12d{x}FA2$DVl7ycxYQ%F#hLM{pR^I*mX zB>k)A{&dw`)S3{^!NNrizAmXNp}tjkPFnbD%rzIaOXM>V?6~Jk06DcUC^-H0k}z%_ zV}STqkswe=KyVTpqSyfiPeel5c2TG5poAPOW$Fw)y~`zLBHOP*I75H>B) zd=Zkhg_PNb(c(FhFObV1kH`!lU+MP}w&DMw4o||rGs6F&^0>ps&m#F560)N(3Mr*r zjp6CchVbxl%(UhIF*hvmpvx!>K1z^XX`BWSx^rs+%ma1wa0{6?Bf*VB-ipMAbZHJ`k%4~_}{=Sy7U6Z7=!eo5niVm7!x3N5JMuhvG{yS*=EQu zLzT~vU#5^_nIXRnRVC)m47EYjo*Bw3YR?RHNc=J9ijk(tf0C?9=Qe@#geM?33(w7z z&Nfps*C^5+iT22}XGO=-Ft#s}Z!zzCBrx7OZ@y>`eNu;StWLir`QNYclyLt40#rd> AZ~y=R diff --git a/tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc b/tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc deleted file mode 100644 index 3213b7edf7e304f99fddfd2d0fb51500befce05b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 5817 zcmcgwO>7iL7OtM|`EMH!WEil6*GV?uA&U(dn_#mfYYaHq;4EPCW3sfGskYr_ew^xd zFe5wbNF02~VUbvYBBGrW>j)zGv`eH!@&*nZXRMZNwS*LDFPobY9FR}@Uib9O^mqu_ zBuaJ9>sMW`s$RW%@9S6dkGi@Lfs&c}*Z5ljLjHplr+ACZlh1&;K`5bIk__8VZkV&` zo?(ymbvB@8jE-cNY)jb$ZcTmI28hOaHp(a-2Xkk9f)(L@F!Uxi}PmqRA!PTlp~B~ zm?_$o_cAh7m|1Vvj>~CvMA7s_X+U7OV&zE-FgJ)oaFMV;?ka)hd4LL(2kNE5DH8XZ z4V5Dg;%LqYI1nxo7xqMl<=h~8Mca6ZcQ2}w%y7D6cIZ0I-3G}!C27iK!-ULusE6`Y zpx#k#)HBLaANAi7Xn+Q3hZw8|fxW z^l16rw8uT_v_N0F4Z7SZ=`-WiV`X`LYni9fyPMs;b0tSCOTi9#-BRw_jBlFc>MP|P zp5h(_aQ|owtnFK8?JC83~C6BkU6O3lEBliB#sH9o&ASEqGXq&jQ7Cbnrq?M@&b|`zaI;h*Qqa~|J zmK7y6Eor)ZF{wxwWj!&bXl)0RnS`9w4z-oWLEqGV2jstS&650%6L8=A=u6>lXrZgn z+HJIU+wUsr6Obos-bsCA*PCB(WLvERU(wa~$&3J94=O&rRe{CPGm4Fq;~#>6<*aB8r1+mL-P z>ht=-aihJr^oG3x608k~qo~(1W2L5SH#BFD{X!~QdbQem?Yfk=?_tT-=;0(h3vELN zXzVmvD#kaM0Wjdih?-Q)jf22tn9Q#9so?v1hFxI@=w|3L12aojunOZD6=HWx*2a?R zMboE^$$LAy%n(y#DqXw+K5PcZlqsr?f;TfnXOjoU6wrbUzd5Fe@!+_$8J0CoVfv8| zC*@?^Zwi`{95I8`(Ew(!g2?n2!_f>tBqT0rlPNQhDt=u}e=$HbgxnUpSngoH`r@d{ zcKn1r5Md$>QFxX-s*4(Rt%bUsM%~U;5(=icdG%3r>&^Hz>AIBrOCj1dcl5hhGx_ebcpo^I9~sXpm+}osxL3V=tbXp;YJk*j zS{Ao0iQ8_{zpJ0Ab~HbX!LzXAgaNeqgwcE=|3M-zZYzk?5NTedk7CW&KhBLTZ{NSP zeScy50b~1tLaf_}buY*ImSTN{*ij>PbgutV?Dbq?zGH!382qN=-bY{ceBJXfkUu$8 zh@CcKr=fjwOTM-+U#g~`F zmv3q}m0Th>bn`;d*%ITx6n0%QfQm^&Oy=w2Hqw~K$UOX%CcXQ^BK1$$)?YS}dTwe23D#Yq`g5xXLD^z~L<`a` zNH{8zp6Y}HQ?jargNK9wgKu5YE7AlSa?jieA5uM1?oyffxRat>DX%U|Lz2R~(-L>h z%GPTn7wPcxQj4FEii8KUVr43l4Wu)A3w}=7tDq`-4aw_B@Uz8s05RX*{WA+g+|5QD zTxv^mu|nxdUlMwKasnGNefVqv9#SQ!@> z<~i&wByS@*3`9dwggP7O-mrkq(Lr)_kVa%RsYw@=kqlF$37Ki|ozEVujd(wO`98{Aikd>V8 zMBZ5?T-`p4JZgC9Zy#K{aQ(t^v~?-kI=`h5-D^bm78*K?hK_tg$13kgTH=qfgN2fyv+k^eo1{2}D;*9(U?`tL_MptbR90P>+<2fr$jYP@uy#ABmceWVobdfjLP zSW|J-8801{04smQvCNGvxeCMtV-|}Wvu60pLz{)F0l?LYgIP7MPV?(kk3#*hM%HZaBA{Y`N#&tZL}Tu z78{CN_Oy(sN|LrXXTNx)CNMVPxGy8Yh+zAXyouxhk`ipPeyqBIlDz}(;@)bU>_tfn zEYUd;EQ@Hl-9pQhoCH{e^PFTW3J*wtbq!8o2^kNm+5NSF@`W8S|JSB#V5HBDk#=@i z!pP6*t?~KaF7WeBh2~D9xpUEPA!ffJ_UA>&)Bdk?&by$0b82z(SLeSz|FAVbc%~3L zYsAjZ^}B#|un-wCB12^u{FQ+<3ZvG)#+JFzy{Xdj=o0zzQ4WC zEP;KFqyg{aw-K{9Vxo=>Q0|*{3Y6 z_$>ZqVFww8j(w5~41(xk8P03g$_8*A0TrlP_*&qgZLdS@49Gwj5baAK@KNTt6;eNI zT`MF!Yh5d3!>n~xFkAA(b$>_X{PXWsUo!``T2KBuWyD$^W5d&OYcEFxaW1>WZ_V>t bSA@uH|JCE4ynprmS?gNmiO>iQmYM$tw?bN% From 8ebfc8bd74d90ea4bcd8aa14adc6d5a6a17b79f0 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 15:57:32 +0800 Subject: [PATCH 08/30] fix(graphrag): harden source verification and streaming --- tools/__pycache__/__init__.cpython-311.pyc | Bin 0 -> 193 bytes .../__pycache__/__init__.cpython-311.pyc | Bin 0 -> 208 bytes .../__pycache__/__init__.cpython-311.pyc | Bin 0 -> 219 bytes .../__pycache__/__init__.cpython-311.pyc | Bin 0 -> 367 bytes .../__pycache__/source_loader.cpython-311.pyc | Bin 0 -> 18888 bytes .../graphrag/ingest/sources/source_loader.py | 58 ++++++++++++++++-- .../manifests/2wikimultihopqa-dev.json | 4 +- .../manifests/hotpotqa-distractor.json | 4 +- .../test_source_loader.cpython-311.pyc | Bin 0 -> 5817 bytes 9 files changed, 57 insertions(+), 9 deletions(-) create mode 100644 tools/__pycache__/__init__.cpython-311.pyc create mode 100644 tools/graphrag/__pycache__/__init__.cpython-311.pyc create mode 100644 tools/graphrag/ingest/__pycache__/__init__.cpython-311.pyc create mode 100644 tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc create mode 100644 tools/graphrag/ingest/sources/__pycache__/source_loader.cpython-311.pyc create mode 100644 tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc diff --git a/tools/__pycache__/__init__.cpython-311.pyc b/tools/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4dd5b3811a02d241794564317270bab78979e51a GIT binary patch literal 193 zcmZ3^%ge<81o;))vy_4KV-N=h7@>^MASKfoQW&BbQW%37G?}Vof>I0ei!)2|iz*dL z^7C^t^U@Ux5|gtN(^HG}{4^PFvB$@!&ryk0@&FAkgB{FKt1RJ$TJpjMD0iUonh2WCb_#v2Su L7htG}1t^MASKfoQW&BbQW%37G?}WD+=~(mGJ+i46;cvQ z5)}%HQVSA`5=%1k^At+*^K&xu()Ii_8E>)2$EV~c$H%W^_zW`Ym$80Eer~FMVpUOY zYD#8_L4|&LYGPVWez`tKSFwIN$kd|5bp81F%)HE!_;|g7%3mBdx%nxjIjMF^MASKfoQW&BbQW%37G?}W@{L|8MGV@Xu+=~(m zGJ+i46*BYEQ;SP7^Yavv^K%RG^HTFliuL?78E>)2$EV~c$H%W^_zW`cmz91-er~FM zVpUOYYD#8_L4|&LYGPVWez|^0etu4|emcm+qQrE4h>7~~@tJvb$Lq#k=Q2=wMJy-w$ literal 0 HcmV?d00001 diff --git a/tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc b/tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bac663e53f4ee792335470224238b243090bc9e6 GIT binary patch literal 367 zcmZusy-ve05VoC^LPK`Oh%FL|p<+Y~NNk9KZdoi7yJ?UUM`u@*PCNt>yaD0~c#VuG zQza%gq;8#X2rP)Z(|14m?$h0K7>wQ-|B7ssQupkPTBqgPn zdKd{$_&G^^F_@ruYk6=bYtUkz71v5edOB4_B_qKtH_}E%&mof$NG>FpV{O1!;nqU>HcvGTyNTe)FqOFs z^-~UB8aTG&9QdGvt@;XUVU=Ur&Aw{RZ-D>^H$LJ8B-Y^jpX~`cdnct=~q{ zhS8$2;{M_>d%u0Gq`zd$(eD^@_B+R1{Vtl)Q9TrAe3Rl#A@|1&Mg0Q))T6&tM};h$ z`GMtQ8uH;!J;)nYc*FLw4)WnoJwAq?U!-rS{xVXo2+9>}%k_IWrjPR5V^2;6BjKTt z5N+l7@U>9H$A^Mkh!=cA0V!u=Yuzd77O0I=o<=;hJ+R$yyY8?g(q4E$Hyl4kRXIO-?bng4o0Hz zbZP`j`EG>CV5#v|du;3#x&aGZypgd@Iz@yQ4WZw&;agCim6|DO;`ep;~x zBa!iF5PL85!o3-XI5-*<1h`?(1fwHf-rvK6d-8Wtb06jPSwXEr4rSH(S&$3aQ_~=s2#Qwuu{j5!6P23lc&ZfSZ{!$@Yoqi^u1r z;5k6KVTd1(g(4aprnUq9I07ZJID8p4W(csj4{}M0$)s44(!*0jCS}Y#HDyv}Oodpk zge!iN&D%VB&JOeAP>R+0NpBQ|5u&_eOzi~4o=ODx&}azuhN2fHM#IrqQF=98g5sAI zi;AYeBp+5RgF#q>feBb!ia{6&Zr`dT{guotDUG@7^{&z?E6dlF2lRDlqM2pjDiqeuX!KrCQAE#r`KW+VQ(465->McXUl&vM3($0A&rM-XK zz!|5R8N(3Gncg(Tnf&L|Y@FpAHGMFXky+co{9_ODeB3#_?CgQmuGIZwIy(^tvR%qo7P7wvPI&0Z4F$>RkEjr z@@QqY6z6!8itF=ntbLQ$j>K7bBswZi<&DO*-e~g5xu3IK>3ZeXos+!TTen=^+{N{I zJuF+V!U|ik~3+_Kl?k0N+zxhc8y2 zxvdj?W0OMEH#i=N27!YHRsK9aa3wSt9n6puXbKwy8csic0}^+sFd+Z(z8?UQnWm?8 zaXOa_#dXmf;-lTne04NKM1WwWh(1Lna&pyyA&ln{s|X#}{X5;Q{e2XVoYt>bEMXxW zL5X89#Csu2(Txkd4-=Nq^$8FuK!Ol}|LWSdD$Lc;4Ml$~I64^;uwkFi$GhS8sbwDr zTpz`b1F;q@{cQtIH5Nb#rETBnILI5qfwpXEKfEoR1#pvErD%Iw7W|W|a(2g@ewCq0 zT)*D1(Dn5WX(yb2MO8r<82nwtdIWUAyK@P$#47CI2&{D6cj;Aa-ZvbIz$V}i#G0RDQKt)Uhbn|d0RIdpAhoU^ zu7~afCvNW)>pRu^T&QBMDi~$ zDyKrc!i#gv*`PBEsIf?@)(Zt$uwImYjViUg#KKuF*? zM3%B*&QOe)cRjRk!Z&MSvS8J&vW)p0z2Ypp(=>PP{a4?8b^gNKt5W5DxpKecJRmy{ z%p9FL`edcV^PXESZxsoVffs;CWZ;Uka<+LPisxeAVxL%bTyma}ohQUIeWI=J2^^~o zmWIT)_VT%!_o9-$R<_q-9jBq`O0lZ@pcM|GeVE z3aG^1_6X_RcL-*fZ(QeROv&`FgRPWk1Ry!j$j&n}#-!DW4ZMO4K=(L27d!v!@h^`* z9F{uI%bn*%xT-El&Wp11qG-FAv^#%1_4d?kJlnhT^!>iiFMM`k{s&U)5xMmUR(M|R z@hS1lIjQQLHJBp5lc2pH*-oZYgsTeNjU=VpWzw`Y!- zJASWNayQ8ChLw%g2}Z2lCHrtExGyP%@=?4Vv_ZWogRSB!dBEFqML8}{0DA*?Q=AH9YP zUHiD3o>mb%5LkaMq4DPu+J;cmZ}EYt=Xauo}eZMHqgXn{>$kdLCrgAm>IHP?Ur>WM#Q+r!gm_g3Fu(&M0;u|gV#8e+ zNLe9c<(j3M`{)?FCwvV+6}6If*Rs84$zGEs;E(prN2*PC>DW=^J0lq|d?n{1<}-prQC+bx>Sf(w>xmJ^WpewzgrtpzrW zt_v2#U*Uq${VfY(vK5!dr!-i1`C+IN^o{Jqg&;#qNLZYPm}AY%oWqhhEvXqPEnhCJ zUn;FnT$f6>%B5Rzpp^v7bR|6-MO)zuk)jq;*%s~~&G^x51UR4{&jJLopS}Va>cPJf zMeL-G@1-s0)S}&N+PlIVTC30@h7-81*Bsh%Je2DL5|rJ{yb(r&#E=|9bMDwCXpRQr z}yvX>2v}m zzBNZK7J$5K*aZ!SyF||1Wna!Wc+ynl5}1)2vn9Y>+&gVBV+|&P58fDtxt6~zt2C~bJ%kTpujSrB0?y)pZ~2>M%qgMbGnI&$DC zy(Lzf%E;OUe8H$M*1p#lilDQ|mlnWKasx+NIEub8aN+Rd*TP)L7tA_#NX_R!t2cs6 z0NiQ<81Alvh!GgS8q>EyUNq*;)HW0Z^Aks$b$kP0&R<515USMQ1Q~=g@y(D*TfHU) z&})rhfmp6oX^QC9xb{?~8?8wSi|!%Cp!Px0k6q=$ykbWAlem-=!z4cnUOdH!YB2bL z6f0zar-BcMu7xmnVl)Vbw*y$0VKOorj$8$$K1+Cg_`rmAGjy6M2H@Ce5;B7}J~58I zLB0n{yYd^wCp;$f1l`o^=o$VqTF>+I`oQ+*Q-)lKrRbNjvva7 zA1*ttE;+7BjxpIWHe*PdifyX@gW%u40^Qz8d=IIri^H%S&t$xV{EKF|bSU4`(x@8+tXlygtG4h&ddySYYtk^W#Eo@Xt*-VkJ z8>agYTRXa%sjr(YJq+{pZg;m;_eTreeSrC+)zacJ8EPm&S%^>M`%^>;UsjUZSrnQFJMhWbcPIa&$y`f_i>M)6VI zMSM$c&;CNcOr&4nj?2~p{j|jO&`-0rpWq@y_0-y{==gy_nElLw;ubIy`f(d1(0cwB zmB>8IP2^88RHy3=Q&b6By;(F&!mj}o#87Zdz+t7V+opY^!Lb1@cp$b#OVjf7lpGgB z>V<1iiGYj_42AdWKX3l5dH$-@bVP1Cax-PE!G|$huCIq5!z@!etbnkLn8%k^hR~)x zh5-_Up~S2m$gRNJ0~{u1Xz>k02!h|D7{P=)f&m<8n^xFy(3BJwtf)ND`3dlbt3+Ql zv=SEu;V0&i&~+}15d**meDEewoD_qa2f7omf~$n)gZ4ETh4=*iuc1^915k~q{0|{F zR{2b1y@oHw0n31zZynw{RgzB12ps%@*7)z5SUb)04k}J{pVOs*u`C2>|m<0y9Q@tPTu&x7O zhjkrD9_>bA+k^BH482$(CJdC*yX@GslNbqyC7G`O>F{^wT1uErRTH2`wrb?+B_$HMlE(3PNaFMa-&6V6RfG z2DV&h;}wzzZ&DFMuIe;Rn;`%TZC9L*V;-Y(+?-2*3oBM=OVg5ct%@6&h~}sOV6oFbFloyJD|t0z&(S|mECk;L>H(6F)R9%j z^1b;fZ`3wmB{_;iR#_wQAcknxZB23JI0ZH$!w|DJUZarceJfCyts&Vi%-GZ?QE$5U zgtb8^-y2jEJ*KInaDP(~EZQDh2s^P7Wg#zyqfn3iw%!=iHXH+v66I^5Ilc}78mM?& zh)?ND{AKu!Ri&L_YV=#w2VO%EmP5Wet^?9BcZzhtv}ubEVxb)f_9EDa;2;7b|D(8`qR)x}DLdG)kPIvOp)ij^ zK6Wjj8n_e|Z`Kjh2jHcAOKi%zhCLL_fCSZ$@=ev|dspwjyi~PatlIt)_nq?D^53yY zsvYJ&Al07PvN?8+yMIA+?2;V2WXG;m%2w1(FO(&{RqtEhwI&Y#toUB>tUg&%cGr5_ zI(PV;;ycB&#YvZU?&$4(v-?)6o0qFQma01zb}#Ofs*lOl$3%DG>)Y1d^Vb$%5?gzu z)*iXFXO%LPb^{4UD0%&J*2KAFMfLj~?{@rb-@Sbx1ee289b4H=uhchvcxB!#)$f+; zcZYWE@K6!j<}M;`X!2+Q$3M^8ygk^YlKT+dB|6><8lAus?aE3-UL0 zr-?UIfW+HKTD2n%R7|Mw5y8rL1m@;HJB<}S$16|%yn!kKt08SBemfO_t8A?E&xDIx8JrB0^ zxIPA34|s>(MiDazQ<&T23CIjYp0&jpK{rc7=nKQKS@O#1Q{u8SikvcafS_J5Iu1Dz zNBxa1`#y}s1d(me+Yd&dBa3-GFsFDMME3PTBqDH7#O#_Q-||m5lyTY=@dGc(DSZX! zjU2#BJit8-5C@h!HcR2Lfl)|E8-?OVl}`c2P1^B3C%&0Pj>?wq55zWFwaNrZ-+6s?ymcCW7rfku=Wnc0a_ww6d0OvqIp;dOo&%%AImhd1! zjX4HBq>cZA5HO6*COoYBDSbJ%A<9oizWOM0gY9D zv%2wpFg^&{RSPMG0pOd!+J1_;`oWRO$W?^_hYycx9)xi~EI2rYI02WK>b(I-VmNqo zCdVqY>W7BI&Org;|A6wLvbS%YWwLX#NcBC(5O?KynQlZ zmfTIUy9qo8b~jMZA|G(Gxw?1kvj%ju_-Br;>MDwW9hO(WU-@q3&#LcL!xI-cR}7BL z$%;+zN@>+Uc`K8i%DLd38;PS!o(9p=uu|n0o6d-R7o?^OQq@Jd>Z0hqn5?Wx98Mfb z9J$v&+nwCF2^1SmtseAFI5s5R)&Vs(qcCYI6@T|UdxxRg=zI~xg zsy`&xAA0BXY}Z`-T>Fz1PlfDh2HTbcEI(-g!423!2q(e*U^q$BDXFSguId%Nz0l5V z*Qyc9B1bS$o{cHLQh9a4lrSgEa(UBo`L?C;f%I z?*QE7jtfRsT9NY42vTXYTuM3xrEn1N49D0+|{#3f1?+C%FC%(Hsv^z8cy|&C-__(jtVWXQpTLtS>jbI-Z9M z&YBWr4r%M-Y*}&X8FkmiO{72OX$x0WIBXt1JfY&2X9vAcTO(g;+s`#@$Pe|W_<=WtJlS$8*pg-i#YUY=w<2p z>o0B`V%Hibz_XmY+Re|{(SA)y>oZEaVP`rzqdX$(tdWb0alOg z*{cV>O~~O^T*dnMUx(LfX1S~)FqhflhIO&MQHw*^+v{=n%Dj^FFvV5f zGlRpTx;M6|PS^s1ziO!g6I3QVFBtKuUw$Y|Aj~N`__GqfFF0>8!<4>>)#g_LpDT{S zSBh|D;^DbiQ{htST5`)OM*IW?qq%=}jA}yk zPpIn?1?DO3)0mWT7QiY-Ki&z6v`P80OZ%BbesEsO@aaY+mq+{ar!%n*tRiP4Vdf=KL z`RM3<)6Y+RcuKv`g=FBdJ&kjh&Vtk7QRQA2(7m?Q)`R72rHQ^ zD!-2Z5706G*8sq61;z%9tx$gj4Pudcko9#05}Eb2A%A_X&)02KYPD5lRvWbL$JWuv z4e1t4F97)0K>9ypJzD{!kiKO=gzL=tU!PdG_O&4GJSp!yDQ-Q9nc7Q5ep_MWtC%5p zPa6ja&XtnXBtS@WU`t!8QKb;FN3xoi2@zse}P5+ zBLKycS>XczpCH3uqGIb1mLSG%QY|_sv)_I2=kONI<7nLm8v`+ycVZ?^96yMl3jul& z{6%W?XGYEA{|i1MeRzPW-vfXzOyJ9Nu+f8ybc`$)#hjrFWaW_l3fQHLtpO9Ny@7-5 zSk<$s#TQUUhyV!#cc=Cv4xM=~m|?}L!>Zd59ftGuU{&s-JHXIfjli`RU4?s5$KCr) zS^3@J+rx9AcOrKpV%h2Y%;(0>jPvZLMGuPLcER;~&u=}8mC_5R)f%zox`Cuesol_H_G-#^7g#xH^sjymbM;}w;sw?cOtL4-on*=<1D?~aJym7 zui9or=ibz4EmGrNxpD7|5l7ST(azikjoD_-=QY#6PBW#R84CvFUXUtU<%-r-%5Lta z9}x%-3mEo)S#zKJ)c?Rg|JrA5_uGEg|B!w7(mxvi$4f6u z`_Ie!&);uL)-*^ZzFGb34g7NC`t9o~?OzDPuj=EFJjZ0uu~~MOU2#;*b$(PUIT~a~ z!zyKj`jg&`@0;E=B|PBOkiCAl$fWx4f+zW(IN|a{v9xKV125qyIRn-V>t@7@6|gIs1%uzOj|5(NFjF zpIbk(;@V$`idCI(NwzN8)+O4yAdyV1pGNhA9?|(iLil9rgQ@#Bq=pydh8JM{SW9Pz z-xDNjjcl!fF_?WxPt|OArUyQMnmL#(^UgqIb^a3j<{qFJtJPI|FN@VP_MSM!PS1Ebj3 zA^G;ozP+OF;)3w|sozd5#-xL1<%4G-Lp*0;Hvo+mC7T(a5lL!5suh$n!s?LJuu*)1n7fcmm*N=1j6Li|os*EUO2iPypC0da(a$fmf#o>#Y`; zJbH-Q(}L}k!MUm{)`$2^SQ^n%ZC)z#84w`;*8%=-=j%(u>)DcVb4)8QM1;Ukj&-8xSR=Kn7VZksm% literal 0 HcmV?d00001 diff --git a/tools/graphrag/ingest/sources/source_loader.py b/tools/graphrag/ingest/sources/source_loader.py index 6406a965d..e56ec240e 100644 --- a/tools/graphrag/ingest/sources/source_loader.py +++ b/tools/graphrag/ingest/sources/source_loader.py @@ -36,6 +36,10 @@ def from_file(cls, path: Path) -> "SourceManifest": raw = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as exc: raise SourceError("cannot read manifest %s: %s" % (path, exc)) from exc + if not isinstance(raw, dict): + raise SourceError("manifest %s must contain a JSON object" % path) + raw = {key: os.path.expandvars(value) if isinstance(value, str) else value + for key, value in raw.items()} required = ("dataset", "dataset_release", "split", "sha256", "parser_version") missing = [key for key in required if not str(raw.get(key, "")).strip()] if missing: @@ -94,9 +98,15 @@ def _resolve_source(self, manifest: SourceManifest) -> Path: ) target = Path(os.path.expandvars(manifest.cache_path or Path(manifest.source_uri).name)).expanduser() target.parent.mkdir(parents=True, exist_ok=True) + temporary = target.with_name(target.name + ".part") try: - urllib.request.urlretrieve(manifest.source_uri, target) + urllib.request.urlretrieve(manifest.source_uri, temporary) + temporary.replace(target) except OSError as exc: + try: + temporary.unlink() + except OSError: + pass raise SourceError("download failed for %s/%s: %s" % (manifest.dataset, manifest.split, exc)) from exc return target @@ -125,10 +135,7 @@ def _records(path: Path) -> Iterable[Any]: first = stream.read(1) stream.seek(0) if first == "[": - values = json.load(stream) - if not isinstance(values, list): - raise SourceError("JSON source %s must contain an array" % path) - yield from values + yield from SourceLoader._array_records(stream, path) else: for line_number, line in enumerate(stream, start=1): if line.strip(): @@ -142,6 +149,47 @@ def _records(path: Path) -> Iterable[Any]: except (OSError, json.JSONDecodeError) as exc: raise SourceError("cannot parse source %s: %s" % (path, exc)) from exc + @staticmethod + def _array_records(stream: Any, path: Path) -> Iterator[Any]: + decoder = json.JSONDecoder() + buffer = "" + first = True + while True: + chunk = stream.read(64 * 1024) + if chunk: + buffer += chunk + end = not chunk + while True: + buffer = buffer.lstrip() + if first: + if not buffer: + break + if buffer[0] != "[": + raise SourceError("JSON source %s must contain an array" % path) + buffer = buffer[1:] + first = False + buffer = buffer.lstrip() + if not buffer: + break + if buffer[0] == "]": + return + try: + value, consumed = decoder.raw_decode(buffer) + except json.JSONDecodeError: + if end: + raise SourceError("truncated JSON array in %s" % path) + break + yield value + buffer = buffer[consumed:].lstrip() + if buffer.startswith(","): + buffer = buffer[1:] + elif buffer.startswith("]"): + return + elif buffer and end: + raise SourceError("invalid JSON array separator in %s" % path) + if end: + raise SourceError("unterminated JSON array in %s" % path) + @staticmethod def _canonical_record(manifest: SourceManifest, number: int, raw: Any) -> Dict[str, Any]: if not isinstance(raw, dict): diff --git a/tools/graphrag/manifests/2wikimultihopqa-dev.json b/tools/graphrag/manifests/2wikimultihopqa-dev.json index 180e49332..03edb9c5f 100644 --- a/tools/graphrag/manifests/2wikimultihopqa-dev.json +++ b/tools/graphrag/manifests/2wikimultihopqa-dev.json @@ -2,8 +2,8 @@ "dataset": "2wikimultihopqa", "dataset_release": "official", "split": "dev", - "source_uri": null, + "source_uri": "https://www.dropbox.com/s/npidmtadreo6df2/data.zip?dl=1", "cache_path": "${GRAPH_RAG_CACHE}/2wikimultihopqa-dev.json", - "sha256": "REPLACE_WITH_OFFICIAL_SHA256", + "sha256": "${TWIKIMULTIHOPQA_DEV_SHA256}", "parser_version": "2wikimultihopqa-parser-v1" } diff --git a/tools/graphrag/manifests/hotpotqa-distractor.json b/tools/graphrag/manifests/hotpotqa-distractor.json index 415cbba99..6478266fc 100644 --- a/tools/graphrag/manifests/hotpotqa-distractor.json +++ b/tools/graphrag/manifests/hotpotqa-distractor.json @@ -2,8 +2,8 @@ "dataset": "hotpotqa", "dataset_release": "official-distractor", "split": "distractor", - "source_uri": null, + "source_uri": "http://curtis.ml.cmu.edu/datasets/hotpot/hotpot_dev_distractor_v1.json", "cache_path": "${GRAPH_RAG_CACHE}/hotpotqa-distractor.json", - "sha256": "REPLACE_WITH_OFFICIAL_SHA256", + "sha256": "${HOTPOTQA_DISTRACTOR_SHA256}", "parser_version": "hotpotqa-parser-v1" } diff --git a/tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc b/tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3213b7edf7e304f99fddfd2d0fb51500befce05b GIT binary patch literal 5817 zcmcgwO>7iL7OtM|`EMH!WEil6*GV?uA&U(dn_#mfYYaHq;4EPCW3sfGskYr_ew^xd zFe5wbNF02~VUbvYBBGrW>j)zGv`eH!@&*nZXRMZNwS*LDFPobY9FR}@Uib9O^mqu_ zBuaJ9>sMW`s$RW%@9S6dkGi@Lfs&c}*Z5ljLjHplr+ACZlh1&;K`5bIk__8VZkV&` zo?(ymbvB@8jE-cNY)jb$ZcTmI28hOaHp(a-2Xkk9f)(L@F!Uxi}PmqRA!PTlp~B~ zm?_$o_cAh7m|1Vvj>~CvMA7s_X+U7OV&zE-FgJ)oaFMV;?ka)hd4LL(2kNE5DH8XZ z4V5Dg;%LqYI1nxo7xqMl<=h~8Mca6ZcQ2}w%y7D6cIZ0I-3G}!C27iK!-ULusE6`Y zpx#k#)HBLaANAi7Xn+Q3hZw8|fxW z^l16rw8uT_v_N0F4Z7SZ=`-WiV`X`LYni9fyPMs;b0tSCOTi9#-BRw_jBlFc>MP|P zp5h(_aQ|owtnFK8?JC83~C6BkU6O3lEBliB#sH9o&ASEqGXq&jQ7Cbnrq?M@&b|`zaI;h*Qqa~|J zmK7y6Eor)ZF{wxwWj!&bXl)0RnS`9w4z-oWLEqGV2jstS&650%6L8=A=u6>lXrZgn z+HJIU+wUsr6Obos-bsCA*PCB(WLvERU(wa~$&3J94=O&rRe{CPGm4Fq;~#>6<*aB8r1+mL-P z>ht=-aihJr^oG3x608k~qo~(1W2L5SH#BFD{X!~QdbQem?Yfk=?_tT-=;0(h3vELN zXzVmvD#kaM0Wjdih?-Q)jf22tn9Q#9so?v1hFxI@=w|3L12aojunOZD6=HWx*2a?R zMboE^$$LAy%n(y#DqXw+K5PcZlqsr?f;TfnXOjoU6wrbUzd5Fe@!+_$8J0CoVfv8| zC*@?^Zwi`{95I8`(Ew(!g2?n2!_f>tBqT0rlPNQhDt=u}e=$HbgxnUpSngoH`r@d{ zcKn1r5Md$>QFxX-s*4(Rt%bUsM%~U;5(=icdG%3r>&^Hz>AIBrOCj1dcl5hhGx_ebcpo^I9~sXpm+}osxL3V=tbXp;YJk*j zS{Ao0iQ8_{zpJ0Ab~HbX!LzXAgaNeqgwcE=|3M-zZYzk?5NTedk7CW&KhBLTZ{NSP zeScy50b~1tLaf_}buY*ImSTN{*ij>PbgutV?Dbq?zGH!382qN=-bY{ceBJXfkUu$8 zh@CcKr=fjwOTM-+U#g~`F zmv3q}m0Th>bn`;d*%ITx6n0%QfQm^&Oy=w2Hqw~K$UOX%CcXQ^BK1$$)?YS}dTwe23D#Yq`g5xXLD^z~L<`a` zNH{8zp6Y}HQ?jargNK9wgKu5YE7AlSa?jieA5uM1?oyffxRat>DX%U|Lz2R~(-L>h z%GPTn7wPcxQj4FEii8KUVr43l4Wu)A3w}=7tDq`-4aw_B@Uz8s05RX*{WA+g+|5QD zTxv^mu|nxdUlMwKasnGNefVqv9#SQ!@> z<~i&wByS@*3`9dwggP7O-mrkq(Lr)_kVa%RsYw@=kqlF$37Ki|ozEVujd(wO`98{Aikd>V8 zMBZ5?T-`p4JZgC9Zy#K{aQ(t^v~?-kI=`h5-D^bm78*K?hK_tg$13kgTH=qfgN2fyv+k^eo1{2}D;*9(U?`tL_MptbR90P>+<2fr$jYP@uy#ABmceWVobdfjLP zSW|J-8801{04smQvCNGvxeCMtV-|}Wvu60pLz{)F0l?LYgIP7MPV?(kk3#*hM%HZaBA{Y`N#&tZL}Tu z78{CN_Oy(sN|LrXXTNx)CNMVPxGy8Yh+zAXyouxhk`ipPeyqBIlDz}(;@)bU>_tfn zEYUd;EQ@Hl-9pQhoCH{e^PFTW3J*wtbq!8o2^kNm+5NSF@`W8S|JSB#V5HBDk#=@i z!pP6*t?~KaF7WeBh2~D9xpUEPA!ffJ_UA>&)Bdk?&by$0b82z(SLeSz|FAVbc%~3L zYsAjZ^}B#|un-wCB12^u{FQ+<3ZvG)#+JFzy{Xdj=o0zzQ4WC zEP;KFqyg{aw-K{9Vxo=>Q0|*{3Y6 z_$>ZqVFww8j(w5~41(xk8P03g$_8*A0TrlP_*&qgZLdS@49Gwj5baAK@KNTt6;eNI zT`MF!Yh5d3!>n~xFkAA(b$>_X{PXWsUo!``T2KBuWyD$^W5d&OYcEFxaW1>WZ_V>t bSA@uH|JCE4ynprmS?gNmiO>iQmYM$tw?bN% literal 0 HcmV?d00001 From a4350afb46a92f7841eb3350d9ea8d067f392446 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 15:58:15 +0800 Subject: [PATCH 09/30] chore(graphrag): remove generated Python caches --- tools/__pycache__/__init__.cpython-311.pyc | Bin 193 -> 0 bytes .../__pycache__/__init__.cpython-311.pyc | Bin 208 -> 0 bytes .../ingest/__pycache__/__init__.cpython-311.pyc | Bin 219 -> 0 bytes .../__pycache__/__init__.cpython-311.pyc | Bin 367 -> 0 bytes .../__pycache__/source_loader.cpython-311.pyc | Bin 18888 -> 0 bytes .../test_source_loader.cpython-311.pyc | Bin 5817 -> 0 bytes 6 files changed, 0 insertions(+), 0 deletions(-) delete mode 100644 tools/__pycache__/__init__.cpython-311.pyc delete mode 100644 tools/graphrag/__pycache__/__init__.cpython-311.pyc delete mode 100644 tools/graphrag/ingest/__pycache__/__init__.cpython-311.pyc delete mode 100644 tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc delete mode 100644 tools/graphrag/ingest/sources/__pycache__/source_loader.cpython-311.pyc delete mode 100644 tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc diff --git a/tools/__pycache__/__init__.cpython-311.pyc b/tools/__pycache__/__init__.cpython-311.pyc deleted file mode 100644 index 4dd5b3811a02d241794564317270bab78979e51a..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 193 zcmZ3^%ge<81o;))vy_4KV-N=h7@>^MASKfoQW&BbQW%37G?}Vof>I0ei!)2|iz*dL z^7C^t^U@Ux5|gtN(^HG}{4^PFvB$@!&ryk0@&FAkgB{FKt1RJ$TJpjMD0iUonh2WCb_#v2Su L7htG}1t^MASKfoQW&BbQW%37G?}WD+=~(mGJ+i46;cvQ z5)}%HQVSA`5=%1k^At+*^K&xu()Ii_8E>)2$EV~c$H%W^_zW`Ym$80Eer~FMVpUOY zYD#8_L4|&LYGPVWez`tKSFwIN$kd|5bp81F%)HE!_;|g7%3mBdx%nxjIjMF^MASKfoQW&BbQW%37G?}W@{L|8MGV@Xu+=~(m zGJ+i46*BYEQ;SP7^Yavv^K%RG^HTFliuL?78E>)2$EV~c$H%W^_zW`cmz91-er~FM zVpUOYYD#8_L4|&LYGPVWez|^0etu4|emcm+qQrE4h>7~~@tJvb$Lq#k=Q2=wMJy-w$ diff --git a/tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc b/tools/graphrag/ingest/sources/__pycache__/__init__.cpython-311.pyc deleted file mode 100644 index bac663e53f4ee792335470224238b243090bc9e6..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 367 zcmZusy-ve05VoC^LPK`Oh%FL|p<+Y~NNk9KZdoi7yJ?UUM`u@*PCNt>yaD0~c#VuG zQza%gq;8#X2rP)Z(|14m?$h0K7>wQ-|B7ssQupkPTBqgPn zdKd{$_&G^^F_@ruYk6=bYtUkz71v5edOB4_B_qKtH_}E%&mof$NG>FpV{O1!;nqU>HcvGTyNTe)FqOFs z^-~UB8aTG&9QdGvt@;XUVU=Ur&Aw{RZ-D>^H$LJ8B-Y^jpX~`cdnct=~q{ zhS8$2;{M_>d%u0Gq`zd$(eD^@_B+R1{Vtl)Q9TrAe3Rl#A@|1&Mg0Q))T6&tM};h$ z`GMtQ8uH;!J;)nYc*FLw4)WnoJwAq?U!-rS{xVXo2+9>}%k_IWrjPR5V^2;6BjKTt z5N+l7@U>9H$A^Mkh!=cA0V!u=Yuzd77O0I=o<=;hJ+R$yyY8?g(q4E$Hyl4kRXIO-?bng4o0Hz zbZP`j`EG>CV5#v|du;3#x&aGZypgd@Iz@yQ4WZw&;agCim6|DO;`ep;~x zBa!iF5PL85!o3-XI5-*<1h`?(1fwHf-rvK6d-8Wtb06jPSwXEr4rSH(S&$3aQ_~=s2#Qwuu{j5!6P23lc&ZfSZ{!$@Yoqi^u1r z;5k6KVTd1(g(4aprnUq9I07ZJID8p4W(csj4{}M0$)s44(!*0jCS}Y#HDyv}Oodpk zge!iN&D%VB&JOeAP>R+0NpBQ|5u&_eOzi~4o=ODx&}azuhN2fHM#IrqQF=98g5sAI zi;AYeBp+5RgF#q>feBb!ia{6&Zr`dT{guotDUG@7^{&z?E6dlF2lRDlqM2pjDiqeuX!KrCQAE#r`KW+VQ(465->McXUl&vM3($0A&rM-XK zz!|5R8N(3Gncg(Tnf&L|Y@FpAHGMFXky+co{9_ODeB3#_?CgQmuGIZwIy(^tvR%qo7P7wvPI&0Z4F$>RkEjr z@@QqY6z6!8itF=ntbLQ$j>K7bBswZi<&DO*-e~g5xu3IK>3ZeXos+!TTen=^+{N{I zJuF+V!U|ik~3+_Kl?k0N+zxhc8y2 zxvdj?W0OMEH#i=N27!YHRsK9aa3wSt9n6puXbKwy8csic0}^+sFd+Z(z8?UQnWm?8 zaXOa_#dXmf;-lTne04NKM1WwWh(1Lna&pyyA&ln{s|X#}{X5;Q{e2XVoYt>bEMXxW zL5X89#Csu2(Txkd4-=Nq^$8FuK!Ol}|LWSdD$Lc;4Ml$~I64^;uwkFi$GhS8sbwDr zTpz`b1F;q@{cQtIH5Nb#rETBnILI5qfwpXEKfEoR1#pvErD%Iw7W|W|a(2g@ewCq0 zT)*D1(Dn5WX(yb2MO8r<82nwtdIWUAyK@P$#47CI2&{D6cj;Aa-ZvbIz$V}i#G0RDQKt)Uhbn|d0RIdpAhoU^ zu7~afCvNW)>pRu^T&QBMDi~$ zDyKrc!i#gv*`PBEsIf?@)(Zt$uwImYjViUg#KKuF*? zM3%B*&QOe)cRjRk!Z&MSvS8J&vW)p0z2Ypp(=>PP{a4?8b^gNKt5W5DxpKecJRmy{ z%p9FL`edcV^PXESZxsoVffs;CWZ;Uka<+LPisxeAVxL%bTyma}ohQUIeWI=J2^^~o zmWIT)_VT%!_o9-$R<_q-9jBq`O0lZ@pcM|GeVE z3aG^1_6X_RcL-*fZ(QeROv&`FgRPWk1Ry!j$j&n}#-!DW4ZMO4K=(L27d!v!@h^`* z9F{uI%bn*%xT-El&Wp11qG-FAv^#%1_4d?kJlnhT^!>iiFMM`k{s&U)5xMmUR(M|R z@hS1lIjQQLHJBp5lc2pH*-oZYgsTeNjU=VpWzw`Y!- zJASWNayQ8ChLw%g2}Z2lCHrtExGyP%@=?4Vv_ZWogRSB!dBEFqML8}{0DA*?Q=AH9YP zUHiD3o>mb%5LkaMq4DPu+J;cmZ}EYt=Xauo}eZMHqgXn{>$kdLCrgAm>IHP?Ur>WM#Q+r!gm_g3Fu(&M0;u|gV#8e+ zNLe9c<(j3M`{)?FCwvV+6}6If*Rs84$zGEs;E(prN2*PC>DW=^J0lq|d?n{1<}-prQC+bx>Sf(w>xmJ^WpewzgrtpzrW zt_v2#U*Uq${VfY(vK5!dr!-i1`C+IN^o{Jqg&;#qNLZYPm}AY%oWqhhEvXqPEnhCJ zUn;FnT$f6>%B5Rzpp^v7bR|6-MO)zuk)jq;*%s~~&G^x51UR4{&jJLopS}Va>cPJf zMeL-G@1-s0)S}&N+PlIVTC30@h7-81*Bsh%Je2DL5|rJ{yb(r&#E=|9bMDwCXpRQr z}yvX>2v}m zzBNZK7J$5K*aZ!SyF||1Wna!Wc+ynl5}1)2vn9Y>+&gVBV+|&P58fDtxt6~zt2C~bJ%kTpujSrB0?y)pZ~2>M%qgMbGnI&$DC zy(Lzf%E;OUe8H$M*1p#lilDQ|mlnWKasx+NIEub8aN+Rd*TP)L7tA_#NX_R!t2cs6 z0NiQ<81Alvh!GgS8q>EyUNq*;)HW0Z^Aks$b$kP0&R<515USMQ1Q~=g@y(D*TfHU) z&})rhfmp6oX^QC9xb{?~8?8wSi|!%Cp!Px0k6q=$ykbWAlem-=!z4cnUOdH!YB2bL z6f0zar-BcMu7xmnVl)Vbw*y$0VKOorj$8$$K1+Cg_`rmAGjy6M2H@Ce5;B7}J~58I zLB0n{yYd^wCp;$f1l`o^=o$VqTF>+I`oQ+*Q-)lKrRbNjvva7 zA1*ttE;+7BjxpIWHe*PdifyX@gW%u40^Qz8d=IIri^H%S&t$xV{EKF|bSU4`(x@8+tXlygtG4h&ddySYYtk^W#Eo@Xt*-VkJ z8>agYTRXa%sjr(YJq+{pZg;m;_eTreeSrC+)zacJ8EPm&S%^>M`%^>;UsjUZSrnQFJMhWbcPIa&$y`f_i>M)6VI zMSM$c&;CNcOr&4nj?2~p{j|jO&`-0rpWq@y_0-y{==gy_nElLw;ubIy`f(d1(0cwB zmB>8IP2^88RHy3=Q&b6By;(F&!mj}o#87Zdz+t7V+opY^!Lb1@cp$b#OVjf7lpGgB z>V<1iiGYj_42AdWKX3l5dH$-@bVP1Cax-PE!G|$huCIq5!z@!etbnkLn8%k^hR~)x zh5-_Up~S2m$gRNJ0~{u1Xz>k02!h|D7{P=)f&m<8n^xFy(3BJwtf)ND`3dlbt3+Ql zv=SEu;V0&i&~+}15d**meDEewoD_qa2f7omf~$n)gZ4ETh4=*iuc1^915k~q{0|{F zR{2b1y@oHw0n31zZynw{RgzB12ps%@*7)z5SUb)04k}J{pVOs*u`C2>|m<0y9Q@tPTu&x7O zhjkrD9_>bA+k^BH482$(CJdC*yX@GslNbqyC7G`O>F{^wT1uErRTH2`wrb?+B_$HMlE(3PNaFMa-&6V6RfG z2DV&h;}wzzZ&DFMuIe;Rn;`%TZC9L*V;-Y(+?-2*3oBM=OVg5ct%@6&h~}sOV6oFbFloyJD|t0z&(S|mECk;L>H(6F)R9%j z^1b;fZ`3wmB{_;iR#_wQAcknxZB23JI0ZH$!w|DJUZarceJfCyts&Vi%-GZ?QE$5U zgtb8^-y2jEJ*KInaDP(~EZQDh2s^P7Wg#zyqfn3iw%!=iHXH+v66I^5Ilc}78mM?& zh)?ND{AKu!Ri&L_YV=#w2VO%EmP5Wet^?9BcZzhtv}ubEVxb)f_9EDa;2;7b|D(8`qR)x}DLdG)kPIvOp)ij^ zK6Wjj8n_e|Z`Kjh2jHcAOKi%zhCLL_fCSZ$@=ev|dspwjyi~PatlIt)_nq?D^53yY zsvYJ&Al07PvN?8+yMIA+?2;V2WXG;m%2w1(FO(&{RqtEhwI&Y#toUB>tUg&%cGr5_ zI(PV;;ycB&#YvZU?&$4(v-?)6o0qFQma01zb}#Ofs*lOl$3%DG>)Y1d^Vb$%5?gzu z)*iXFXO%LPb^{4UD0%&J*2KAFMfLj~?{@rb-@Sbx1ee289b4H=uhchvcxB!#)$f+; zcZYWE@K6!j<}M;`X!2+Q$3M^8ygk^YlKT+dB|6><8lAus?aE3-UL0 zr-?UIfW+HKTD2n%R7|Mw5y8rL1m@;HJB<}S$16|%yn!kKt08SBemfO_t8A?E&xDIx8JrB0^ zxIPA34|s>(MiDazQ<&T23CIjYp0&jpK{rc7=nKQKS@O#1Q{u8SikvcafS_J5Iu1Dz zNBxa1`#y}s1d(me+Yd&dBa3-GFsFDMME3PTBqDH7#O#_Q-||m5lyTY=@dGc(DSZX! zjU2#BJit8-5C@h!HcR2Lfl)|E8-?OVl}`c2P1^B3C%&0Pj>?wq55zWFwaNrZ-+6s?ymcCW7rfku=Wnc0a_ww6d0OvqIp;dOo&%%AImhd1! zjX4HBq>cZA5HO6*COoYBDSbJ%A<9oizWOM0gY9D zv%2wpFg^&{RSPMG0pOd!+J1_;`oWRO$W?^_hYycx9)xi~EI2rYI02WK>b(I-VmNqo zCdVqY>W7BI&Org;|A6wLvbS%YWwLX#NcBC(5O?KynQlZ zmfTIUy9qo8b~jMZA|G(Gxw?1kvj%ju_-Br;>MDwW9hO(WU-@q3&#LcL!xI-cR}7BL z$%;+zN@>+Uc`K8i%DLd38;PS!o(9p=uu|n0o6d-R7o?^OQq@Jd>Z0hqn5?Wx98Mfb z9J$v&+nwCF2^1SmtseAFI5s5R)&Vs(qcCYI6@T|UdxxRg=zI~xg zsy`&xAA0BXY}Z`-T>Fz1PlfDh2HTbcEI(-g!423!2q(e*U^q$BDXFSguId%Nz0l5V z*Qyc9B1bS$o{cHLQh9a4lrSgEa(UBo`L?C;f%I z?*QE7jtfRsT9NY42vTXYTuM3xrEn1N49D0+|{#3f1?+C%FC%(Hsv^z8cy|&C-__(jtVWXQpTLtS>jbI-Z9M z&YBWr4r%M-Y*}&X8FkmiO{72OX$x0WIBXt1JfY&2X9vAcTO(g;+s`#@$Pe|W_<=WtJlS$8*pg-i#YUY=w<2p z>o0B`V%Hibz_XmY+Re|{(SA)y>oZEaVP`rzqdX$(tdWb0alOg z*{cV>O~~O^T*dnMUx(LfX1S~)FqhflhIO&MQHw*^+v{=n%Dj^FFvV5f zGlRpTx;M6|PS^s1ziO!g6I3QVFBtKuUw$Y|Aj~N`__GqfFF0>8!<4>>)#g_LpDT{S zSBh|D;^DbiQ{htST5`)OM*IW?qq%=}jA}yk zPpIn?1?DO3)0mWT7QiY-Ki&z6v`P80OZ%BbesEsO@aaY+mq+{ar!%n*tRiP4Vdf=KL z`RM3<)6Y+RcuKv`g=FBdJ&kjh&Vtk7QRQA2(7m?Q)`R72rHQ^ zD!-2Z5706G*8sq61;z%9tx$gj4Pudcko9#05}Eb2A%A_X&)02KYPD5lRvWbL$JWuv z4e1t4F97)0K>9ypJzD{!kiKO=gzL=tU!PdG_O&4GJSp!yDQ-Q9nc7Q5ep_MWtC%5p zPa6ja&XtnXBtS@WU`t!8QKb;FN3xoi2@zse}P5+ zBLKycS>XczpCH3uqGIb1mLSG%QY|_sv)_I2=kONI<7nLm8v`+ycVZ?^96yMl3jul& z{6%W?XGYEA{|i1MeRzPW-vfXzOyJ9Nu+f8ybc`$)#hjrFWaW_l3fQHLtpO9Ny@7-5 zSk<$s#TQUUhyV!#cc=Cv4xM=~m|?}L!>Zd59ftGuU{&s-JHXIfjli`RU4?s5$KCr) zS^3@J+rx9AcOrKpV%h2Y%;(0>jPvZLMGuPLcER;~&u=}8mC_5R)f%zox`Cuesol_H_G-#^7g#xH^sjymbM;}w;sw?cOtL4-on*=<1D?~aJym7 zui9or=ibz4EmGrNxpD7|5l7ST(azikjoD_-=QY#6PBW#R84CvFUXUtU<%-r-%5Lta z9}x%-3mEo)S#zKJ)c?Rg|JrA5_uGEg|B!w7(mxvi$4f6u z`_Ie!&);uL)-*^ZzFGb34g7NC`t9o~?OzDPuj=EFJjZ0uu~~MOU2#;*b$(PUIT~a~ z!zyKj`jg&`@0;E=B|PBOkiCAl$fWx4f+zW(IN|a{v9xKV125qyIRn-V>t@7@6|gIs1%uzOj|5(NFjF zpIbk(;@V$`idCI(NwzN8)+O4yAdyV1pGNhA9?|(iLil9rgQ@#Bq=pydh8JM{SW9Pz z-xDNjjcl!fF_?WxPt|OArUyQMnmL#(^UgqIb^a3j<{qFJtJPI|FN@VP_MSM!PS1Ebj3 zA^G;ozP+OF;)3w|sozd5#-xL1<%4G-Lp*0;Hvo+mC7T(a5lL!5suh$n!s?LJuu*)1n7fcmm*N=1j6Li|os*EUO2iPypC0da(a$fmf#o>#Y`; zJbH-Q(}L}k!MUm{)`$2^SQ^n%ZC)z#84w`;*8%=-=j%(u>)DcVb4)8QM1;Ukj&-8xSR=Kn7VZksm% diff --git a/tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc b/tools/graphrag/tests/__pycache__/test_source_loader.cpython-311.pyc deleted file mode 100644 index 3213b7edf7e304f99fddfd2d0fb51500befce05b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 5817 zcmcgwO>7iL7OtM|`EMH!WEil6*GV?uA&U(dn_#mfYYaHq;4EPCW3sfGskYr_ew^xd zFe5wbNF02~VUbvYBBGrW>j)zGv`eH!@&*nZXRMZNwS*LDFPobY9FR}@Uib9O^mqu_ zBuaJ9>sMW`s$RW%@9S6dkGi@Lfs&c}*Z5ljLjHplr+ACZlh1&;K`5bIk__8VZkV&` zo?(ymbvB@8jE-cNY)jb$ZcTmI28hOaHp(a-2Xkk9f)(L@F!Uxi}PmqRA!PTlp~B~ zm?_$o_cAh7m|1Vvj>~CvMA7s_X+U7OV&zE-FgJ)oaFMV;?ka)hd4LL(2kNE5DH8XZ z4V5Dg;%LqYI1nxo7xqMl<=h~8Mca6ZcQ2}w%y7D6cIZ0I-3G}!C27iK!-ULusE6`Y zpx#k#)HBLaANAi7Xn+Q3hZw8|fxW z^l16rw8uT_v_N0F4Z7SZ=`-WiV`X`LYni9fyPMs;b0tSCOTi9#-BRw_jBlFc>MP|P zp5h(_aQ|owtnFK8?JC83~C6BkU6O3lEBliB#sH9o&ASEqGXq&jQ7Cbnrq?M@&b|`zaI;h*Qqa~|J zmK7y6Eor)ZF{wxwWj!&bXl)0RnS`9w4z-oWLEqGV2jstS&650%6L8=A=u6>lXrZgn z+HJIU+wUsr6Obos-bsCA*PCB(WLvERU(wa~$&3J94=O&rRe{CPGm4Fq;~#>6<*aB8r1+mL-P z>ht=-aihJr^oG3x608k~qo~(1W2L5SH#BFD{X!~QdbQem?Yfk=?_tT-=;0(h3vELN zXzVmvD#kaM0Wjdih?-Q)jf22tn9Q#9so?v1hFxI@=w|3L12aojunOZD6=HWx*2a?R zMboE^$$LAy%n(y#DqXw+K5PcZlqsr?f;TfnXOjoU6wrbUzd5Fe@!+_$8J0CoVfv8| zC*@?^Zwi`{95I8`(Ew(!g2?n2!_f>tBqT0rlPNQhDt=u}e=$HbgxnUpSngoH`r@d{ zcKn1r5Md$>QFxX-s*4(Rt%bUsM%~U;5(=icdG%3r>&^Hz>AIBrOCj1dcl5hhGx_ebcpo^I9~sXpm+}osxL3V=tbXp;YJk*j zS{Ao0iQ8_{zpJ0Ab~HbX!LzXAgaNeqgwcE=|3M-zZYzk?5NTedk7CW&KhBLTZ{NSP zeScy50b~1tLaf_}buY*ImSTN{*ij>PbgutV?Dbq?zGH!382qN=-bY{ceBJXfkUu$8 zh@CcKr=fjwOTM-+U#g~`F zmv3q}m0Th>bn`;d*%ITx6n0%QfQm^&Oy=w2Hqw~K$UOX%CcXQ^BK1$$)?YS}dTwe23D#Yq`g5xXLD^z~L<`a` zNH{8zp6Y}HQ?jargNK9wgKu5YE7AlSa?jieA5uM1?oyffxRat>DX%U|Lz2R~(-L>h z%GPTn7wPcxQj4FEii8KUVr43l4Wu)A3w}=7tDq`-4aw_B@Uz8s05RX*{WA+g+|5QD zTxv^mu|nxdUlMwKasnGNefVqv9#SQ!@> z<~i&wByS@*3`9dwggP7O-mrkq(Lr)_kVa%RsYw@=kqlF$37Ki|ozEVujd(wO`98{Aikd>V8 zMBZ5?T-`p4JZgC9Zy#K{aQ(t^v~?-kI=`h5-D^bm78*K?hK_tg$13kgTH=qfgN2fyv+k^eo1{2}D;*9(U?`tL_MptbR90P>+<2fr$jYP@uy#ABmceWVobdfjLP zSW|J-8801{04smQvCNGvxeCMtV-|}Wvu60pLz{)F0l?LYgIP7MPV?(kk3#*hM%HZaBA{Y`N#&tZL}Tu z78{CN_Oy(sN|LrXXTNx)CNMVPxGy8Yh+zAXyouxhk`ipPeyqBIlDz}(;@)bU>_tfn zEYUd;EQ@Hl-9pQhoCH{e^PFTW3J*wtbq!8o2^kNm+5NSF@`W8S|JSB#V5HBDk#=@i z!pP6*t?~KaF7WeBh2~D9xpUEPA!ffJ_UA>&)Bdk?&by$0b82z(SLeSz|FAVbc%~3L zYsAjZ^}B#|un-wCB12^u{FQ+<3ZvG)#+JFzy{Xdj=o0zzQ4WC zEP;KFqyg{aw-K{9Vxo=>Q0|*{3Y6 z_$>ZqVFww8j(w5~41(xk8P03g$_8*A0TrlP_*&qgZLdS@49Gwj5baAK@KNTt6;eNI zT`MF!Yh5d3!>n~xFkAA(b$>_X{PXWsUo!``T2KBuWyD$^W5d&OYcEFxaW1>WZ_V>t bSA@uH|JCE4ynprmS?gNmiO>iQmYM$tw?bN% From 87f1889432536868c9fdc3a8abb89981e5654e71 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 16:02:51 +0800 Subject: [PATCH 10/30] fix(graphrag): verify staged downloads and strict JSON arrays --- .../graphrag/ingest/sources/source_loader.py | 55 ++++++++++++++----- tools/graphrag/tests/test_source_loader.py | 21 +++++++ 2 files changed, 62 insertions(+), 14 deletions(-) diff --git a/tools/graphrag/ingest/sources/source_loader.py b/tools/graphrag/ingest/sources/source_loader.py index e56ec240e..485f94684 100644 --- a/tools/graphrag/ingest/sources/source_loader.py +++ b/tools/graphrag/ingest/sources/source_loader.py @@ -10,6 +10,7 @@ import hashlib import json import os +import tempfile import urllib.request from dataclasses import dataclass from pathlib import Path @@ -49,6 +50,8 @@ def from_file(cls, path: Path) -> "SourceManifest": checksum = str(raw["sha256"]).lower() if len(checksum) != 64 or any(char not in "0123456789abcdef" for char in checksum): raise SourceError("manifest %s has invalid sha256" % path) + if raw.get("dataset") == "2wikimultihopqa" and raw.get("split") not in ("dev", "test"): + raise SourceError("2wikimultihopqa manifest split must be dev or test") return cls( dataset=str(raw["dataset"]), dataset_release=str(raw["dataset_release"]), @@ -98,15 +101,21 @@ def _resolve_source(self, manifest: SourceManifest) -> Path: ) target = Path(os.path.expandvars(manifest.cache_path or Path(manifest.source_uri).name)).expanduser() target.parent.mkdir(parents=True, exist_ok=True) - temporary = target.with_name(target.name + ".part") + descriptor, temporary_name = tempfile.mkstemp(prefix=target.name + ".", suffix=".part", + dir=str(target.parent)) + os.close(descriptor) + temporary = Path(temporary_name) try: urllib.request.urlretrieve(manifest.source_uri, temporary) + self._verify_checksum(temporary, manifest) temporary.replace(target) - except OSError as exc: + except (OSError, SourceError) as exc: try: temporary.unlink() except OSError: pass + if isinstance(exc, SourceError): + raise raise SourceError("download failed for %s/%s: %s" % (manifest.dataset, manifest.split, exc)) from exc return target @@ -153,7 +162,10 @@ def _records(path: Path) -> Iterable[Any]: def _array_records(stream: Any, path: Path) -> Iterator[Any]: decoder = json.JSONDecoder() buffer = "" - first = True + opened = False + expect_value = True + after_comma = False + closed = False while True: chunk = stream.read(64 * 1024) if chunk: @@ -161,18 +173,35 @@ def _array_records(stream: Any, path: Path) -> Iterator[Any]: end = not chunk while True: buffer = buffer.lstrip() - if first: + if closed: + if buffer: + raise SourceError("trailing data after JSON array in %s" % path) + break + if not opened: if not buffer: break if buffer[0] != "[": raise SourceError("JSON source %s must contain an array" % path) buffer = buffer[1:] - first = False + opened = True buffer = buffer.lstrip() if not buffer: break - if buffer[0] == "]": - return + if expect_value and buffer[0] == "]" and not after_comma: + buffer = buffer[1:] + closed = True + continue + if not expect_value: + if buffer[0] != ",": + if buffer[0] == "]": + buffer = buffer[1:] + closed = True + continue + raise SourceError("invalid JSON array separator in %s" % path) + buffer = buffer[1:] + expect_value = True + after_comma = True + continue try: value, consumed = decoder.raw_decode(buffer) except json.JSONDecodeError: @@ -180,14 +209,12 @@ def _array_records(stream: Any, path: Path) -> Iterator[Any]: raise SourceError("truncated JSON array in %s" % path) break yield value - buffer = buffer[consumed:].lstrip() - if buffer.startswith(","): - buffer = buffer[1:] - elif buffer.startswith("]"): - return - elif buffer and end: - raise SourceError("invalid JSON array separator in %s" % path) + buffer = buffer[consumed:] + expect_value = False + after_comma = False if end: + if closed: + return raise SourceError("unterminated JSON array in %s" % path) @staticmethod diff --git a/tools/graphrag/tests/test_source_loader.py b/tools/graphrag/tests/test_source_loader.py index 1a4f16e19..3be0adecc 100644 --- a/tools/graphrag/tests/test_source_loader.py +++ b/tools/graphrag/tests/test_source_loader.py @@ -39,6 +39,27 @@ def test_malformed_record_reports_location(self): with self.assertRaisesRegex(SourceError, "d/dev record 1 field question"): list(SourceLoader().load(manifest)) + def test_streams_json_array(self): + with tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "records.json" + records = [{"id": str(index), "question": "q", "context": [["T", ["x"]]]} + for index in range(1200)] + source.write_text(json.dumps(records), encoding="utf-8") + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + manifest = SourceManifest("d", "r", "s", None, str(source), checksum, "p") + batches = list(SourceLoader(batch_size=100).load_batches(manifest)) + self.assertEqual([100] * 12, [len(batch) for batch in batches]) + + def test_rejects_trailing_comma_and_trailing_data(self): + for contents in ("[{} ,]", "[{}]garbage"): + with self.subTest(contents=contents), tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "records.json" + source.write_text(contents, encoding="utf-8") + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + manifest = SourceManifest("d", "r", "s", None, str(source), checksum, "p") + with self.assertRaises(SourceError): + list(SourceLoader().load(manifest)) + if __name__ == "__main__": unittest.main() From 3400599f6283f125863204a93cae30018eb15ac0 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 16:05:17 +0800 Subject: [PATCH 11/30] feat(graphrag): add deterministic document normalization and chunking --- tools/graphrag/ingest/normalize.py | 105 +++++++++++++++++++++++++ tools/graphrag/tests/test_normalize.py | 33 ++++++++ 2 files changed, 138 insertions(+) create mode 100644 tools/graphrag/ingest/normalize.py create mode 100644 tools/graphrag/tests/test_normalize.py diff --git a/tools/graphrag/ingest/normalize.py b/tools/graphrag/ingest/normalize.py new file mode 100644 index 000000000..132eb215b --- /dev/null +++ b/tools/graphrag/ingest/normalize.py @@ -0,0 +1,105 @@ +"""Canonical document normalization and deterministic character chunking.""" + +from __future__ import annotations + +import hashlib +import re +import unicodedata +from dataclasses import dataclass +from typing import Any, Dict, Iterable, Iterator, List + + +NORMALIZATION_VERSION = "unicode-nfc-whitespace-v1" +CHUNKING_VERSION = "character-window-v1" + + +@dataclass(frozen=True) +class ChunkingConfig: + size: int = 1200 + overlap: int = 120 + token_characters: int = 4 + + def __post_init__(self) -> None: + if self.size < 1 or self.overlap < 0 or self.overlap >= self.size: + raise ValueError("chunk size must be positive and overlap must be in [0, size)") + if self.token_characters < 1: + raise ValueError("token_characters must be positive") + + +def normalize_text(value: str) -> str: + text = unicodedata.normalize("NFC", value) + text = text.replace("\r\n", "\n").replace("\r", "\n") + lines = [re.sub(r"\s+", " ", line).strip() for line in text.split("\n")] + return "\n".join(line for line in lines if line) + + +def documents(records: Iterable[Dict[str, Any]]) -> Iterator[Dict[str, Any]]: + for record in records: + paragraphs = record["paragraphs"] + title = normalize_text(" ".join(item["title"] for item in paragraphs if item["title"])) + body = "\n\n".join(normalize_text(item["text"]) for item in paragraphs) + text = normalize_text((title + "\n\n" + body) if title else body) + if not text: + raise ValueError("record %s normalizes to an empty document" % record["document_id"]) + encoded = text.encode("utf-8") + digest = hashlib.sha256(encoded).hexdigest() + source_id = "%s:%s:%s" % (record["dataset"], record["split"], record["document_id"]) + yield { + "document_id": source_id, + "dataset": record["dataset"], + "dataset_release": record["dataset_release"], + "split": record["split"], + "title": title, + "source_hash": digest, + "text": text, + "question": normalize_text(record["question"]), + "answers": record.get("answers", []), + "supporting_facts": record.get("supporting_facts", []), + "normalization_version": NORMALIZATION_VERSION, + } + + +def chunk_document(document: Dict[str, Any], config: ChunkingConfig) -> List[Dict[str, Any]]: + text = document["text"] + result: List[Dict[str, Any]] = [] + start = 0 + index = 0 + while start < len(text): + end = min(start + config.size, len(text)) + if end < len(text): + boundary = max(text.rfind("\n", start + config.size // 2, end), + text.rfind(" ", start + config.size // 2, end)) + if boundary > start: + end = boundary + chunk_text = text[start:end].strip() + leading = len(text[start:end]) - len(text[start:end].lstrip()) + actual_start = start + leading + actual_end = actual_start + len(chunk_text) + if chunk_text: + digest = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest() + chunk_id = hashlib.sha256((document["document_id"] + "\0" + str(index) + "\0" + digest) + .encode("utf-8")).hexdigest() + result.append({ + "chunk_id": chunk_id, + "document_id": document["document_id"], + "chunk_index": index, + "start_offset": actual_start, + "end_offset": actual_end, + "token_estimate": (len(chunk_text) + config.token_characters - 1) + // config.token_characters, + "text": chunk_text, + "policy_version": CHUNKING_VERSION, + "text_hash": digest, + "source_hash": document["source_hash"], + }) + index += 1 + if end == len(text): + break + start = max(end - config.overlap, start + 1) + return result + + +def chunks(documents_iter: Iterable[Dict[str, Any]], config: ChunkingConfig + ) -> Iterator[Dict[str, Any]]: + for document in documents_iter: + yield from chunk_document(document, config) diff --git a/tools/graphrag/tests/test_normalize.py b/tools/graphrag/tests/test_normalize.py new file mode 100644 index 000000000..23b12f1ed --- /dev/null +++ b/tools/graphrag/tests/test_normalize.py @@ -0,0 +1,33 @@ +import unittest + +from tools.graphrag.ingest.normalize import ChunkingConfig, chunk_document, documents, normalize_text + + +class NormalizeTest(unittest.TestCase): + + def test_normalizes_unicode_and_whitespace(self): + self.assertEqual("Cafe\u0301".replace("e\u0301", "é"), normalize_text("Cafe\u0301 \r\n \t")) + + def test_chunks_have_stable_ids_offsets_and_overlap(self): + record = { + "dataset": "d", "dataset_release": "r", "split": "dev", "document_id": "x", + "question": " q ? ", "paragraphs": [{"title": "T", "text": "alpha beta gamma delta"}], + } + document = list(documents([record]))[0] + config = ChunkingConfig(size=12, overlap=3, token_characters=4) + first = chunk_document(document, config) + second = chunk_document(document, config) + self.assertEqual(first, second) + self.assertGreater(len(first), 1) + self.assertLessEqual(first[1]["start_offset"], first[0]["end_offset"]) + for chunk in first: + self.assertEqual(document["document_id"], chunk["document_id"]) + self.assertEqual(chunk["text"], document["text"][chunk["start_offset"]:chunk["end_offset"]]) + + def test_rejects_invalid_overlap(self): + with self.assertRaises(ValueError): + ChunkingConfig(size=10, overlap=10) + + +if __name__ == "__main__": + unittest.main() From bee429df1bc8e08e8c6e222daa961b11d0fa28d1 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 16:12:45 +0800 Subject: [PATCH 12/30] fix(graphrag): preserve source identity and Java-compatible offsets --- tools/graphrag/ingest/normalize.py | 17 +++++++++++------ tools/graphrag/ingest/sources/source_loader.py | 6 +++++- tools/graphrag/tests/test_normalize.py | 10 ++++++++++ 3 files changed, 26 insertions(+), 7 deletions(-) diff --git a/tools/graphrag/ingest/normalize.py b/tools/graphrag/ingest/normalize.py index 132eb215b..51bf42796 100644 --- a/tools/graphrag/ingest/normalize.py +++ b/tools/graphrag/ingest/normalize.py @@ -33,6 +33,10 @@ def normalize_text(value: str) -> str: return "\n".join(line for line in lines if line) +def _utf16_length(value: str) -> int: + return len(value.encode("utf-16-le")) // 2 + + def documents(records: Iterable[Dict[str, Any]]) -> Iterator[Dict[str, Any]]: for record in records: paragraphs = record["paragraphs"] @@ -41,16 +45,17 @@ def documents(records: Iterable[Dict[str, Any]]) -> Iterator[Dict[str, Any]]: text = normalize_text((title + "\n\n" + body) if title else body) if not text: raise ValueError("record %s normalizes to an empty document" % record["document_id"]) - encoded = text.encode("utf-8") - digest = hashlib.sha256(encoded).hexdigest() - source_id = "%s:%s:%s" % (record["dataset"], record["split"], record["document_id"]) + digest = record.get("source_hash") + if not digest: + digest = hashlib.sha256(text.encode("utf-8")).hexdigest() yield { - "document_id": source_id, + "document_id": record["document_id"], "dataset": record["dataset"], "dataset_release": record["dataset_release"], "split": record["split"], "title": title, "source_hash": digest, + "text_hash": hashlib.sha256(text.encode("utf-8")).hexdigest(), "text": text, "question": normalize_text(record["question"]), "answers": record.get("answers", []), @@ -83,8 +88,8 @@ def chunk_document(document: Dict[str, Any], config: ChunkingConfig) -> List[Dic "chunk_id": chunk_id, "document_id": document["document_id"], "chunk_index": index, - "start_offset": actual_start, - "end_offset": actual_end, + "start_offset": _utf16_length(text[:actual_start]), + "end_offset": _utf16_length(text[:actual_end]), "token_estimate": (len(chunk_text) + config.token_characters - 1) // config.token_characters, "text": chunk_text, diff --git a/tools/graphrag/ingest/sources/source_loader.py b/tools/graphrag/ingest/sources/source_loader.py index 485f94684..47cc19485 100644 --- a/tools/graphrag/ingest/sources/source_loader.py +++ b/tools/graphrag/ingest/sources/source_loader.py @@ -259,7 +259,7 @@ def _canonical_record(manifest: SourceManifest, number: int, raw: Any) -> Dict[s if not isinstance(answers, list) or any(not isinstance(answer, str) for answer in answers): raise SourceError("%s/%s record %d field answer has invalid shape" % ( manifest.dataset, manifest.split, number)) - return { + canonical = { "dataset": manifest.dataset, "dataset_release": manifest.dataset_release, "split": manifest.split, @@ -270,3 +270,7 @@ def _canonical_record(manifest: SourceManifest, number: int, raw: Any) -> Dict[s "paragraphs": paragraphs, "supporting_facts": raw.get("supporting_facts") or raw.get("evidence") or [], } + canonical["source_hash"] = hashlib.sha256( + json.dumps(canonical, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") + ).hexdigest() + return canonical diff --git a/tools/graphrag/tests/test_normalize.py b/tools/graphrag/tests/test_normalize.py index 23b12f1ed..8a3a616fb 100644 --- a/tools/graphrag/tests/test_normalize.py +++ b/tools/graphrag/tests/test_normalize.py @@ -24,6 +24,16 @@ def test_chunks_have_stable_ids_offsets_and_overlap(self): self.assertEqual(document["document_id"], chunk["document_id"]) self.assertEqual(chunk["text"], document["text"][chunk["start_offset"]:chunk["end_offset"]]) + def test_offsets_use_utf16_units_for_non_bmp_text(self): + record = { + "dataset": "d", "dataset_release": "r", "split": "dev", "document_id": "x", + "question": "q", "paragraphs": [{"title": "T", "text": "a😀b"}], + } + document = list(documents([record]))[0] + chunk = chunk_document(document, ChunkingConfig(size=20, overlap=0))[0] + self.assertEqual(0, chunk["start_offset"]) + self.assertEqual(6, chunk["end_offset"]) + def test_rejects_invalid_overlap(self): with self.assertRaises(ValueError): ChunkingConfig(size=10, overlap=10) From b825c9fa42e5948574495dd3c9be5cdb38316f1a Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 16:19:42 +0800 Subject: [PATCH 13/30] feat(graphrag): build deterministic entity graph artifacts --- tools/graphrag/ingest/entity_graph.py | 136 ++++++++++++++++++++++ tools/graphrag/tests/test_entity_graph.py | 26 +++++ 2 files changed, 162 insertions(+) create mode 100644 tools/graphrag/ingest/entity_graph.py create mode 100644 tools/graphrag/tests/test_entity_graph.py diff --git a/tools/graphrag/ingest/entity_graph.py b/tools/graphrag/ingest/entity_graph.py new file mode 100644 index 000000000..b6c669952 --- /dev/null +++ b/tools/graphrag/ingest/entity_graph.py @@ -0,0 +1,136 @@ +"""Deterministic dataset-evidence entity mapping and graph JSONL writer.""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +import unicodedata +from pathlib import Path +from typing import Any, Dict, Iterable, Iterator, List, Tuple + + +EXTRACTOR_VERSION = "dataset-evidence-entity-v1" +GRAPH_SCHEMA_VERSION = "chunk-entity-evidence-v1" + + +def canonical_name(value: str) -> str: + normalized = unicodedata.normalize("NFKC", value) + return " ".join(normalized.casefold().split()) + + +def stable_id(kind: str, *parts: str) -> str: + payload = "\0".join((kind,) + parts).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def _entity(name: str, dataset: str, release: str) -> Dict[str, Any]: + normalized = canonical_name(name) + return { + "entity_id": stable_id("entity", dataset, release, normalized), + "canonical_name": " ".join(unicodedata.normalize("NFC", name).split()), + "normalized_name": normalized, + "entity_type": "CONCEPT", + } + + +def graph_records(records: Iterable[Dict[str, Any]], + chunks: Iterable[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + chunk_list = list(chunks) + by_document: Dict[str, List[Dict[str, Any]]] = {} + for chunk in chunk_list: + by_document.setdefault(chunk["document_id"], []).append(chunk) + vertices: Dict[str, Dict[str, Any]] = {} + edges: Dict[str, Dict[str, Any]] = {} + for record in records: + dataset = record["dataset"] + release = record["dataset_release"] + document_id = record["document_id"] + document_chunks = by_document.get(document_id, []) + chunk_entities: Dict[str, set] = {} + for paragraph in record.get("paragraphs", []): + name = paragraph.get("title", "").strip() + if not name: + continue + entity = _entity(name, dataset, release) + vertices[entity["entity_id"]] = dict(entity, vertex_id=entity["entity_id"], + vertex_type="ENTITY", dataset=dataset, dataset_release=release, + source_document_ids=[document_id]) + for chunk in document_chunks: + if canonical_name(name) in canonical_name(chunk["text"]): + chunk_entities.setdefault(chunk["chunk_id"], set()).add(entity["entity_id"]) + for chunk in document_chunks: + vertex_id = stable_id("chunk-vertex", chunk["chunk_id"]) + vertices[vertex_id] = { + "vertex_id": vertex_id, "vertex_type": "CHUNK", "chunk_id": chunk["chunk_id"], + "document_id": document_id, "dataset": dataset, + "source_document_ids": [document_id], + } + for entity_id in sorted(chunk_entities.get(chunk["chunk_id"], set())): + edge_id = stable_id("mentions", vertex_id, entity_id) + edges[edge_id] = { + "edge_id": edge_id, "source_vertex_id": vertex_id, + "target_vertex_id": entity_id, "relation_type": "MENTIONS", + "source_chunk_ids": [chunk["chunk_id"]], "source_document_ids": [document_id], + } + facts = record.get("evidences", []) or record.get("supporting_facts", []) + for fact in facts: + if not isinstance(fact, (list, tuple)) or len(fact) < 2: + continue + if len(fact) >= 3: + subject, relation, object_name = (str(fact[0]), str(fact[1]), str(fact[2])) + else: + subject, relation, object_name = str(fact[0]), "SUPPORTS", document_id + source_entity = _entity(subject, dataset, release) + target_entity = _entity(object_name, dataset, release) + for entity in (source_entity, target_entity): + vertices[entity["entity_id"]] = dict(entity, vertex_id=entity["entity_id"], + vertex_type="ENTITY", dataset=dataset, dataset_release=release, + source_document_ids=[document_id]) + source_chunk = document_chunks[0]["chunk_id"] if document_chunks else "" + edge_id = stable_id("relation", source_entity["entity_id"], canonical_name(relation), + target_entity["entity_id"], source_chunk) + edges[edge_id] = { + "edge_id": edge_id, "source_vertex_id": source_entity["entity_id"], + "target_vertex_id": target_entity["entity_id"], "relation_type": canonical_name(relation), + "source_chunk_ids": [source_chunk] if source_chunk else [], + "source_document_ids": [document_id], + } + return ([vertices[key] for key in sorted(vertices)], [edges[key] for key in sorted(edges)]) + + +def write_graph_artifact(directory: Path, graph_version: str, + vertices: Iterable[Dict[str, Any]], + edges: Iterable[Dict[str, Any]]) -> Path: + """Write immutable vertex/edge JSONL and publish a manifest last.""" + target = directory / graph_version + target.mkdir(parents=True, exist_ok=True) + staged: List[Tuple[Path, Path]] = [] + try: + for name, values in (("vertices.jsonl", vertices), ("edges.jsonl", edges)): + descriptor, temp_name = tempfile.mkstemp(prefix=name + ".", suffix=".tmp", dir=target) + temp_path = Path(temp_name) + with os.fdopen(descriptor, "w", encoding="utf-8") as output: + for value in values: + output.write(json.dumps(value, sort_keys=True, ensure_ascii=False, + separators=(",", ":")) + "\n") + output.flush() + os.fsync(output.fileno()) + staged.append((temp_path, target / name)) + for temporary, final in staged: + os.replace(temporary, final) + manifest = { + "graph_version": graph_version, "schema_version": GRAPH_SCHEMA_VERSION, + "extractor_version": EXTRACTOR_VERSION, + "vertex_count": sum(1 for _ in (target / "vertices.jsonl").open(encoding="utf-8")), + "edge_count": sum(1 for _ in (target / "edges.jsonl").open(encoding="utf-8")), + } + manifest_path = target / "manifest.json" + temporary = target / "manifest.json.tmp" + temporary.write_text(json.dumps(manifest, sort_keys=True, indent=2) + "\n", encoding="utf-8") + os.replace(temporary, manifest_path) + return manifest_path + finally: + for temporary, _ in staged: + temporary.unlink(missing_ok=True) diff --git a/tools/graphrag/tests/test_entity_graph.py b/tools/graphrag/tests/test_entity_graph.py new file mode 100644 index 000000000..615a8afb8 --- /dev/null +++ b/tools/graphrag/tests/test_entity_graph.py @@ -0,0 +1,26 @@ +import unittest + +from tools.graphrag.ingest.entity_graph import graph_records, stable_id + + +class EntityGraphTest(unittest.TestCase): + + def test_ids_and_records_are_deterministic_with_provenance(self): + record = { + "dataset": "2wiki", "dataset_release": "r1", "document_id": "doc1", + "paragraphs": [{"title": "Alpha", "text": "Alpha has a link."}], + "evidences": [["Alpha", "related to", "Beta"]], + } + chunks = [{"chunk_id": "chunk1", "document_id": "doc1", "text": "Alpha has a link."}] + vertices, edges = graph_records([record], chunks) + again_vertices, again_edges = graph_records([record], chunks) + self.assertEqual(vertices, again_vertices) + self.assertEqual(edges, again_edges) + self.assertTrue(all(vertex.get("source_document_ids") for vertex in vertices)) + self.assertTrue(all(edge.get("source_document_ids") for edge in edges)) + self.assertEqual(stable_id("entity", "2wiki", "r1", "alpha"), + next(v["vertex_id"] for v in vertices if v.get("canonical_name") == "Alpha")) + + +if __name__ == "__main__": + unittest.main() From a12560135dd51241ee0a4563068e5cf114502864 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 16:20:48 +0800 Subject: [PATCH 14/30] fix(graphrag): preserve dataset evidence triples --- tools/graphrag/ingest/sources/source_loader.py | 1 + tools/graphrag/tests/test_source_loader.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/graphrag/ingest/sources/source_loader.py b/tools/graphrag/ingest/sources/source_loader.py index 47cc19485..181e0ac5b 100644 --- a/tools/graphrag/ingest/sources/source_loader.py +++ b/tools/graphrag/ingest/sources/source_loader.py @@ -269,6 +269,7 @@ def _canonical_record(manifest: SourceManifest, number: int, raw: Any) -> Dict[s "answers": answers, "paragraphs": paragraphs, "supporting_facts": raw.get("supporting_facts") or raw.get("evidence") or [], + "evidences": raw.get("evidences") or [], } canonical["source_hash"] = hashlib.sha256( json.dumps(canonical, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8") diff --git a/tools/graphrag/tests/test_source_loader.py b/tools/graphrag/tests/test_source_loader.py index 3be0adecc..3be423389 100644 --- a/tools/graphrag/tests/test_source_loader.py +++ b/tools/graphrag/tests/test_source_loader.py @@ -13,7 +13,8 @@ def test_loads_jsonl_in_stable_batches(self): with tempfile.TemporaryDirectory() as directory: source = Path(directory) / "records.jsonl" records = [ - {"id": "a", "question": "q1", "context": [["T", ["text"]]], "answer": "a1"}, + {"id": "a", "question": "q1", "context": [["T", ["text"]]], "answer": "a1", + "evidences": [["T", "rel", "U"]]}, {"id": "b", "question": "q2", "context": [{"title": "U", "text": "text"}], "answer": ["a2"]}, ] source.write_text("\n".join(json.dumps(record) for record in records), encoding="utf-8") @@ -21,6 +22,7 @@ def test_loads_jsonl_in_stable_batches(self): manifest = SourceManifest("d", "r", "s", None, str(source), checksum, "p") batches = list(SourceLoader(batch_size=1).load_batches(manifest)) self.assertEqual([["a"], ["b"]], [[item["document_id"] for item in batch] for batch in batches]) + self.assertEqual([["T", "rel", "U"]], batches[0][0]["evidences"]) def test_checksum_mismatch_fails_before_parsing(self): with tempfile.TemporaryDirectory() as directory: From c4157df71d9f74f6ff8fad9b7e6cc1b03bf3d20f Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 16:29:34 +0800 Subject: [PATCH 15/30] fix(graphrag): make graph publication deterministic and staged --- tools/graphrag/ingest/entity_graph.py | 76 +++++++++++++++++++-------- 1 file changed, 55 insertions(+), 21 deletions(-) diff --git a/tools/graphrag/ingest/entity_graph.py b/tools/graphrag/ingest/entity_graph.py index b6c669952..31ac7b8de 100644 --- a/tools/graphrag/ingest/entity_graph.py +++ b/tools/graphrag/ingest/entity_graph.py @@ -13,6 +13,10 @@ EXTRACTOR_VERSION = "dataset-evidence-entity-v1" GRAPH_SCHEMA_VERSION = "chunk-entity-evidence-v1" +SUPPORTED_RELATIONS = frozenset({ + "supports", "related to", "relates to", "country", "country of origin", + "country of citizenship", "mentions", +}) def canonical_name(value: str) -> str: @@ -35,14 +39,12 @@ def _entity(name: str, dataset: str, release: str) -> Dict[str, Any]: } -def graph_records(records: Iterable[Dict[str, Any]], - chunks: Iterable[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: +def graph_record_stream(records: Iterable[Dict[str, Any]], + chunks: Iterable[Dict[str, Any]]) -> Iterator[Tuple[Dict[str, Any], Dict[str, Any]]]: chunk_list = list(chunks) by_document: Dict[str, List[Dict[str, Any]]] = {} for chunk in chunk_list: by_document.setdefault(chunk["document_id"], []).append(chunk) - vertices: Dict[str, Dict[str, Any]] = {} - edges: Dict[str, Dict[str, Any]] = {} for record in records: dataset = record["dataset"] release = record["dataset_release"] @@ -54,26 +56,28 @@ def graph_records(records: Iterable[Dict[str, Any]], if not name: continue entity = _entity(name, dataset, release) - vertices[entity["entity_id"]] = dict(entity, vertex_id=entity["entity_id"], + entity_vertex = dict(entity, vertex_id=entity["entity_id"], vertex_type="ENTITY", dataset=dataset, dataset_release=release, source_document_ids=[document_id]) + yield (entity_vertex, None) for chunk in document_chunks: if canonical_name(name) in canonical_name(chunk["text"]): chunk_entities.setdefault(chunk["chunk_id"], set()).add(entity["entity_id"]) for chunk in document_chunks: vertex_id = stable_id("chunk-vertex", chunk["chunk_id"]) - vertices[vertex_id] = { + chunk_vertex = { "vertex_id": vertex_id, "vertex_type": "CHUNK", "chunk_id": chunk["chunk_id"], "document_id": document_id, "dataset": dataset, "source_document_ids": [document_id], } + yield (chunk_vertex, None) for entity_id in sorted(chunk_entities.get(chunk["chunk_id"], set())): edge_id = stable_id("mentions", vertex_id, entity_id) - edges[edge_id] = { + yield (None, { "edge_id": edge_id, "source_vertex_id": vertex_id, "target_vertex_id": entity_id, "relation_type": "MENTIONS", "source_chunk_ids": [chunk["chunk_id"]], "source_document_ids": [document_id], - } + }) facts = record.get("evidences", []) or record.get("supporting_facts", []) for fact in facts: if not isinstance(fact, (list, tuple)) or len(fact) < 2: @@ -82,21 +86,44 @@ def graph_records(records: Iterable[Dict[str, Any]], subject, relation, object_name = (str(fact[0]), str(fact[1]), str(fact[2])) else: subject, relation, object_name = str(fact[0]), "SUPPORTS", document_id + relation = canonical_name(relation) + if not subject.strip() or not object_name.strip() or relation not in SUPPORTED_RELATIONS: + continue source_entity = _entity(subject, dataset, release) target_entity = _entity(object_name, dataset, release) for entity in (source_entity, target_entity): - vertices[entity["entity_id"]] = dict(entity, vertex_id=entity["entity_id"], + yield (dict(entity, vertex_id=entity["entity_id"], vertex_type="ENTITY", dataset=dataset, dataset_release=release, - source_document_ids=[document_id]) - source_chunk = document_chunks[0]["chunk_id"] if document_chunks else "" + source_document_ids=[document_id]), None) + matching = [chunk["chunk_id"] for chunk in document_chunks + if canonical_name(subject) in canonical_name(chunk["text"])] + source_chunk = matching[0] if matching else "" edge_id = stable_id("relation", source_entity["entity_id"], canonical_name(relation), target_entity["entity_id"], source_chunk) - edges[edge_id] = { + yield (None, { "edge_id": edge_id, "source_vertex_id": source_entity["entity_id"], "target_vertex_id": target_entity["entity_id"], "relation_type": canonical_name(relation), "source_chunk_ids": [source_chunk] if source_chunk else [], "source_document_ids": [document_id], - } + }) + + +def graph_records(records: Iterable[Dict[str, Any]], + chunks: Iterable[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]: + vertices: Dict[str, Dict[str, Any]] = {} + edges: Dict[str, Dict[str, Any]] = {} + for vertex, edge in graph_record_stream(records, chunks): + if vertex is not None: + existing = vertices.get(vertex["vertex_id"]) + if existing is None: + vertices[vertex["vertex_id"]] = vertex + else: + existing["canonical_name"] = min(existing.get("canonical_name", ""), + vertex.get("canonical_name", "")) + existing["source_document_ids"] = sorted(set(existing["source_document_ids"]) + | set(vertex["source_document_ids"])) + if edge is not None: + edges[edge["edge_id"]] = edge return ([vertices[key] for key in sorted(vertices)], [edges[key] for key in sorted(edges)]) @@ -105,11 +132,13 @@ def write_graph_artifact(directory: Path, graph_version: str, edges: Iterable[Dict[str, Any]]) -> Path: """Write immutable vertex/edge JSONL and publish a manifest last.""" target = directory / graph_version - target.mkdir(parents=True, exist_ok=True) + if target.exists(): + raise FileExistsError("graph version already exists: %s" % graph_version) + temporary_directory = Path(tempfile.mkdtemp(prefix=graph_version + ".", dir=directory)) staged: List[Tuple[Path, Path]] = [] try: for name, values in (("vertices.jsonl", vertices), ("edges.jsonl", edges)): - descriptor, temp_name = tempfile.mkstemp(prefix=name + ".", suffix=".tmp", dir=target) + descriptor, temp_name = tempfile.mkstemp(prefix=name + ".", suffix=".tmp", dir=temporary_directory) temp_path = Path(temp_name) with os.fdopen(descriptor, "w", encoding="utf-8") as output: for value in values: @@ -117,20 +146,25 @@ def write_graph_artifact(directory: Path, graph_version: str, separators=(",", ":")) + "\n") output.flush() os.fsync(output.fileno()) - staged.append((temp_path, target / name)) + staged.append((temp_path, temporary_directory / name)) for temporary, final in staged: os.replace(temporary, final) manifest = { "graph_version": graph_version, "schema_version": GRAPH_SCHEMA_VERSION, "extractor_version": EXTRACTOR_VERSION, - "vertex_count": sum(1 for _ in (target / "vertices.jsonl").open(encoding="utf-8")), - "edge_count": sum(1 for _ in (target / "edges.jsonl").open(encoding="utf-8")), + "vertex_count": sum(1 for _ in (temporary_directory / "vertices.jsonl").open(encoding="utf-8")), + "edge_count": sum(1 for _ in (temporary_directory / "edges.jsonl").open(encoding="utf-8")), } - manifest_path = target / "manifest.json" - temporary = target / "manifest.json.tmp" + manifest_path = temporary_directory / "manifest.json" + temporary = temporary_directory / "manifest.json.tmp" temporary.write_text(json.dumps(manifest, sort_keys=True, indent=2) + "\n", encoding="utf-8") os.replace(temporary, manifest_path) - return manifest_path + os.replace(temporary_directory, target) + return target / "manifest.json" finally: for temporary, _ in staged: temporary.unlink(missing_ok=True) + if temporary_directory.exists() and not target.exists(): + for child in temporary_directory.iterdir(): + child.unlink(missing_ok=True) + temporary_directory.rmdir() From 985fba00a9befda62f2179c1bfada2f20dbd0c1e Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 16:38:06 +0800 Subject: [PATCH 16/30] feat(ai): add deterministic Lucene BM25 index builder --- .../index/bm25/LuceneBm25IndexBuilder.java | 120 ++++++++++++++++++ .../bm25/LuceneBm25IndexBuilderTest.java | 53 ++++++++ 2 files changed, 173 insertions(+) create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java new file mode 100644 index 000000000..eda448acb --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java @@ -0,0 +1,120 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index.bm25; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.index.Bm25IndexBuilder; +import org.apache.geaflow.ai.retrieval.index.IndexArtifact; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.metadata.IndexBuildMetadata; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.version.IndexVersion; +import org.apache.lucene.analysis.standard.StandardAnalyzer; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.StringField; +import org.apache.lucene.document.TextField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +/** Builds a deterministic Lucene 8 BM25 artifact from canonical chunks. */ +public final class LuceneBm25IndexBuilder implements Bm25IndexBuilder { + + private static final String INDEX_NAME = "bm25"; + private static final String BUILDER_VERSION = "lucene-8.11.2-standard-v1"; + private final Path outputDirectory; + + public LuceneBm25IndexBuilder(Path outputDirectory) { + this.outputDirectory = Objects.requireNonNull(outputDirectory, "outputDirectory"); + } + + @Override + public IndexArtifact build(IngestionContext context, List chunks) throws IOException { + Objects.requireNonNull(context, "context"); + Objects.requireNonNull(chunks, "chunks"); + List ordered = new ArrayList<>(chunks); + ordered.sort(Comparator.comparing(TextChunk::getChunkId)); + String version = context.getGraphVersion().getVersion(); + Path staging = Files.createTempDirectory(outputDirectory, INDEX_NAME + "-" + version + "-"); + Path published = outputDirectory.resolve(INDEX_NAME + "-" + version); + try (Directory directory = FSDirectory.open(staging)) { + IndexWriterConfig config = new IndexWriterConfig(new StandardAnalyzer()); + config.setSimilarity(new org.apache.lucene.search.similarities.BM25Similarity()); + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (TextChunk chunk : ordered) { + Document document = new Document(); + document.add(new StringField("chunkId", chunk.getChunkId(), Field.Store.YES)); + document.add(new StringField("documentId", chunk.getDocumentId(), Field.Store.YES)); + document.add(new TextField("text", chunk.getText(), Field.Store.YES)); + document.add(new StoredField("startOffset", chunk.getStartOffset())); + document.add(new StoredField("endOffset", chunk.getEndOffset())); + document.add(new StoredField("textHash", chunk.getTextHash())); + writer.addDocument(document); + } + writer.commit(); + } + try (DirectoryReader reader = DirectoryReader.open(directory)) { + if (reader.numDocs() != ordered.size()) { + throw new IOException("BM25 document count mismatch"); + } + } + } + if (Files.exists(published)) { + throw new IOException("BM25 artifact already exists: " + published); + } + Files.move(staging, published, StandardCopyOption.ATOMIC_MOVE); + IndexVersion indexVersion = new IndexVersion(INDEX_NAME, version, version); + IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, + INDEX_NAME, BUILDER_VERSION, published.toString(), true); + return new LuceneIndexArtifact(metadata, published); + } + + private static final class LuceneIndexArtifact implements IndexArtifact { + private final IndexBuildMetadata metadata; + private final Path path; + + private LuceneIndexArtifact(IndexBuildMetadata metadata, Path path) { + this.metadata = metadata; + this.path = path; + } + + @Override + public IndexBuildMetadata getMetadata() { + return metadata; + } + + @Override + public void close() throws IOException { + if (!Files.exists(path.resolve("segments.gen")) && !Files.exists(path.resolve("segments_1"))) { + throw new IOException("BM25 artifact is not readable: " + path); + } + } + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java new file mode 100644 index 000000000..de3c48afa --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java @@ -0,0 +1,53 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index.bm25; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; +import org.apache.geaflow.ai.retrieval.metadata.ChunkingConfiguration; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.junit.jupiter.api.Test; + +class LuceneBm25IndexBuilderTest { + + @Test + void buildsReadableDeterministicArtifact() throws Exception { + Path directory = Files.createTempDirectory("bm25-test"); + GraphVersion version = new GraphVersion("g", "v1"); + DatasetManifest manifest = new DatasetManifest("v1", "d", "r", "dev", null, + "cache", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "p", new ChunkingConfiguration(100, 10), "graph-v1", "offline", "v1", 1L); + IngestionContext context = new IngestionContext(manifest, version, "test"); + TextChunk first = new TextChunk("b", "doc", 1, 2, 3, 1, "beta", "p", "hash-b"); + TextChunk second = new TextChunk("a", "doc", 0, 0, 2, 1, "alpha", "p", "hash-a"); + LuceneBm25IndexBuilder builder = new LuceneBm25IndexBuilder(directory); + try (org.apache.geaflow.ai.retrieval.index.IndexArtifact artifact = + builder.build(context, Arrays.asList(first, second))) { + assertEquals("bm25", artifact.getMetadata().getIndexType()); + assertEquals(2, Files.list(directory.resolve("bm25-v1")).count()); + } + } +} From deea988379795fddb50c900acb8b4319702460f8 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 16:42:12 +0800 Subject: [PATCH 17/30] feat(ai): add offline versioned vector index builder --- .../vector/OfflineVectorIndexBuilder.java | 146 ++++++++++++++++++ .../vector/OfflineVectorIndexBuilderTest.java | 58 +++++++ 2 files changed, 204 insertions(+) create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java new file mode 100644 index 000000000..f72d2b85c --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java @@ -0,0 +1,146 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index.vector; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.index.IndexArtifact; +import org.apache.geaflow.ai.retrieval.index.VectorIndexBuilder; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.metadata.IndexBuildMetadata; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.version.IndexVersion; + +/** Builds a deterministic, service-independent vector artifact from precomputed vectors. */ +public final class OfflineVectorIndexBuilder implements VectorIndexBuilder { + + private static final String INDEX_NAME = "vector"; + private static final String BUILDER_VERSION = "offline-vector-v1"; + private final Path outputDirectory; + private final Map vectors; + private final String vectorSource; + private final String vectorVersion; + + public OfflineVectorIndexBuilder(Path outputDirectory, Map vectors, + String vectorSource, String vectorVersion) { + this.outputDirectory = Objects.requireNonNull(outputDirectory, "outputDirectory"); + this.vectors = Objects.requireNonNull(vectors, "vectors"); + this.vectorSource = requireText(vectorSource, "vectorSource"); + this.vectorVersion = requireText(vectorVersion, "vectorVersion"); + } + + @Override + public IndexArtifact build(IngestionContext context, List chunks) throws IOException { + Objects.requireNonNull(context, "context"); + Objects.requireNonNull(chunks, "chunks"); + List ordered = new ArrayList<>(chunks); + ordered.sort(Comparator.comparing(TextChunk::getChunkId)); + int dimensions = validate(ordered); + Files.createDirectories(outputDirectory); + Path staging = Files.createTempFile(outputDirectory, INDEX_NAME + "-", ".tmp"); + Path published = outputDirectory.resolve(INDEX_NAME + "-" + context.getGraphVersion().getVersion() + ".bin"); + try (DataOutputStream output = new DataOutputStream(Files.newOutputStream(staging))) { + output.writeUTF("GEAFLOW-VECTOR-1"); + output.writeUTF(vectorSource); + output.writeUTF(vectorVersion); + output.writeInt(dimensions); + output.writeInt(ordered.size()); + for (TextChunk chunk : ordered) { + output.writeUTF(chunk.getChunkId()); + float[] vector = vectors.get(chunk.getChunkId()); + for (float value : vector) { + output.writeFloat(value); + } + } + } + if (Files.exists(published)) { + Files.deleteIfExists(staging); + throw new IOException("vector artifact already exists: " + published); + } + Files.move(staging, published); + IndexVersion indexVersion = new IndexVersion(INDEX_NAME, + context.getGraphVersion().getVersion(), context.getGraphVersion().getVersion()); + IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, + INDEX_NAME, BUILDER_VERSION + ":" + vectorSource + ":" + vectorVersion, + published.toString(), true); + return new VectorArtifact(metadata, published); + } + + private int validate(List chunks) throws IOException { + int dimensions = -1; + for (TextChunk chunk : chunks) { + float[] vector = vectors.get(chunk.getChunkId()); + if (vector == null) { + throw new IOException("missing vector for chunk " + chunk.getChunkId()); + } + if (dimensions < 0) { + dimensions = vector.length; + } + if (vector.length != dimensions || vector.length == 0) { + throw new IOException("vector dimensions do not align for chunk " + chunk.getChunkId()); + } + for (float value : vector) { + if (Float.isNaN(value) || Float.isInfinite(value)) { + throw new IOException("non-finite vector value for chunk " + chunk.getChunkId()); + } + } + } + if (vectors.size() != chunks.size()) { + throw new IOException("vector count does not match chunk count"); + } + return dimensions; + } + + private static String requireText(String value, String name) { + if (value == null || value.trim().isEmpty()) { + throw new IllegalArgumentException(name + " must not be blank"); + } + return value; + } + + private static final class VectorArtifact implements IndexArtifact { + private final IndexBuildMetadata metadata; + private final Path path; + + private VectorArtifact(IndexBuildMetadata metadata, Path path) { + this.metadata = metadata; + this.path = path; + } + + @Override + public IndexBuildMetadata getMetadata() { + return metadata; + } + + @Override + public void close() throws IOException { + if (!Files.isRegularFile(path) || Files.size(path) == 0) { + throw new IOException("vector artifact is not readable: " + path); + } + } + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java new file mode 100644 index 000000000..0ab8f834c --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index.vector; + +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.HashMap; +import java.util.Map; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.metadata.ChunkingConfiguration; +import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.junit.jupiter.api.Test; + +class OfflineVectorIndexBuilderTest { + + @Test + void rejectsNonFiniteAndMissingVectors() throws Exception { + Path directory = Files.createTempDirectory("vector-test"); + TextChunk chunk = new TextChunk("c1", "d", 0, 0, 1, 1, "x"); + IngestionContext context = context(); + Map missing = new HashMap<>(); + OfflineVectorIndexBuilder missingBuilder = new OfflineVectorIndexBuilder(directory, missing, "fixture", "v1"); + assertThrows(java.io.IOException.class, () -> missingBuilder.build(context, Arrays.asList(chunk))); + Map invalid = new HashMap<>(); + invalid.put("c1", new float[] {Float.NaN}); + OfflineVectorIndexBuilder invalidBuilder = new OfflineVectorIndexBuilder(directory, invalid, "fixture", "v1"); + assertThrows(java.io.IOException.class, () -> invalidBuilder.build(context, Arrays.asList(chunk))); + } + + private static IngestionContext context() { + DatasetManifest manifest = new DatasetManifest("v1", "d", "r", "dev", null, + "cache", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "p", new ChunkingConfiguration(100, 10), "graph-v1", "fixture", "v1", 1L); + return new IngestionContext(manifest, new GraphVersion("g", "v1"), "test"); + } +} From 9a0d2c399b41096060373b8b1585124bfd073c1b Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 19:05:05 +0800 Subject: [PATCH 18/30] feat(graphrag): complete week2 ingestion and readiness pipeline --- .../index/bm25/LuceneBm25IndexBuilder.java | 10 ++- .../vector/OfflineVectorIndexBuilder.java | 28 +++++++- .../retrieval/ingest/IngestionPipeline.java | 6 ++ tools/graphrag/ingest/entity_graph.py | 8 ++- tools/graphrag/ingest/pipeline.py | 69 +++++++++++++++++++ tools/graphrag/tests/test_pipeline.py | 55 +++++++++++++++ 6 files changed, 169 insertions(+), 7 deletions(-) create mode 100644 tools/graphrag/ingest/pipeline.py create mode 100644 tools/graphrag/tests/test_pipeline.py diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java index eda448acb..295df752b 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java @@ -36,6 +36,7 @@ import org.apache.lucene.analysis.standard.StandardAnalyzer; import org.apache.lucene.document.Document; import org.apache.lucene.document.Field; +import org.apache.lucene.document.StoredField; import org.apache.lucene.document.StringField; import org.apache.lucene.document.TextField; import org.apache.lucene.index.DirectoryReader; @@ -61,6 +62,7 @@ public IndexArtifact build(IngestionContext context, List chunks) thr Objects.requireNonNull(chunks, "chunks"); List ordered = new ArrayList<>(chunks); ordered.sort(Comparator.comparing(TextChunk::getChunkId)); + Files.createDirectories(outputDirectory); String version = context.getGraphVersion().getVersion(); Path staging = Files.createTempDirectory(outputDirectory, INDEX_NAME + "-" + version + "-"); Path published = outputDirectory.resolve(INDEX_NAME + "-" + version); @@ -75,7 +77,9 @@ public IndexArtifact build(IngestionContext context, List chunks) thr document.add(new TextField("text", chunk.getText(), Field.Store.YES)); document.add(new StoredField("startOffset", chunk.getStartOffset())); document.add(new StoredField("endOffset", chunk.getEndOffset())); - document.add(new StoredField("textHash", chunk.getTextHash())); + if (chunk.getTextHash() != null) { + document.add(new StoredField("textHash", chunk.getTextHash())); + } writer.addDocument(document); } writer.commit(); @@ -87,7 +91,9 @@ public IndexArtifact build(IngestionContext context, List chunks) thr } } if (Files.exists(published)) { - throw new IOException("BM25 artifact already exists: " + published); + return new LuceneIndexArtifact(new IndexBuildMetadata(context.getGraphVersion(), + new IndexVersion(INDEX_NAME, version, version), INDEX_NAME, BUILDER_VERSION, + published.toString(), true), published); } Files.move(staging, published, StandardCopyOption.ATOMIC_MOVE); IndexVersion indexVersion = new IndexVersion(INDEX_NAME, version, version); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java index f72d2b85c..9957936a8 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java @@ -19,6 +19,7 @@ package org.apache.geaflow.ai.retrieval.index.vector; +import java.io.DataInputStream; import java.io.DataOutputStream; import java.io.IOException; import java.nio.file.Files; @@ -59,7 +60,7 @@ public IndexArtifact build(IngestionContext context, List chunks) thr Objects.requireNonNull(chunks, "chunks"); List ordered = new ArrayList<>(chunks); ordered.sort(Comparator.comparing(TextChunk::getChunkId)); - int dimensions = validate(ordered); + int dimensions = validate(context, ordered); Files.createDirectories(outputDirectory); Path staging = Files.createTempFile(outputDirectory, INDEX_NAME + "-", ".tmp"); Path published = outputDirectory.resolve(INDEX_NAME + "-" + context.getGraphVersion().getVersion() + ".bin"); @@ -71,6 +72,7 @@ public IndexArtifact build(IngestionContext context, List chunks) thr output.writeInt(ordered.size()); for (TextChunk chunk : ordered) { output.writeUTF(chunk.getChunkId()); + output.writeUTF(chunk.getDocumentId()); float[] vector = vectors.get(chunk.getChunkId()); for (float value : vector) { output.writeFloat(value); @@ -79,7 +81,11 @@ public IndexArtifact build(IngestionContext context, List chunks) thr } if (Files.exists(published)) { Files.deleteIfExists(staging); - throw new IOException("vector artifact already exists: " + published); + return new VectorArtifact(new IndexBuildMetadata(context.getGraphVersion(), + new IndexVersion(INDEX_NAME, context.getGraphVersion().getVersion(), + context.getGraphVersion().getVersion()), INDEX_NAME, + BUILDER_VERSION + ":" + vectorSource + ":" + vectorVersion, + published.toString(), true), published); } Files.move(staging, published); IndexVersion indexVersion = new IndexVersion(INDEX_NAME, @@ -90,9 +96,20 @@ public IndexArtifact build(IngestionContext context, List chunks) thr return new VectorArtifact(metadata, published); } - private int validate(List chunks) throws IOException { + private int validate(IngestionContext context, List chunks) throws IOException { + if (!vectorSource.equals(context.getManifest().getVectorSource()) + || !vectorVersion.equals(context.getManifest().getVectorVersion())) { + throw new IOException("vector source/version does not match manifest"); + } + if (chunks.isEmpty()) { + throw new IOException("vector artifact cannot be empty"); + } int dimensions = -1; + java.util.Set chunkIds = new java.util.HashSet<>(); for (TextChunk chunk : chunks) { + if (!chunkIds.add(chunk.getChunkId())) { + throw new IOException("duplicate chunk ID " + chunk.getChunkId()); + } float[] vector = vectors.get(chunk.getChunkId()); if (vector == null) { throw new IOException("missing vector for chunk " + chunk.getChunkId()); @@ -141,6 +158,11 @@ public void close() throws IOException { if (!Files.isRegularFile(path) || Files.size(path) == 0) { throw new IOException("vector artifact is not readable: " + path); } + try (DataInputStream input = new DataInputStream(Files.newInputStream(path))) { + if (!"GEAFLOW-VECTOR-1".equals(input.readUTF())) { + throw new IOException("invalid vector artifact header"); + } + } } } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java index d3ca791ec..3f92dedac 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java @@ -76,6 +76,12 @@ public ImportMetadata run(IngestionContext context) throws IOException { indexArtifacts.add(vector); context.setQualityCounters(new QualityCounters(documents.size(), chunks.size(), extraction.getEntities().size(), extraction.getEdges().size(), 0, 0)); + for (IndexArtifact artifact : indexArtifacts) { + artifact.close(); + } + for (GraphArtifact artifact : graphArtifacts) { + artifact.close(); + } return metadataPublisher.publish(context, graph, Collections.unmodifiableList(indexArtifacts)); } catch (Exception failure) { diff --git a/tools/graphrag/ingest/entity_graph.py b/tools/graphrag/ingest/entity_graph.py index 31ac7b8de..a007cbf48 100644 --- a/tools/graphrag/ingest/entity_graph.py +++ b/tools/graphrag/ingest/entity_graph.py @@ -149,11 +149,15 @@ def write_graph_artifact(directory: Path, graph_version: str, staged.append((temp_path, temporary_directory / name)) for temporary, final in staged: os.replace(temporary, final) + with (temporary_directory / "vertices.jsonl").open(encoding="utf-8") as vertices_file: + vertex_count = sum(1 for _ in vertices_file) + with (temporary_directory / "edges.jsonl").open(encoding="utf-8") as edges_file: + edge_count = sum(1 for _ in edges_file) manifest = { "graph_version": graph_version, "schema_version": GRAPH_SCHEMA_VERSION, "extractor_version": EXTRACTOR_VERSION, - "vertex_count": sum(1 for _ in (temporary_directory / "vertices.jsonl").open(encoding="utf-8")), - "edge_count": sum(1 for _ in (temporary_directory / "edges.jsonl").open(encoding="utf-8")), + "vertex_count": vertex_count, + "edge_count": edge_count, } manifest_path = temporary_directory / "manifest.json" temporary = temporary_directory / "manifest.json.tmp" diff --git a/tools/graphrag/ingest/pipeline.py b/tools/graphrag/ingest/pipeline.py new file mode 100644 index 000000000..f86fbf796 --- /dev/null +++ b/tools/graphrag/ingest/pipeline.py @@ -0,0 +1,69 @@ +"""Restartable offline ingestion orchestration for the Week 2 artifacts.""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from pathlib import Path +from typing import Any, Callable, Dict, Iterable, List, Optional + +from .entity_graph import graph_records, write_graph_artifact +from .normalize import ChunkingConfig, chunks, documents +from .sources.source_loader import SourceLoader, SourceManifest + + +class ImportFailure(RuntimeError): + """Raised after an import is recorded as FAILED.""" + + +def run_import(manifest: SourceManifest, output_directory: Path, loader: SourceLoader, + chunking: ChunkingConfig, bm25_builder: Callable[[Path, List[Dict[str, Any]]], Path], + vector_builder: Callable[[Path, List[Dict[str, Any]]], Path]) -> Dict[str, Any]: + """Build one immutable version; publish the metadata file only after every artifact succeeds.""" + output_directory.mkdir(parents=True, exist_ok=True) + version = hashlib.sha256((manifest.dataset + "\0" + manifest.split + "\0" + manifest.sha256 + + "\0" + manifest.parser_version + "\0" + repr(chunking)) + .encode("utf-8")).hexdigest()[:16] + final = output_directory / version + if final.exists() and (final / "READY.json").is_file(): + return json.loads((final / "READY.json").read_text(encoding="utf-8")) + staging = Path(tempfile.mkdtemp(prefix=version + ".", dir=output_directory)) + metadata: Dict[str, Any] = { + "dataset": manifest.dataset, "dataset_release": manifest.dataset_release, + "split": manifest.split, "sha256": manifest.sha256, "version": version, + "state": "IMPORTING", "parser_version": manifest.parser_version, + "errors": [], + } + try: + source_records = list(loader.load(manifest)) + normalized = list(documents(source_records)) + normalized_chunks = list(chunks(normalized, chunking)) + (staging / "chunks.jsonl").write_text( + "".join(json.dumps(item, sort_keys=True, ensure_ascii=False) + "\n" + for item in normalized_chunks), encoding="utf-8") + vertices, edges = graph_records(source_records, normalized_chunks) + graph_manifest = write_graph_artifact(staging, "graph", vertices, edges) + metadata["state"] = "INDEXING" + bm25_path = bm25_builder(staging, normalized_chunks) + vector_path = vector_builder(staging, normalized_chunks) + metadata.update({ + "state": "READY", "documents": len(normalized), "chunks": len(normalized_chunks), + "vertices": len(vertices), "edges": len(edges), "graph_manifest": "graph/manifest.json", + "bm25_artifact": Path(bm25_path).name, "vector_artifact": Path(vector_path).name, + }) + for relative in (metadata["graph_manifest"], metadata["bm25_artifact"], metadata["vector_artifact"]): + if not (staging / relative).exists(): + raise ImportFailure("artifact does not exist before publication: %s" % relative) + ready_path = staging / "READY.json" + ready_path.write_text(json.dumps(metadata, sort_keys=True, indent=2) + "\n", encoding="utf-8") + os.replace(staging, final) + return metadata + except Exception as failure: + metadata["state"] = "FAILED" + metadata["errors"] = [str(failure)] + (staging / "FAILED.json").write_text(json.dumps(metadata, sort_keys=True, indent=2) + "\n", + encoding="utf-8") + raise ImportFailure("import failed for %s/%s: %s" % + (manifest.dataset, manifest.split, failure)) from failure diff --git a/tools/graphrag/tests/test_pipeline.py b/tools/graphrag/tests/test_pipeline.py new file mode 100644 index 000000000..701dd1a41 --- /dev/null +++ b/tools/graphrag/tests/test_pipeline.py @@ -0,0 +1,55 @@ +import hashlib +import json +import tempfile +import unittest +from pathlib import Path + +from tools.graphrag.ingest.normalize import ChunkingConfig +from tools.graphrag.ingest.pipeline import ImportFailure, run_import +from tools.graphrag.ingest.sources.source_loader import SourceLoader, SourceManifest + + +class PipelineTest(unittest.TestCase): + + def test_publishes_ready_and_is_idempotent(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source.jsonl" + source.write_text(json.dumps({"id": "d", "question": "q", "context": [["A", ["A text"]]]}), + encoding="utf-8") + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + manifest = SourceManifest("d", "r", "dev", None, str(source), checksum, "p") + + def build(name): + return lambda output, chunks: _artifact(output, name, chunks) + + first = run_import(manifest, root / "versions", SourceLoader(), ChunkingConfig(50, 5), + build("bm25"), build("vector")) + second = run_import(manifest, root / "versions", SourceLoader(), ChunkingConfig(50, 5), + build("bm25"), build("vector")) + self.assertEqual(first, second) + self.assertEqual("READY", first["state"]) + + def test_failed_index_is_not_ready(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source.jsonl" + source.write_text(json.dumps({"id": "d", "question": "q", "context": [["A", ["A text"]]]}), + encoding="utf-8") + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + manifest = SourceManifest("d", "r", "dev", None, str(source), checksum, "p") + with self.assertRaises(ImportFailure): + run_import(manifest, root / "versions", SourceLoader(), ChunkingConfig(50, 5), + lambda output, chunks: (_ for _ in ()).throw(ValueError("bad index")), + lambda output, chunks: output / "vector") + self.assertEqual(0, len(list((root / "versions").glob("*/READY.json")))) + + +def _artifact(output, name, chunks): + path = output / (name + ".artifact") + path.write_text(str(len(chunks)), encoding="utf-8") + return path + + +if __name__ == "__main__": + unittest.main() From ace2021c8350ead281330b7caaea7385e86e1bca Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 23:06:46 +0800 Subject: [PATCH 19/30] fix(graphrag): harden ingestion and index artifact validation --- .../index/bm25/LuceneBm25IndexBuilder.java | 13 +++++++++--- .../vector/OfflineVectorIndexBuilder.java | 13 +++++++++++- tools/graphrag/ingest/entity_graph.py | 12 +++++++++-- tools/graphrag/ingest/pipeline.py | 3 ++- .../graphrag/ingest/sources/source_loader.py | 21 ++++++++++++++++++- .../manifests/2wikimultihopqa-dev.json | 2 +- 6 files changed, 55 insertions(+), 9 deletions(-) diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java index 295df752b..bd0bddf78 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java @@ -95,7 +95,11 @@ public IndexArtifact build(IngestionContext context, List chunks) thr new IndexVersion(INDEX_NAME, version, version), INDEX_NAME, BUILDER_VERSION, published.toString(), true), published); } - Files.move(staging, published, StandardCopyOption.ATOMIC_MOVE); + try { + Files.move(staging, published, StandardCopyOption.ATOMIC_MOVE); + } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { + Files.move(staging, published); + } IndexVersion indexVersion = new IndexVersion(INDEX_NAME, version, version); IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, INDEX_NAME, BUILDER_VERSION, published.toString(), true); @@ -118,8 +122,11 @@ public IndexBuildMetadata getMetadata() { @Override public void close() throws IOException { - if (!Files.exists(path.resolve("segments.gen")) && !Files.exists(path.resolve("segments_1"))) { - throw new IOException("BM25 artifact is not readable: " + path); + try (Directory directory = FSDirectory.open(path); + DirectoryReader ignored = DirectoryReader.open(directory)) { + if (ignored.numDocs() < 0) { + throw new IOException("BM25 artifact is not readable: " + path); + } } } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java index 9957936a8..6739a9186 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java @@ -87,7 +87,11 @@ public IndexArtifact build(IngestionContext context, List chunks) thr BUILDER_VERSION + ":" + vectorSource + ":" + vectorVersion, published.toString(), true), published); } - Files.move(staging, published); + try { + Files.move(staging, published, java.nio.file.StandardCopyOption.ATOMIC_MOVE); + } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { + Files.move(staging, published); + } IndexVersion indexVersion = new IndexVersion(INDEX_NAME, context.getGraphVersion().getVersion(), context.getGraphVersion().getVersion()); IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, @@ -162,6 +166,13 @@ public void close() throws IOException { if (!"GEAFLOW-VECTOR-1".equals(input.readUTF())) { throw new IOException("invalid vector artifact header"); } + input.readUTF(); + input.readUTF(); + int dimensions = input.readInt(); + int count = input.readInt(); + if (dimensions <= 0 || count <= 0) { + throw new IOException("invalid vector artifact dimensions/count"); + } } } } diff --git a/tools/graphrag/ingest/entity_graph.py b/tools/graphrag/ingest/entity_graph.py index a007cbf48..22e2a24b2 100644 --- a/tools/graphrag/ingest/entity_graph.py +++ b/tools/graphrag/ingest/entity_graph.py @@ -75,7 +75,7 @@ def graph_record_stream(records: Iterable[Dict[str, Any]], edge_id = stable_id("mentions", vertex_id, entity_id) yield (None, { "edge_id": edge_id, "source_vertex_id": vertex_id, - "target_vertex_id": entity_id, "relation_type": "MENTIONS", + "target_vertex_id": entity_id, "relation_type": "mentions", "source_chunk_ids": [chunk["chunk_id"]], "source_document_ids": [document_id], }) facts = record.get("evidences", []) or record.get("supporting_facts", []) @@ -123,7 +123,14 @@ def graph_records(records: Iterable[Dict[str, Any]], existing["source_document_ids"] = sorted(set(existing["source_document_ids"]) | set(vertex["source_document_ids"])) if edge is not None: - edges[edge["edge_id"]] = edge + existing_edge = edges.get(edge["edge_id"]) + if existing_edge is None: + edges[edge["edge_id"]] = edge + else: + existing_edge["source_document_ids"] = sorted( + set(existing_edge["source_document_ids"]) | set(edge["source_document_ids"])) + existing_edge["source_chunk_ids"] = sorted( + set(existing_edge["source_chunk_ids"]) | set(edge["source_chunk_ids"])) return ([vertices[key] for key in sorted(vertices)], [edges[key] for key in sorted(edges)]) @@ -131,6 +138,7 @@ def write_graph_artifact(directory: Path, graph_version: str, vertices: Iterable[Dict[str, Any]], edges: Iterable[Dict[str, Any]]) -> Path: """Write immutable vertex/edge JSONL and publish a manifest last.""" + directory.mkdir(parents=True, exist_ok=True) target = directory / graph_version if target.exists(): raise FileExistsError("graph version already exists: %s" % graph_version) diff --git a/tools/graphrag/ingest/pipeline.py b/tools/graphrag/ingest/pipeline.py index f86fbf796..667378195 100644 --- a/tools/graphrag/ingest/pipeline.py +++ b/tools/graphrag/ingest/pipeline.py @@ -51,7 +51,8 @@ def run_import(manifest: SourceManifest, output_directory: Path, loader: SourceL metadata.update({ "state": "READY", "documents": len(normalized), "chunks": len(normalized_chunks), "vertices": len(vertices), "edges": len(edges), "graph_manifest": "graph/manifest.json", - "bm25_artifact": Path(bm25_path).name, "vector_artifact": Path(vector_path).name, + "bm25_artifact": str(Path(bm25_path).relative_to(staging)), + "vector_artifact": str(Path(vector_path).relative_to(staging)), }) for relative in (metadata["graph_manifest"], metadata["bm25_artifact"], metadata["vector_artifact"]): if not (staging / relative).exists(): diff --git a/tools/graphrag/ingest/sources/source_loader.py b/tools/graphrag/ingest/sources/source_loader.py index 181e0ac5b..9ebb31b0f 100644 --- a/tools/graphrag/ingest/sources/source_loader.py +++ b/tools/graphrag/ingest/sources/source_loader.py @@ -12,6 +12,7 @@ import os import tempfile import urllib.request +import zipfile from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Iterable, Iterator, List, Optional @@ -138,6 +139,25 @@ def _verify_checksum(path: Path, manifest: SourceManifest) -> None: @staticmethod def _records(path: Path) -> Iterable[Any]: + if path.suffix == ".zip": + try: + with zipfile.ZipFile(path) as archive: + names = sorted(name for name in archive.namelist() + if name.endswith((".json", ".jsonl", ".json.gz", ".jsonl.gz"))) + if len(names) != 1: + raise SourceError("ZIP source %s must contain exactly one JSON data file" % path) + descriptor, temporary_name = tempfile.mkstemp(suffix=Path(names[0]).suffix) + os.close(descriptor) + extracted = Path(temporary_name) + try: + with archive.open(names[0]) as source, extracted.open("wb") as target: + target.write(source.read()) + yield from SourceLoader._records(extracted) + finally: + extracted.unlink(missing_ok=True) + return + except (OSError, zipfile.BadZipFile) as exc: + raise SourceError("cannot parse ZIP source %s: %s" % (path, exc)) from exc opener = gzip.open if path.suffix == ".gz" else open try: with opener(path, "rt", encoding="utf-8") as stream: @@ -263,7 +283,6 @@ def _canonical_record(manifest: SourceManifest, number: int, raw: Any) -> Dict[s "dataset": manifest.dataset, "dataset_release": manifest.dataset_release, "split": manifest.split, - "record_number": number, "document_id": str(document_id), "question": question, "answers": answers, diff --git a/tools/graphrag/manifests/2wikimultihopqa-dev.json b/tools/graphrag/manifests/2wikimultihopqa-dev.json index 03edb9c5f..d66e46efb 100644 --- a/tools/graphrag/manifests/2wikimultihopqa-dev.json +++ b/tools/graphrag/manifests/2wikimultihopqa-dev.json @@ -3,7 +3,7 @@ "dataset_release": "official", "split": "dev", "source_uri": "https://www.dropbox.com/s/npidmtadreo6df2/data.zip?dl=1", - "cache_path": "${GRAPH_RAG_CACHE}/2wikimultihopqa-dev.json", + "cache_path": "${GRAPH_RAG_CACHE}/2wikimultihopqa-dev.zip", "sha256": "${TWIKIMULTIHOPQA_DEV_SHA256}", "parser_version": "2wikimultihopqa-parser-v1" } From e2f59fe7123718bf05b2f8dfbda4fc2893edcf43 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 26 Sep 2026 23:08:52 +0800 Subject: [PATCH 20/30] test(ai): align index fixtures with chunking contract --- .../ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java | 2 +- .../retrieval/index/vector/OfflineVectorIndexBuilderTest.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java index de3c48afa..20d63557c 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java @@ -39,7 +39,7 @@ void buildsReadableDeterministicArtifact() throws Exception { GraphVersion version = new GraphVersion("g", "v1"); DatasetManifest manifest = new DatasetManifest("v1", "d", "r", "dev", null, "cache", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", - "p", new ChunkingConfiguration(100, 10), "graph-v1", "offline", "v1", 1L); + "p", new ChunkingConfiguration("chunk-v1", 100, 10), "graph-v1", "offline", "v1", 1L); IngestionContext context = new IngestionContext(manifest, version, "test"); TextChunk first = new TextChunk("b", "doc", 1, 2, 3, 1, "beta", "p", "hash-b"); TextChunk second = new TextChunk("a", "doc", 0, 0, 2, 1, "alpha", "p", "hash-a"); diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java index 0ab8f834c..ebefa6594 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java @@ -52,7 +52,7 @@ void rejectsNonFiniteAndMissingVectors() throws Exception { private static IngestionContext context() { DatasetManifest manifest = new DatasetManifest("v1", "d", "r", "dev", null, "cache", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", - "p", new ChunkingConfiguration(100, 10), "graph-v1", "fixture", "v1", 1L); + "p", new ChunkingConfiguration("chunk-v1", 100, 10), "graph-v1", "fixture", "v1", 1L); return new IngestionContext(manifest, new GraphVersion("g", "v1"), "test"); } } From c9e9d023e9bab67528676ccb8e29f147edbf2519 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Thu, 8 Oct 2026 21:02:32 +0800 Subject: [PATCH 21/30] feat(graphrag): publish validated retrieval artifacts --- .../ai/retrieval/index/ArtifactIdentity.java | 134 +++++++++ .../index/bm25/LuceneBm25IndexBuilder.java | 128 ++++++--- .../index/graph/GraphArtifactReader.java | 197 ++++++++++++++ .../graph/OfflineGraphArtifactBuilder.java | 178 ++++++++++++ .../vector/OfflineVectorIndexBuilder.java | 154 ++++++++--- .../retrieval/ingest/IngestionPipeline.java | 6 - .../ingest/MetadataStorePublisher.java | 145 ++++++++++ .../ingest/SourceStreamProvider.java | 30 +++ .../retrieval/index/ArtifactIdentityTest.java | 45 ++++ .../bm25/LuceneBm25IndexBuilderTest.java | 54 +++- .../index/graph/GraphArtifactTest.java | 254 ++++++++++++++++++ .../vector/OfflineVectorIndexBuilderTest.java | 69 ++++- .../ingest/IngestionPipelineTest.java | 62 ++++- tools/graphrag/ingest/entity_graph.py | 6 + .../graphrag/ingest/sources/source_loader.py | 75 +++++- tools/graphrag/tests/test_source_loader.py | 22 ++ 16 files changed, 1470 insertions(+), 89 deletions(-) create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentity.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactReader.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/OfflineGraphArtifactBuilder.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataStorePublisher.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/SourceStreamProvider.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentityTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactTest.java diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentity.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentity.java new file mode 100644 index 000000000..5c9015148 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentity.java @@ -0,0 +1,134 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.index; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Path; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.TreeMap; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; + +/** Persisted graph/index identity and canonical chunk fingerprint shared by both indexes. */ +public final class ArtifactIdentity { + private static final String SAFE_VERSION = "[A-Za-z0-9][A-Za-z0-9._-]{0,127}"; + private final String graphName; + private final String graphVersion; + private final String indexVersion; + private final String chunkFingerprint; + + public ArtifactIdentity(String graphName, String graphVersion, String indexVersion, String chunkFingerprint) { + if (blank(graphName) || blank(graphVersion) || blank(indexVersion) || !graphVersion.equals(indexVersion) + || chunkFingerprint == null || !chunkFingerprint.matches("[0-9a-f]{64}")) { + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, "artifact identity missing; rebuild index"); + } + this.graphName = graphName; + this.graphVersion = graphVersion; + this.indexVersion = indexVersion; + this.chunkFingerprint = chunkFingerprint; + } + + public static ArtifactIdentity fromMap(Map metadata) { + return new ArtifactIdentity(metadata.get("graphName"), metadata.get("graphVersion"), + metadata.get("indexVersion"), metadata.get("chunkFingerprint")); + } + + /** Resolves an artifact target while keeping the version inside its output directory. */ + public static Path resolveVersionedPath(Path outputDirectory, String prefix, String version, String suffix) + throws IOException { + if (version == null || !version.matches(SAFE_VERSION)) { + throw new IOException("invalid artifact version for output path"); + } + Path directory = outputDirectory.toAbsolutePath().normalize(); + Path target = directory.resolve(prefix + version + suffix).normalize(); + if (!directory.equals(target.getParent())) { + throw new IOException("artifact path escapes output directory"); + } + return target; + } + + public Map toMap() { + Map metadata = new TreeMap<>(); + metadata.put("graphName", graphName); + metadata.put("graphVersion", graphVersion); + metadata.put("indexVersion", indexVersion); + metadata.put("chunkFingerprint", chunkFingerprint); + return metadata; + } + + public void validate(String name, String graph, String index, List chunks) { + if (!graphName.equals(name) || !graphVersion.equals(graph) || !indexVersion.equals(index) + || !chunkFingerprint.equals(fingerprint(chunks))) { + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, "artifact/fixture identity mismatch"); + } + } + + public void validateVersions(String graph, String index) { + if ((graph != null && !graphVersion.equals(graph)) || (index != null && !indexVersion.equals(index))) { + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, "artifact version mismatch"); + } + } + + public static String fingerprint(List chunks) { + try { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + List ordered = new ArrayList<>(chunks); + ordered.sort(Comparator.comparing(TextChunk::getDocumentId).thenComparing(TextChunk::getChunkId)); + for (TextChunk chunk : ordered) { + append(digest, chunk.getDocumentId()); + append(digest, chunk.getChunkId()); + append(digest, chunk.getText()); + append(digest, Integer.toString(chunk.getChunkIndex())); + append(digest, Integer.toString(chunk.getTokenEstimate())); + append(digest, chunk.getPolicyVersion()); + append(digest, chunk.getTextHash()); + append(digest, Integer.toString(chunk.getStartOffset())); + append(digest, Integer.toString(chunk.getEndOffset())); + } + StringBuilder result = new StringBuilder(); + for (byte value : digest.digest()) { + result.append(String.format(java.util.Locale.ROOT, "%02x", value & 255)); + } + return result.toString(); + } catch (NoSuchAlgorithmException error) { + throw new IllegalStateException("SHA-256 unavailable", error); + } + } + + private static void append(MessageDigest digest, String value) { + if (value == null) { + digest.update(java.nio.ByteBuffer.allocate(4).putInt(-1).array()); + return; + } + byte[] bytes = value.getBytes(StandardCharsets.UTF_8); + digest.update(java.nio.ByteBuffer.allocate(4).putInt(bytes.length).array()); + digest.update(bytes); + } + + private static boolean blank(String value) { + return value == null || value.trim().isEmpty(); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java index bd0bddf78..290fffb54 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java @@ -25,8 +25,11 @@ import java.nio.file.StandardCopyOption; import java.util.ArrayList; import java.util.Comparator; +import java.util.HashSet; import java.util.List; import java.util.Objects; +import java.util.Set; +import org.apache.geaflow.ai.retrieval.index.ArtifactIdentity; import org.apache.geaflow.ai.retrieval.index.Bm25IndexBuilder; import org.apache.geaflow.ai.retrieval.index.IndexArtifact; import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; @@ -62,48 +65,107 @@ public IndexArtifact build(IngestionContext context, List chunks) thr Objects.requireNonNull(chunks, "chunks"); List ordered = new ArrayList<>(chunks); ordered.sort(Comparator.comparing(TextChunk::getChunkId)); - Files.createDirectories(outputDirectory); + Set chunkIds = new HashSet<>(); + for (TextChunk chunk : ordered) { + if (!chunkIds.add(chunk.getChunkId())) { + throw new IOException("duplicate chunk ID " + chunk.getChunkId()); + } + } String version = context.getGraphVersion().getVersion(); + Path published = ArtifactIdentity.resolveVersionedPath(outputDirectory, INDEX_NAME + "-", version, ""); + Files.createDirectories(outputDirectory); Path staging = Files.createTempDirectory(outputDirectory, INDEX_NAME + "-" + version + "-"); - Path published = outputDirectory.resolve(INDEX_NAME + "-" + version); - try (Directory directory = FSDirectory.open(staging)) { - IndexWriterConfig config = new IndexWriterConfig(new StandardAnalyzer()); - config.setSimilarity(new org.apache.lucene.search.similarities.BM25Similarity()); - try (IndexWriter writer = new IndexWriter(directory, config)) { - for (TextChunk chunk : ordered) { - Document document = new Document(); - document.add(new StringField("chunkId", chunk.getChunkId(), Field.Store.YES)); - document.add(new StringField("documentId", chunk.getDocumentId(), Field.Store.YES)); - document.add(new TextField("text", chunk.getText(), Field.Store.YES)); - document.add(new StoredField("startOffset", chunk.getStartOffset())); - document.add(new StoredField("endOffset", chunk.getEndOffset())); - if (chunk.getTextHash() != null) { - document.add(new StoredField("textHash", chunk.getTextHash())); + boolean publishedStaging = false; + try { + try (Directory directory = FSDirectory.open(staging)) { + IndexWriterConfig config = new IndexWriterConfig(new StandardAnalyzer()); + config.setSimilarity(new org.apache.lucene.search.similarities.BM25Similarity()); + try (IndexWriter writer = new IndexWriter(directory, config)) { + for (TextChunk chunk : ordered) { + Document document = new Document(); + document.add(new StringField("chunkId", chunk.getChunkId(), Field.Store.YES)); + document.add(new StringField("documentId", chunk.getDocumentId(), Field.Store.YES)); + document.add(new TextField("text", chunk.getText(), Field.Store.YES)); + document.add(new StoredField("startOffset", chunk.getStartOffset())); + document.add(new StoredField("endOffset", chunk.getEndOffset())); + if (chunk.getTextHash() != null) { + document.add(new StoredField("textHash", chunk.getTextHash())); + } + writer.addDocument(document); } - writer.addDocument(document); + writer.setLiveCommitData(new ArtifactIdentity(context.getGraphVersion().getGraphName(), + version, version, ArtifactIdentity.fingerprint(ordered)).toMap().entrySet()); + writer.commit(); } - writer.commit(); - } - try (DirectoryReader reader = DirectoryReader.open(directory)) { - if (reader.numDocs() != ordered.size()) { - throw new IOException("BM25 document count mismatch"); + try (DirectoryReader reader = DirectoryReader.open(directory)) { + if (reader.numDocs() != ordered.size()) { + throw new IOException("BM25 document count mismatch"); + } } } + if (Files.exists(published)) { + validateExisting(published, ordered, context); + return new LuceneIndexArtifact(new IndexBuildMetadata(context.getGraphVersion(), + new IndexVersion(INDEX_NAME, version, version), INDEX_NAME, BUILDER_VERSION, + published.toString(), true), published); + } + try { + Files.move(staging, published, StandardCopyOption.ATOMIC_MOVE); + } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { + Files.move(staging, published); + } + publishedStaging = true; + IndexVersion indexVersion = new IndexVersion(INDEX_NAME, version, version); + IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, + INDEX_NAME, BUILDER_VERSION, published.toString(), true); + return new LuceneIndexArtifact(metadata, published); + } finally { + if (!publishedStaging) { + deleteRecursively(staging); + } } - if (Files.exists(published)) { - return new LuceneIndexArtifact(new IndexBuildMetadata(context.getGraphVersion(), - new IndexVersion(INDEX_NAME, version, version), INDEX_NAME, BUILDER_VERSION, - published.toString(), true), published); + } + + private static void deleteRecursively(Path path) { + if (!Files.exists(path)) { + return; } - try { - Files.move(staging, published, StandardCopyOption.ATOMIC_MOVE); - } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { - Files.move(staging, published); + try (java.util.stream.Stream paths = Files.walk(path)) { + paths.sorted(Comparator.reverseOrder()).forEach(candidate -> { + try { + Files.deleteIfExists(candidate); + } catch (IOException ignored) { + // Best-effort cleanup must not hide the original build failure. + } + }); + } catch (IOException ignored) { + // Best-effort cleanup must not hide the original build failure. + } + } + + private static void validateExisting(Path path, List chunks, IngestionContext context) throws IOException { + try (Directory directory = FSDirectory.open(path); + DirectoryReader reader = DirectoryReader.open(directory)) { + ArtifactIdentity.fromMap(reader.getIndexCommit().getUserData()).validate( + context.getGraphVersion().getGraphName(), context.getGraphVersion().getVersion(), + context.getGraphVersion().getVersion(), chunks); + if (reader.numDocs() != chunks.size()) { + throw new IOException("existing BM25 document count mismatch"); + } + for (TextChunk chunk : chunks) { + org.apache.lucene.search.IndexSearcher searcher = new org.apache.lucene.search.IndexSearcher(reader); + org.apache.lucene.search.TopDocs found = searcher.search(new org.apache.lucene.search.TermQuery( + new org.apache.lucene.index.Term("chunkId", chunk.getChunkId())), 2); + if (found.totalHits.value != 1) { + throw new IOException("existing BM25 chunk mismatch: " + chunk.getChunkId()); + } + Document document = searcher.doc(found.scoreDocs[0].doc); + if (!chunk.getDocumentId().equals(document.get("documentId")) + || !chunk.getText().equals(document.get("text"))) { + throw new IOException("existing BM25 content mismatch: " + chunk.getChunkId()); + } + } } - IndexVersion indexVersion = new IndexVersion(INDEX_NAME, version, version); - IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, - INDEX_NAME, BUILDER_VERSION, published.toString(), true); - return new LuceneIndexArtifact(metadata, published); } private static final class LuceneIndexArtifact implements IndexArtifact { diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactReader.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactReader.java new file mode 100644 index 000000000..d8d3e71b3 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactReader.java @@ -0,0 +1,197 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index.graph; + +import java.io.Closeable; +import java.io.DataInputStream; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.channel.graph.GraphNeighborProvider; +import org.apache.geaflow.ai.retrieval.channel.graph.InMemoryGraphNeighborProvider; +import org.apache.geaflow.ai.retrieval.index.ArtifactIdentity; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; + +/** Reader for the deterministic GEAFLOW-GRAPH-1 binary artifact. */ +public final class GraphArtifactReader implements Closeable { + private static final String MAGIC = "GEAFLOW-GRAPH-1"; + private final ArtifactIdentity identity; + private final List entities; + private final List edges; + private final GraphNeighborProvider neighbors; + + public GraphArtifactReader(Path artifact, String expectedGraphVersion, String expectedIndexVersion) { + if (artifact == null || !Files.isRegularFile(artifact)) { + throw notReady("missing graph artifact: " + artifact); + } + try (DataInputStream input = new DataInputStream(Files.newInputStream(artifact))) { + if (!MAGIC.equals(input.readUTF())) { + throw new IOException("invalid graph artifact magic"); + } + identity = new ArtifactIdentity(input.readUTF(), input.readUTF(), input.readUTF(), input.readUTF()); + identity.validateVersions(expectedGraphVersion, expectedIndexVersion); + int entityCount = input.readInt(); + int edgeCount = input.readInt(); + if (entityCount < 0 || edgeCount < 0) { + throw new IOException("invalid graph record counts"); + } + long minimumRecordBytes = (long) entityCount * 14L + (long) edgeCount * 12L; + if (minimumRecordBytes > input.available()) { + throw new IOException("graph record counts exceed artifact size"); + } + entities = new ArrayList<>(entityCount); + Set entityIds = new HashSet<>(); + String previousEntityId = null; + for (int index = 0; index < entityCount; index++) { + String entityId = input.readUTF(); + final String canonical = input.readUTF(); + final String type = input.readUTF(); + int aliasCount = input.readInt(); + if (aliasCount < 0 || !entityIds.add(entityId) + || previousEntityId != null && previousEntityId.compareTo(entityId) >= 0) { + throw new IOException("invalid or duplicate graph entity"); + } + validateStringCount(input, aliasCount); + previousEntityId = entityId; + List aliases = readStrings(input, aliasCount); + int sourceCount = input.readInt(); + if (sourceCount < 0) { + throw new IOException("invalid graph entity provenance count"); + } + validateStringCount(input, sourceCount); + List sources = readStrings(input, sourceCount); + entities.add(new EntityRef(entityId, canonical, aliases, type, sources)); + } + edges = new ArrayList<>(edgeCount); + Set edgeIds = new HashSet<>(); + String previousEdgeId = null; + for (int index = 0; index < edgeCount; index++) { + String edgeId = input.readUTF(); + final String label = input.readUTF(); + String source = input.readUTF(); + String target = input.readUTF(); + int sourceCount = input.readInt(); + if (sourceCount < 0 || !edgeIds.add(edgeId) + || !entityIds.contains(source) || !entityIds.contains(target) + || previousEdgeId != null && previousEdgeId.compareTo(edgeId) >= 0) { + throw new IOException("invalid or duplicate graph edge"); + } + validateStringCount(input, sourceCount); + previousEdgeId = edgeId; + edges.add(new GraphEdgeRef(edgeId, label, source, target, readStrings(input, sourceCount))); + } + for (GraphEdgeRef edge : edges) { + if (!entityIds.contains(edge.getSourceEntityId()) + || !entityIds.contains(edge.getTargetEntityId())) { + throw new IOException("graph edge references unknown entity"); + } + } + if (input.read() != -1) { + throw new IOException("trailing graph artifact data"); + } + neighbors = new InMemoryGraphNeighborProvider(edges); + } catch (RetrievalException error) { + throw error; + } catch (Exception error) { + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, + "unable to read graph artifact", error); + } + } + + private static List readStrings(DataInputStream input, int count) throws IOException { + validateStringCount(input, count); + List values = new ArrayList<>(count); + Set seen = new HashSet<>(); + for (int index = 0; index < count; index++) { + String value = input.readUTF(); + if (!seen.add(value)) { + throw new IOException("duplicate graph reference"); + } + values.add(value); + } + return values; + } + + private static void validateStringCount(DataInputStream input, int count) throws IOException { + if (count < 0 || (long) count * 2L > input.available()) { + throw new IOException("graph string count exceeds remaining artifact data"); + } + } + + private static RetrievalException notReady(String message) { + return new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, message); + } + + public ArtifactIdentity getIdentity() { + return identity; + } + + public List getEntities() { + return java.util.Collections.unmodifiableList(entities); + } + + public List getEdges() { + return java.util.Collections.unmodifiableList(edges); + } + + public GraphNeighborProvider getNeighbors() { + return neighbors; + } + + /** Validates that all graph provenance references point at the supplied chunks. */ + public void validateReferences(List chunks) { + if (chunks == null) { + throw notReady("graph fixture chunks are missing"); + } + Set chunkIds = new HashSet<>(); + for (TextChunk chunk : chunks) { + if (chunk == null || !chunkIds.add(chunk.getChunkId())) { + throw notReady("invalid or duplicate fixture chunk"); + } + } + for (EntityRef entity : entities) { + validateChunkReferences(entity.getSourceChunkIds(), chunkIds, "entity " + entity.getEntityId()); + } + for (GraphEdgeRef edge : edges) { + validateChunkReferences(edge.getSourceChunkIds(), chunkIds, "edge " + edge.getEdgeId()); + } + } + + private static void validateChunkReferences(List references, Set chunkIds, String owner) { + for (String chunkId : references) { + if (!chunkIds.contains(chunkId)) { + throw notReady(owner + " references unknown chunk " + chunkId); + } + } + } + + @Override + public void close() { + // The artifact is loaded into immutable memory; no resources remain open. + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/OfflineGraphArtifactBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/OfflineGraphArtifactBuilder.java new file mode 100644 index 000000000..c796cb27e --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/OfflineGraphArtifactBuilder.java @@ -0,0 +1,178 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index.graph; + +import java.io.DataOutputStream; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import org.apache.geaflow.ai.retrieval.index.ArtifactIdentity; +import org.apache.geaflow.ai.retrieval.ingest.GraphArtifact; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.metadata.GraphBuildMetadata; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; + +/** Builds and atomically publishes a deterministic graph artifact. */ +public final class OfflineGraphArtifactBuilder { + private static final String MAGIC = "GEAFLOW-GRAPH-1"; + private final Path outputDirectory; + + public OfflineGraphArtifactBuilder(Path outputDirectory) { + this.outputDirectory = outputDirectory; + } + + public GraphArtifact build(IngestionContext context, List chunks, + List entities, List edges) throws IOException { + Path published = ArtifactIdentity.resolveVersionedPath(outputDirectory, "graph-", + context.getGraphVersion().getVersion(), ".bin"); + Files.createDirectories(outputDirectory); + List orderedEntities = new ArrayList<>(entities); + orderedEntities.sort(Comparator.comparing(EntityRef::getEntityId)); + List orderedEdges = new ArrayList<>(edges); + orderedEdges.sort(Comparator.comparing(GraphEdgeRef::getEdgeId)); + validate(orderedEntities, orderedEdges, chunks); + Path staging = Files.createTempFile(outputDirectory, "graph-", ".tmp"); + try { + try (DataOutputStream output = new DataOutputStream(Files.newOutputStream(staging))) { + output.writeUTF(MAGIC); + output.writeUTF(context.getGraphVersion().getGraphName()); + output.writeUTF(context.getGraphVersion().getVersion()); + output.writeUTF(context.getGraphVersion().getVersion()); + output.writeUTF(ArtifactIdentity.fingerprint(chunks)); + output.writeInt(orderedEntities.size()); + output.writeInt(orderedEdges.size()); + for (EntityRef entity : orderedEntities) { + output.writeUTF(entity.getEntityId()); + output.writeUTF(entity.getCanonicalName()); + output.writeUTF(entity.getType()); + writeStrings(output, entity.getAliases()); + writeStrings(output, entity.getSourceChunkIds()); + } + for (GraphEdgeRef edge : orderedEdges) { + output.writeUTF(edge.getEdgeId()); + output.writeUTF(edge.getLabel()); + output.writeUTF(edge.getSourceEntityId()); + output.writeUTF(edge.getTargetEntityId()); + writeStrings(output, edge.getSourceChunkIds()); + } + } + if (Files.exists(published)) { + validateExisting(published, context, orderedEntities, orderedEdges, chunks); + return new Artifact(new GraphBuildMetadata(context.getGraphVersion(), published.toString(), + "offline-graph", "offline-graph-v1", true), published); + } + try { + Files.move(staging, published, java.nio.file.StandardCopyOption.ATOMIC_MOVE); + } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { + Files.move(staging, published); + } + return new Artifact(new GraphBuildMetadata(context.getGraphVersion(), published.toString(), + "offline-graph", "offline-graph-v1", true), published); + } finally { + Files.deleteIfExists(staging); + } + } + + private static void writeStrings(DataOutputStream output, List values) throws IOException { + output.writeInt(values.size()); + for (String value : values) { + output.writeUTF(value); + } + } + + private static void validate(List entities, List edges, List chunks) + throws IOException { + Set entityIds = new HashSet<>(); + for (EntityRef entity : entities) { + if (!entityIds.add(entity.getEntityId())) { + throw new IOException("duplicate graph entity"); + } + } + Set chunkIds = new HashSet<>(); + for (TextChunk chunk : chunks) { + if (chunk == null || !chunkIds.add(chunk.getChunkId())) { + throw new IOException("invalid or duplicate chunk"); + } + } + for (EntityRef entity : entities) { + for (String chunkId : entity.getSourceChunkIds()) { + if (!chunkIds.contains(chunkId)) { + throw new IOException("graph entity references unknown chunk"); + } + } + } + Set edgeIds = new HashSet<>(); + for (GraphEdgeRef edge : edges) { + if (!edgeIds.add(edge.getEdgeId()) || !entityIds.contains(edge.getSourceEntityId()) + || !entityIds.contains(edge.getTargetEntityId())) { + throw new IOException("invalid graph edge"); + } + for (String chunkId : edge.getSourceChunkIds()) { + if (!chunkIds.contains(chunkId)) { + throw new IOException("graph edge references unknown chunk"); + } + } + } + } + + private static void validateExisting(Path path, IngestionContext context, + List entities, List edges, + List chunks) throws IOException { + String version = context.getGraphVersion().getVersion(); + try (GraphArtifactReader reader = new GraphArtifactReader(path, version, version)) { + reader.getIdentity().validate(context.getGraphVersion().getGraphName(), version, version, chunks); + reader.validateReferences(chunks); + if (!entities.equals(reader.getEntities()) || !edges.equals(reader.getEdges())) { + throw new IOException("existing graph artifact content mismatch"); + } + } catch (org.apache.geaflow.ai.retrieval.api.model.RetrievalException error) { + throw new IOException("existing graph artifact is invalid", error); + } + } + + private static final class Artifact implements GraphArtifact { + private final GraphBuildMetadata metadata; + private final Path path; + + private Artifact(GraphBuildMetadata metadata, Path path) { + this.metadata = metadata; + this.path = path; + } + + @Override + public GraphBuildMetadata getMetadata() { + return metadata; + } + + @Override + public void close() throws IOException { + if (!Files.isRegularFile(path) || Files.size(path) == 0) { + throw new IOException("graph artifact is not readable: " + path); + } + } + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java index 6739a9186..c9596a3d5 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java @@ -26,9 +26,12 @@ import java.nio.file.Path; import java.util.ArrayList; import java.util.Comparator; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; +import org.apache.geaflow.ai.retrieval.index.ArtifactIdentity; import org.apache.geaflow.ai.retrieval.index.IndexArtifact; import org.apache.geaflow.ai.retrieval.index.VectorIndexBuilder; import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; @@ -40,7 +43,7 @@ public final class OfflineVectorIndexBuilder implements VectorIndexBuilder { private static final String INDEX_NAME = "vector"; - private static final String BUILDER_VERSION = "offline-vector-v1"; + private static final String BUILDER_VERSION = "offline-vector-v2"; private final Path outputDirectory; private final Map vectors; private final String vectorSource; @@ -58,46 +61,63 @@ public OfflineVectorIndexBuilder(Path outputDirectory, Map vect public IndexArtifact build(IngestionContext context, List chunks) throws IOException { Objects.requireNonNull(context, "context"); Objects.requireNonNull(chunks, "chunks"); + String version = context.getGraphVersion().getVersion(); + Path published = ArtifactIdentity.resolveVersionedPath(outputDirectory, INDEX_NAME + "-", version, ".bin"); List ordered = new ArrayList<>(chunks); ordered.sort(Comparator.comparing(TextChunk::getChunkId)); int dimensions = validate(context, ordered); Files.createDirectories(outputDirectory); Path staging = Files.createTempFile(outputDirectory, INDEX_NAME + "-", ".tmp"); - Path published = outputDirectory.resolve(INDEX_NAME + "-" + context.getGraphVersion().getVersion() + ".bin"); - try (DataOutputStream output = new DataOutputStream(Files.newOutputStream(staging))) { - output.writeUTF("GEAFLOW-VECTOR-1"); - output.writeUTF(vectorSource); - output.writeUTF(vectorVersion); - output.writeInt(dimensions); - output.writeInt(ordered.size()); - for (TextChunk chunk : ordered) { - output.writeUTF(chunk.getChunkId()); - output.writeUTF(chunk.getDocumentId()); - float[] vector = vectors.get(chunk.getChunkId()); - for (float value : vector) { - output.writeFloat(value); + boolean publishedStaging = false; + try { + try (DataOutputStream output = new DataOutputStream(Files.newOutputStream(staging))) { + output.writeUTF("GEAFLOW-VECTOR-2"); + output.writeUTF(context.getGraphVersion().getGraphName()); + output.writeUTF(context.getGraphVersion().getVersion()); + output.writeUTF(context.getGraphVersion().getVersion()); + output.writeUTF(ArtifactIdentity.fingerprint(ordered)); + output.writeUTF(vectorSource); + output.writeUTF(vectorVersion); + output.writeInt(dimensions); + output.writeInt(ordered.size()); + for (TextChunk chunk : ordered) { + output.writeUTF(chunk.getChunkId()); + output.writeUTF(chunk.getDocumentId()); + float[] vector = vectors.get(chunk.getChunkId()); + for (float value : vector) { + output.writeFloat(value); + } + } + } + if (Files.exists(published)) { + VectorArtifact.validateArtifact(published, ordered, dimensions, vectorSource, vectorVersion, context, vectors); + return new VectorArtifact(new IndexBuildMetadata(context.getGraphVersion(), + new IndexVersion(INDEX_NAME, context.getGraphVersion().getVersion(), + context.getGraphVersion().getVersion()), INDEX_NAME, + BUILDER_VERSION + ":" + vectorSource + ":" + vectorVersion, + published.toString(), true), published); + } + try { + Files.move(staging, published, java.nio.file.StandardCopyOption.ATOMIC_MOVE); + } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { + Files.move(staging, published); + } + publishedStaging = true; + IndexVersion indexVersion = new IndexVersion(INDEX_NAME, + context.getGraphVersion().getVersion(), context.getGraphVersion().getVersion()); + IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, + INDEX_NAME, BUILDER_VERSION + ":" + vectorSource + ":" + vectorVersion, + published.toString(), true); + return new VectorArtifact(metadata, published); + } finally { + if (!publishedStaging) { + try { + Files.deleteIfExists(staging); + } catch (IOException ignored) { + // Best-effort cleanup must not hide the original build failure. } } } - if (Files.exists(published)) { - Files.deleteIfExists(staging); - return new VectorArtifact(new IndexBuildMetadata(context.getGraphVersion(), - new IndexVersion(INDEX_NAME, context.getGraphVersion().getVersion(), - context.getGraphVersion().getVersion()), INDEX_NAME, - BUILDER_VERSION + ":" + vectorSource + ":" + vectorVersion, - published.toString(), true), published); - } - try { - Files.move(staging, published, java.nio.file.StandardCopyOption.ATOMIC_MOVE); - } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { - Files.move(staging, published); - } - IndexVersion indexVersion = new IndexVersion(INDEX_NAME, - context.getGraphVersion().getVersion(), context.getGraphVersion().getVersion()); - IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, - INDEX_NAME, BUILDER_VERSION + ":" + vectorSource + ":" + vectorVersion, - published.toString(), true); - return new VectorArtifact(metadata, published); } private int validate(IngestionContext context, List chunks) throws IOException { @@ -124,11 +144,16 @@ private int validate(IngestionContext context, List chunks) throws IO if (vector.length != dimensions || vector.length == 0) { throw new IOException("vector dimensions do not align for chunk " + chunk.getChunkId()); } + double normSquared = 0.0; for (float value : vector) { + normSquared += (double) value * value; if (Float.isNaN(value) || Float.isInfinite(value)) { throw new IOException("non-finite vector value for chunk " + chunk.getChunkId()); } } + if (normSquared == 0.0) { + throw new IOException("zero-norm vector for chunk " + chunk.getChunkId()); + } } if (vectors.size() != chunks.size()) { throw new IOException("vector count does not match chunk count"); @@ -163,16 +188,73 @@ public void close() throws IOException { throw new IOException("vector artifact is not readable: " + path); } try (DataInputStream input = new DataInputStream(Files.newInputStream(path))) { - if (!"GEAFLOW-VECTOR-1".equals(input.readUTF())) { + if (!"GEAFLOW-VECTOR-2".equals(input.readUTF())) { throw new IOException("invalid vector artifact header"); } - input.readUTF(); - input.readUTF(); + for (int field = 0; field < 6; field++) { + input.readUTF(); + } int dimensions = input.readInt(); int count = input.readInt(); if (dimensions <= 0 || count <= 0) { throw new IOException("invalid vector artifact dimensions/count"); } + for (int index = 0; index < count; index++) { + input.readUTF(); + input.readUTF(); + for (int dimension = 0; dimension < dimensions; dimension++) { + float value = input.readFloat(); + if (Float.isNaN(value) || Float.isInfinite(value)) { + throw new IOException("invalid vector value"); + } + } + } + if (input.read() != -1) { + throw new IOException("trailing data in vector artifact"); + } + } + } + + private static void validateArtifact(Path artifact, List chunks, int dimensions, + String expectedSource, String expectedVersion, IngestionContext context, + Map vectors) throws IOException { + try (DataInputStream input = new DataInputStream(Files.newInputStream(artifact))) { + if (!"GEAFLOW-VECTOR-2".equals(input.readUTF())) { + throw new IOException("legacy vector artifact; rebuild index"); + } + new ArtifactIdentity(input.readUTF(), input.readUTF(), input.readUTF(), input.readUTF()).validate( + context.getGraphVersion().getGraphName(), context.getGraphVersion().getVersion(), + context.getGraphVersion().getVersion(), chunks); + if (!expectedSource.equals(input.readUTF()) || !expectedVersion.equals(input.readUTF())) { + throw new IOException("vector artifact metadata mismatch"); + } + if (input.readInt() != dimensions || input.readInt() != chunks.size()) { + throw new IOException("vector artifact shape mismatch"); + } + Map expectedChunks = new java.util.HashMap<>(); + Set expected = new HashSet<>(); + for (TextChunk chunk : chunks) { + expected.add(chunk.getChunkId()); + expectedChunks.put(chunk.getChunkId(), chunk); + } + Set actual = new HashSet<>(); + for (int index = 0; index < chunks.size(); index++) { + String chunkId = input.readUTF(); + String documentId = input.readUTF(); + if (!actual.add(chunkId) || !expected.contains(chunkId) + || !expectedChunks.get(chunkId).getDocumentId().equals(documentId)) { + throw new IOException("vector artifact chunk mismatch"); + } + for (int dimension = 0; dimension < dimensions; dimension++) { + float value = input.readFloat(); + if (Float.compare(value, vectors.get(chunkId)[dimension]) != 0) { + throw new IOException("existing vector content mismatch"); + } + } + } + if (!actual.equals(expected) || input.read() != -1) { + throw new IOException("vector artifact content mismatch"); + } } } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java index 3f92dedac..d3ca791ec 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java @@ -76,12 +76,6 @@ public ImportMetadata run(IngestionContext context) throws IOException { indexArtifacts.add(vector); context.setQualityCounters(new QualityCounters(documents.size(), chunks.size(), extraction.getEntities().size(), extraction.getEdges().size(), 0, 0)); - for (IndexArtifact artifact : indexArtifacts) { - artifact.close(); - } - for (GraphArtifact artifact : graphArtifacts) { - artifact.close(); - } return metadataPublisher.publish(context, graph, Collections.unmodifiableList(indexArtifacts)); } catch (Exception failure) { diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataStorePublisher.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataStorePublisher.java new file mode 100644 index 000000000..f4a03e60c --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataStorePublisher.java @@ -0,0 +1,145 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.io.IOException; +import java.io.InputStream; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.index.IndexArtifact; +import org.apache.geaflow.ai.retrieval.metadata.GraphBuildMetadata; +import org.apache.geaflow.ai.retrieval.metadata.ImportMetadata; +import org.apache.geaflow.ai.retrieval.metadata.ImportState; +import org.apache.geaflow.ai.retrieval.metadata.MetadataException; +import org.apache.geaflow.ai.retrieval.metadata.MetadataStore; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; + +/** + * Metadata publisher backed by the metadata store lifecycle. + * + *

The publisher verifies the source before recording graph and index metadata and captures the + * published pointer at begin time for compare-and-set publication. The source stream is opened + * only for verification and is always closed by this publisher.

+ */ +public final class MetadataStorePublisher implements MetadataPublisher { + + private static final List REQUIRED_INDEXES = Collections.unmodifiableList( + Arrays.asList("bm25", "vector")); + + private final MetadataStore store; + private final SourceStreamProvider sourceProvider; + private final Map attempts = new HashMap<>(); + + public MetadataStorePublisher(MetadataStore store, SourceStreamProvider sourceProvider) { + this.store = Objects.requireNonNull(store, "store"); + this.sourceProvider = Objects.requireNonNull(sourceProvider, "sourceProvider"); + } + + @Override + public synchronized void begin(IngestionContext context) { + Objects.requireNonNull(context, "context"); + GraphVersion version = context.getGraphVersion(); + if (attempts.containsKey(version)) { + throw new MetadataException(MetadataException.Code.VERSION_CONFLICT, + "ingestion attempt already started"); + } + GraphBuildMetadata graph = new GraphBuildMetadata(version, null, "graph", "pending", + false, REQUIRED_INDEXES); + GraphVersion expected = store.getPublishedVersion(version.getGraphName()).orElse(null); + store.begin(context.getManifest(), context.getImporterVersion(), graph); + attempts.put(version, new Attempt(expected)); + } + + @Override + public synchronized ImportMetadata publish(IngestionContext context, GraphArtifact graph, + List indexes) throws IOException { + Objects.requireNonNull(context, "context"); + Objects.requireNonNull(graph, "graph"); + Objects.requireNonNull(indexes, "indexes"); + GraphVersion version = context.getGraphVersion(); + final Attempt attempt = requireAttempt(version); + GraphBuildMetadata graphMetadata = graph.getMetadata(); + if (!version.equals(graphMetadata.getGraphVersion()) + || !REQUIRED_INDEXES.equals(graphMetadata.getRequiredIndexes())) { + throw new MetadataException(MetadataException.Code.INVALID_METADATA, + "graph metadata must match the ingestion version and required indexes"); + } + try (InputStream source = openSource(context)) { + store.verifySource(version, source); + } + store.finishImport(version, graphMetadata, context.getQualityCounters()); + for (IndexArtifact index : indexes) { + Objects.requireNonNull(index, "index"); + store.putIndex(version, index.getMetadata()); + } + ImportMetadata published = store.publish(version, attempt.expectedPublishedVersion); + attempts.remove(version); + return published; + } + + @Override + public synchronized void fail(IngestionContext context, Exception failure) throws IOException { + Objects.requireNonNull(context, "context"); + Objects.requireNonNull(failure, "failure"); + GraphVersion version = context.getGraphVersion(); + Attempt attempt = attempts.remove(version); + if (attempt == null) { + return; + } + ImportMetadata current = store.find(version).orElse(null); + if (current == null || current.getState() == ImportState.FAILED + || current.getState() == ImportState.READY) { + return; + } + MetadataException.Code code = failure instanceof MetadataException + ? ((MetadataException) failure).getCode() : MetadataException.Code.INVALID_METADATA; + String reason = failure.getMessage() == null ? failure.getClass().getName() : failure.getMessage(); + store.fail(version, code, reason); + } + + private InputStream openSource(IngestionContext context) throws IOException { + InputStream source = sourceProvider.open(context); + if (source == null) { + throw new IOException("source provider returned null"); + } + return source; + } + + private Attempt requireAttempt(GraphVersion version) { + Attempt attempt = attempts.get(version); + if (attempt == null) { + throw new MetadataException(MetadataException.Code.INVALID_TRANSITION, + "ingestion attempt has not started"); + } + return attempt; + } + + private static final class Attempt { + private final GraphVersion expectedPublishedVersion; + + private Attempt(GraphVersion expectedPublishedVersion) { + this.expectedPublishedVersion = expectedPublishedVersion; + } + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/SourceStreamProvider.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/SourceStreamProvider.java new file mode 100644 index 000000000..f05228056 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/SourceStreamProvider.java @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.io.IOException; +import java.io.InputStream; + +/** Opens the exact source bytes used to verify an ingestion attempt. */ +@FunctionalInterface +public interface SourceStreamProvider { + + InputStream open(IngestionContext context) throws IOException; +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentityTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentityTest.java new file mode 100644 index 000000000..368d8ff83 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentityTest.java @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.io.IOException; +import java.nio.file.Path; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class ArtifactIdentityTest { + @TempDir + Path tempDir; + + @Test + void resolvesSafeVersionAndRejectsPathTraversal() throws Exception { + assertEquals(tempDir.resolve("graph-v1.2.bin"), ArtifactIdentity.resolveVersionedPath( + tempDir, "graph-", "v1.2", ".bin")); + assertThrows(IOException.class, () -> ArtifactIdentity.resolveVersionedPath( + tempDir, "graph-", "../outside", ".bin")); + assertThrows(IOException.class, () -> ArtifactIdentity.resolveVersionedPath( + tempDir, "graph-", "..", ".bin")); + assertThrows(IOException.class, () -> ArtifactIdentity.resolveVersionedPath( + tempDir, "graph-", "v1/../../outside", ".bin")); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java index 20d63557c..e0d356cc8 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilderTest.java @@ -20,6 +20,8 @@ package org.apache.geaflow.ai.retrieval.index.bm25; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import java.nio.file.Files; import java.nio.file.Path; @@ -29,6 +31,8 @@ import org.apache.geaflow.ai.retrieval.metadata.ChunkingConfiguration; import org.apache.geaflow.ai.retrieval.model.document.TextChunk; import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; import org.junit.jupiter.api.Test; class LuceneBm25IndexBuilderTest { @@ -47,7 +51,55 @@ void buildsReadableDeterministicArtifact() throws Exception { try (org.apache.geaflow.ai.retrieval.index.IndexArtifact artifact = builder.build(context, Arrays.asList(first, second))) { assertEquals("bm25", artifact.getMetadata().getIndexType()); - assertEquals(2, Files.list(directory.resolve("bm25-v1")).count()); + try (org.apache.lucene.store.Directory luceneDirectory = + org.apache.lucene.store.FSDirectory.open(directory.resolve("bm25-v1")); + org.apache.lucene.index.DirectoryReader reader = + org.apache.lucene.index.DirectoryReader.open(luceneDirectory)) { + assertEquals(2, reader.numDocs()); + } } } + + @Test + void readerReportsMissingArtifactAsIndexNotReady() { + RetrievalException exception = org.junit.jupiter.api.Assertions.assertThrows( + RetrievalException.class, + () -> new Bm25IndexReader(java.nio.file.Paths.get("/tmp/missing-bm25-artifact"))); + assertEquals(RetrievalErrorCode.INDEX_NOT_READY, exception.getCode()); + } + + @Test + void rejectsPathTraversalVersionBeforeCreatingOutputDirectory() { + Path directory = java.nio.file.Paths.get(System.getProperty("java.io.tmpdir"), + "bm25-unsafe-" + java.util.UUID.randomUUID()); + TextChunk chunk = new TextChunk("c1", "doc", 0, 0, 1, 1, "text"); + + assertThrows(java.io.IOException.class, + () -> new LuceneBm25IndexBuilder(directory).build(context("../outside"), Arrays.asList(chunk))); + assertFalse(Files.exists(directory)); + } + + @Test + void removesStagingDirectoryWhenPublishedArtifactValidationFails() throws Exception { + Path directory = Files.createTempDirectory("bm25-cleanup-test"); + Path published = Files.createDirectory(directory.resolve("bm25-v1")); + Files.write(published.resolve("invalid"), new byte[] {1}); + TextChunk chunk = new TextChunk("c1", "doc", 0, 0, 1, 1, "text"); + assertThrows(java.io.IOException.class, + () -> new LuceneBm25IndexBuilder(directory).build(context(), Arrays.asList(chunk))); + try (java.util.stream.Stream paths = Files.list(directory)) { + assertFalse(paths.anyMatch(path -> path.getFileName().toString().startsWith("bm25-v1-"))); + } + } + + private static IngestionContext context() { + return context("v1"); + } + + private static IngestionContext context(String graphVersion) { + DatasetManifest manifest = new DatasetManifest("v1", "d", "r", "dev", null, + "cache", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "p", new ChunkingConfiguration("chunk-v1", 100, 10), "graph-v1", "offline", "v1", 1L); + return new IngestionContext(manifest, new GraphVersion("g", graphVersion), "test"); + } } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactTest.java new file mode 100644 index 000000000..c04b176cd --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactTest.java @@ -0,0 +1,254 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index.graph; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.io.ByteArrayInputStream; +import java.io.DataInputStream; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.metadata.ChunkingConfiguration; +import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.apache.geaflow.ai.retrieval.service.RecallService; +import org.apache.geaflow.ai.retrieval.service.RetrievalFixtureRegistry; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class GraphArtifactTest { + private static final String GRAPH = "artifact-graph"; + private static final String VERSION = "v1"; + + @TempDir + Path tempDir; + + @Test + void builderAndReaderRoundTripDeterministicGraphAndRegistryRetrievesIt() throws Exception { + List chunks = chunks(); + List entities = Arrays.asList( + new EntityRef("e2", "Astronomy", "topic"), + new EntityRef("e1", "Confucius", Collections.singletonList("kongzi"), "person", + Collections.singletonList("c1"))); + List edges = Collections.singletonList( + new GraphEdgeRef("edge-1", "studies", "e1", "e2", Collections.singletonList("c2"))); + Path firstDirectory = tempDir.resolve("first"); + Path secondDirectory = tempDir.resolve("second"); + Path firstArtifact; + try (org.apache.geaflow.ai.retrieval.ingest.GraphArtifact ignored = + new OfflineGraphArtifactBuilder(firstDirectory).build(context(), chunks, entities, edges)) { + firstArtifact = firstDirectory.resolve("graph-v1.bin"); + } + try (org.apache.geaflow.ai.retrieval.ingest.GraphArtifact ignored = + new OfflineGraphArtifactBuilder(secondDirectory).build(context(), chunks, + Arrays.asList(entities.get(1), entities.get(0)), edges)) { + // A different input order must produce the same published bytes. + } + assertTrue(Arrays.equals(Files.readAllBytes(firstArtifact), + Files.readAllBytes(secondDirectory.resolve("graph-v1.bin")))); + + try (GraphArtifactReader reader = new GraphArtifactReader(firstArtifact, VERSION, VERSION)) { + reader.getIdentity().validate(GRAPH, VERSION, VERSION, chunks); + assertEquals(Arrays.asList("e1", "e2"), Arrays.asList( + reader.getEntities().get(0).getEntityId(), reader.getEntities().get(1).getEntityId())); + assertEquals(Collections.singletonList("kongzi"), reader.getEntities().get(0).getAliases()); + assertEquals(Collections.singletonList("c1"), reader.getEntities().get(0).getSourceChunkIds()); + assertEquals(edges, reader.getNeighbors().neighbors("e1")); + reader.validateReferences(chunks); + assertThrows(RetrievalException.class, () -> reader.getIdentity().validate( + GRAPH, VERSION, VERSION, Collections.singletonList( + new TextChunk("c1", "doc-1", 0, 0, 1, 1, "changed")))); + } + + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + try { + registry.loadGraph(GRAPH, VERSION, VERSION, firstArtifact, chunks); + RecallService service = new RecallService(registry); + org.apache.geaflow.ai.retrieval.api.model.RetrievalRequest request = + new org.apache.geaflow.ai.retrieval.api.model.RetrievalRequest(); + request.setGraphName(GRAPH); + request.setGraphVersion(VERSION); + request.setIndexVersion(VERSION); + request.setQuery("Confucius"); + request.setMode("GRAPH_ONLY"); + request.setBudget(new org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget(5, 1000, 5, 100)); + org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse response = service.retrieve(request); + assertEquals(1, response.getEvidence().size()); + assertEquals("c2", response.getEvidence().get(0).getChunks().get(0).getChunkId()); + assertEquals(VERSION, response.getTrace().getGraphVersion()); + } finally { + registry.close(); + } + } + + @Test + void rejectsVersionCorruptionAndInvalidBuildReferences() throws Exception { + List chunks = chunks(); + Path artifact = tempDir.resolve("graph-v1.bin"); + try (org.apache.geaflow.ai.retrieval.ingest.GraphArtifact ignored = + new OfflineGraphArtifactBuilder(tempDir).build(context(), chunks, entities(), edges())) { + // The graph artifact is published for reader validation below. + } + assertArtifactError(RetrievalErrorCode.INDEX_NOT_READY, + () -> new GraphArtifactReader(artifact, "v2", VERSION)); + + Path corrupt = tempDir.resolve("corrupt.bin"); + Files.write(corrupt, new byte[] {1, 2, 3}); + assertArtifactError(RetrievalErrorCode.INDEX_NOT_READY, + () -> new GraphArtifactReader(corrupt, VERSION, VERSION)); + + List unknownEndpoint = Collections.singletonList( + new GraphEdgeRef("bad-edge", "studies", "missing", "e2", Collections.singletonList("c2"))); + assertThrows(IOException.class, + () -> new OfflineGraphArtifactBuilder(tempDir.resolve("invalid-endpoint")) + .build(context(), chunks, entities(), unknownEndpoint)); + List unknownChunk = Collections.singletonList( + new GraphEdgeRef("bad-edge", "studies", "e1", "e2", Collections.singletonList("missing"))); + assertThrows(IOException.class, + () -> new OfflineGraphArtifactBuilder(tempDir.resolve("invalid-chunk")) + .build(context(), chunks, entities(), unknownChunk)); + List unknownEntityChunk = Collections.singletonList(new EntityRef( + "e1", "Confucius", Collections.emptyList(), "person", Collections.singletonList("missing"))); + assertThrows(IOException.class, + () -> new OfflineGraphArtifactBuilder(tempDir.resolve("invalid-entity-chunk")) + .build(context(), chunks, unknownEntityChunk, Collections.emptyList())); + } + + @Test + void rejectsPathTraversalVersionBeforeCreatingOutputDirectory() { + Path output = tempDir.resolve("unsafe-output"); + assertThrows(IOException.class, () -> new OfflineGraphArtifactBuilder(output).build( + context("../outside"), chunks(), entities(), edges())); + assertTrue(Files.notExists(output)); + } + + @Test + void rejectsCountsThatCannotFitInArtifactBeforeAllocating() throws Exception { + List chunks = chunks(); + Path artifact = tempDir.resolve("graph-v1.bin"); + try (org.apache.geaflow.ai.retrieval.ingest.GraphArtifact ignored = + new OfflineGraphArtifactBuilder(tempDir).build(context(), chunks, entities(), edges())) { + // Build a valid artifact to mutate its count fields below. + } + byte[] original = Files.readAllBytes(artifact); + int[] offsets = graphCountOffsets(original); + + byte[] invalidEntityCount = original.clone(); + ByteBuffer.wrap(invalidEntityCount).putInt(offsets[0], Integer.MAX_VALUE); + Path invalidEntityArtifact = tempDir.resolve("invalid-entity-count.bin"); + Files.write(invalidEntityArtifact, invalidEntityCount); + assertArtifactError(RetrievalErrorCode.INDEX_NOT_READY, + () -> new GraphArtifactReader(invalidEntityArtifact, VERSION, VERSION)); + + byte[] invalidAliasCount = original.clone(); + ByteBuffer.wrap(invalidAliasCount).putInt(offsets[1], Integer.MAX_VALUE); + Path invalidAliasArtifact = tempDir.resolve("invalid-alias-count.bin"); + Files.write(invalidAliasArtifact, invalidAliasCount); + assertArtifactError(RetrievalErrorCode.INDEX_NOT_READY, + () -> new GraphArtifactReader(invalidAliasArtifact, VERSION, VERSION)); + } + + @Test + void reusesIdenticalPublishedArtifactAndRejectsDifferentContent() throws Exception { + List chunks = chunks(); + OfflineGraphArtifactBuilder builder = new OfflineGraphArtifactBuilder(tempDir); + try (org.apache.geaflow.ai.retrieval.ingest.GraphArtifact ignored = + builder.build(context(), chunks, entities(), edges())) { + // Publish initial content. + } + Path published = tempDir.resolve("graph-v1.bin"); + byte[] original = Files.readAllBytes(published); + try (org.apache.geaflow.ai.retrieval.ingest.GraphArtifact ignored = + builder.build(context(), chunks, entities(), edges())) { + // Identical content is accepted as an idempotent rebuild. + } + assertTrue(Arrays.equals(original, Files.readAllBytes(published))); + + List changed = Collections.singletonList( + new GraphEdgeRef("edge-2", "teaches", "e1", "e2", Collections.singletonList("c2"))); + assertThrows(IOException.class, () -> builder.build(context(), chunks, entities(), changed)); + assertNotEquals(0, Files.size(published)); + } + + private static void assertArtifactError(RetrievalErrorCode code, + org.junit.jupiter.api.function.Executable operation) { + RetrievalException error = assertThrows(RetrievalException.class, operation); + assertEquals(code, error.getCode()); + } + + private static int[] graphCountOffsets(byte[] artifact) throws IOException { + try (DataInputStream input = new DataInputStream(new ByteArrayInputStream(artifact))) { + for (int index = 0; index < 5; index++) { + input.readUTF(); + } + int entityCountOffset = artifact.length - input.available(); + input.readInt(); + input.readInt(); + input.readUTF(); + input.readUTF(); + input.readUTF(); + int aliasCountOffset = artifact.length - input.available(); + return new int[] {entityCountOffset, aliasCountOffset}; + } + } + + private static List chunks() { + return Arrays.asList( + new TextChunk("c1", "doc-1", 0, 0, 10, 2, "Confucius"), + new TextChunk("c2", "doc-2", 0, 0, 9, 2, "philosophy")); + } + + private static List entities() { + return Arrays.asList( + new EntityRef("e1", "Confucius", Collections.singletonList("kongzi"), "person", + Collections.singletonList("c1")), + new EntityRef("e2", "Astronomy", "topic")); + } + + private static List edges() { + return Collections.singletonList( + new GraphEdgeRef("edge-1", "studies", "e1", "e2", Collections.singletonList("c2"))); + } + + private IngestionContext context() { + return context(VERSION); + } + + private IngestionContext context(String graphVersion) { + DatasetManifest manifest = new DatasetManifest(VERSION, "dataset", "revision", "dev", null, + "cache", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "policy", new ChunkingConfiguration("chunk-v1", 100, 10), "graph", "fixture", VERSION, 7L); + return new IngestionContext(manifest, new GraphVersion(GRAPH, graphVersion), "test"); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java index ebefa6594..be6cb00ae 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilderTest.java @@ -19,11 +19,14 @@ package org.apache.geaflow.ai.retrieval.index.vector; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; +import java.io.DataOutputStream; import java.nio.file.Files; import java.nio.file.Path; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.Map; import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; @@ -31,6 +34,8 @@ import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; import org.apache.geaflow.ai.retrieval.model.document.TextChunk; import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; import org.junit.jupiter.api.Test; class OfflineVectorIndexBuilderTest { @@ -49,10 +54,72 @@ void rejectsNonFiniteAndMissingVectors() throws Exception { assertThrows(java.io.IOException.class, () -> invalidBuilder.build(context, Arrays.asList(chunk))); } + @Test + void rejectsPathTraversalVersionBeforeCreatingOutputDirectory() { + Path directory = tempDir.resolve("unsafe-vector"); + TextChunk chunk = new TextChunk("c1", "d", 0, 0, 1, 1, "x"); + Map vectors = Collections.singletonMap("c1", new float[] {1.0F}); + OfflineVectorIndexBuilder builder = new OfflineVectorIndexBuilder(directory, vectors, "fixture", "v1"); + + assertThrows(java.io.IOException.class, + () -> builder.build(context("../outside"), Collections.singletonList(chunk))); + assertFalse(Files.exists(directory)); + } + + @Test + void readerRejectsCorruptArtifact() throws Exception { + Path artifact = Files.createTempFile("vector-corrupt", ".bin"); + try (DataOutputStream output = new DataOutputStream(Files.newOutputStream(artifact))) { + output.writeUTF("bad-magic"); + } + RetrievalException corrupt = assertThrows(RetrievalException.class, + () -> new VectorIndexReader(artifact)); + org.junit.jupiter.api.Assertions.assertEquals(RetrievalErrorCode.INDEX_NOT_READY, + corrupt.getCode()); + } + + @Test + void rejectsZeroNormAndChangedExistingVectors() throws Exception { + Path directory = Files.createTempDirectory("vector-existing"); + TextChunk chunk = new TextChunk("chunk", "doc", 0, 0, 4, 1, "text"); + Map vectors = new HashMap<>(); + vectors.put("chunk", new float[] {0.0F, 0.0F}); + assertThrows(java.io.IOException.class, () -> new OfflineVectorIndexBuilder(directory, vectors, "fixture", "v1") + .build(context(), Collections.singletonList(chunk))); + vectors.put("chunk", new float[] {1.0F, 0.0F}); + new OfflineVectorIndexBuilder(directory, vectors, "fixture", "v1").build(context(), Collections.singletonList(chunk)).close(); + vectors.put("chunk", new float[] {0.0F, 1.0F}); + assertThrows(java.io.IOException.class, () -> new OfflineVectorIndexBuilder(directory, vectors, "fixture", "v1") + .build(context(), Collections.singletonList(chunk))); + } + + @Test + void removesStagingFileWhenExistingArtifactValidationFails() throws Exception { + Path directory = Files.createTempDirectory("vector-cleanup"); + TextChunk chunk = new TextChunk("chunk", "doc", 0, 0, 4, 1, "text"); + Map vectors = new HashMap<>(); + vectors.put("chunk", new float[] {1.0F, 0.0F}); + OfflineVectorIndexBuilder builder = new OfflineVectorIndexBuilder(directory, vectors, "fixture", "v1"); + builder.build(context(), Collections.singletonList(chunk)).close(); + vectors.put("chunk", new float[] {0.0F, 1.0F}); + assertThrows(java.io.IOException.class, + () -> builder.build(context(), Collections.singletonList(chunk))); + try (java.util.stream.Stream paths = Files.list(directory)) { + assertFalse(paths.anyMatch(path -> path.getFileName().toString().endsWith(".tmp"))); + } + } + private static IngestionContext context() { + return context("v1"); + } + + private static IngestionContext context(String graphVersion) { DatasetManifest manifest = new DatasetManifest("v1", "d", "r", "dev", null, "cache", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", "p", new ChunkingConfiguration("chunk-v1", 100, 10), "graph-v1", "fixture", "v1", 1L); - return new IngestionContext(manifest, new GraphVersion("g", "v1"), "test"); + return new IngestionContext(manifest, new GraphVersion("g", graphVersion), "test"); } + + @org.junit.jupiter.api.io.TempDir + Path tempDir; } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java index 4ef738e2c..bdd4dfcdd 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java @@ -24,6 +24,7 @@ import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertTrue; +import java.io.ByteArrayInputStream; import java.util.Arrays; import java.util.Collections; import java.util.List; @@ -35,9 +36,11 @@ import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; import org.apache.geaflow.ai.retrieval.metadata.GraphBuildMetadata; import org.apache.geaflow.ai.retrieval.metadata.ImportMetadata; +import org.apache.geaflow.ai.retrieval.metadata.ImportState; +import org.apache.geaflow.ai.retrieval.metadata.IndexBuildMetadata; +import org.apache.geaflow.ai.retrieval.metadata.InMemoryMetadataStore; import org.apache.geaflow.ai.retrieval.metadata.MetadataException; import org.apache.geaflow.ai.retrieval.metadata.QualityCounters; -import org.apache.geaflow.ai.retrieval.metadata.IndexBuildMetadata; import org.apache.geaflow.ai.retrieval.model.document.SourceDocument; import org.apache.geaflow.ai.retrieval.model.document.TextChunk; import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; @@ -161,6 +164,40 @@ public void fail(IngestionContext value, Exception failure) { assertTrue(bm25.closed); } + @Test + public void metadataStorePublisherCompletesReadyLifecycle() throws Exception { + String source = "source"; + String checksum = "41cf6794ba4200b839c53531555f0f3998df4cbb01a4d5cb0b94e3ca5e23947d"; + GraphVersion version = new GraphVersion("graph", "g1"); + DatasetManifest manifest = new DatasetManifest("v1", "dataset", "release", "dev", "uri", null, + checksum, "preprocess-v1", new ChunkingConfiguration("chunk-v1", 100, 10), + "schema-v1", "model", "v1", 7); + IngestionContext context = new IngestionContext(manifest, version, "importer-v1"); + InMemoryMetadataStore store = new InMemoryMetadataStore(); + MetadataStorePublisher publisher = new MetadataStorePublisher(store, + value -> new ByteArrayInputStream(source.getBytes(java.nio.charset.StandardCharsets.UTF_8))); + TrackingGraphArtifact graph = new TrackingGraphArtifact(new java.util.ArrayList<>(), "graph", + version, Arrays.asList("bm25", "vector")); + TrackingIndexArtifact bm25 = new TrackingIndexArtifact(new java.util.ArrayList<>(), "bm25", version); + TrackingIndexArtifact vector = new TrackingIndexArtifact(new java.util.ArrayList<>(), "vector", version); + IngestionPipeline pipeline = new IngestionPipeline( + value -> Collections.singletonList(document()), + (documents, value) -> Collections.singletonList(chunk()), + (chunks, value) -> new ExtractionResult(Collections.singletonList(entity()), + Collections.singletonList(edge())), + (value, documents, chunks, extraction) -> graph, + (value, chunks) -> bm25, + (value, chunks) -> vector, + publisher); + + ImportMetadata result = pipeline.run(context); + + assertEquals(ImportState.READY, result.getState()); + assertEquals(version, store.getPublishedVersion("graph").get()); + assertEquals(Arrays.asList("bm25", "vector"), result.getGraph().getRequiredIndexes()); + assertEquals(2, result.getIndexes().size()); + } + private static IngestionContext context() { DatasetManifest manifest = new DatasetManifest("v1", "hotpotqa", "release-1", "train", null, "cache/hotpotqa", SHA256, "normalizer-1", @@ -187,17 +224,26 @@ private static GraphEdgeRef edge() { private static final class TrackingGraphArtifact implements GraphArtifact { private final List calls; private final String name; + private final GraphVersion version; + private final List requiredIndexes; private boolean closed; private TrackingGraphArtifact(List calls, String name) { + this(calls, name, new GraphVersion("graph-1", "version-1"), Collections.emptyList()); + } + + private TrackingGraphArtifact(List calls, String name, GraphVersion version, + List requiredIndexes) { this.calls = calls; this.name = name; + this.version = version; + this.requiredIndexes = requiredIndexes; } @Override public GraphBuildMetadata getMetadata() { - return new GraphBuildMetadata(new GraphVersion("graph-1", "version-1"), - "memory:" + name, "fixture", "builder-1", true); + return new GraphBuildMetadata(version, "memory:" + name, "fixture", "builder-1", true, + requiredIndexes); } @Override @@ -210,17 +256,23 @@ public void close() { private static final class TrackingIndexArtifact implements IndexArtifact { private final List calls; private final String name; + private final GraphVersion version; private boolean closed; private TrackingIndexArtifact(List calls, String name) { + this(calls, name, new GraphVersion("graph-1", "version-1")); + } + + private TrackingIndexArtifact(List calls, String name, GraphVersion version) { this.calls = calls; this.name = name; + this.version = version; } @Override public IndexBuildMetadata getMetadata() { - return new IndexBuildMetadata(new GraphVersion("graph-1", "version-1"), - new org.apache.geaflow.ai.retrieval.model.version.IndexVersion(name, "v1", "version-1"), + return new IndexBuildMetadata(version, + new org.apache.geaflow.ai.retrieval.model.version.IndexVersion(name, "v1", version.getVersion()), name, "builder-1", "memory:" + name, true); } diff --git a/tools/graphrag/ingest/entity_graph.py b/tools/graphrag/ingest/entity_graph.py index 22e2a24b2..5cfcaf9fd 100644 --- a/tools/graphrag/ingest/entity_graph.py +++ b/tools/graphrag/ingest/entity_graph.py @@ -5,6 +5,7 @@ import hashlib import json import os +import re import tempfile import unicodedata from pathlib import Path @@ -138,8 +139,13 @@ def write_graph_artifact(directory: Path, graph_version: str, vertices: Iterable[Dict[str, Any]], edges: Iterable[Dict[str, Any]]) -> Path: """Write immutable vertex/edge JSONL and publish a manifest last.""" + if not isinstance(graph_version, str) or not re.fullmatch(r"[A-Za-z0-9._-]{1,128}", graph_version): + raise ValueError("graph_version contains unsafe path characters") directory.mkdir(parents=True, exist_ok=True) + root = directory.resolve() target = directory / graph_version + if target.resolve().parent != root: + raise ValueError("graph_version must remain within output directory") if target.exists(): raise FileExistsError("graph version already exists: %s" % graph_version) temporary_directory = Path(tempfile.mkdtemp(prefix=graph_version + ".", dir=directory)) diff --git a/tools/graphrag/ingest/sources/source_loader.py b/tools/graphrag/ingest/sources/source_loader.py index 9ebb31b0f..4f70ed278 100644 --- a/tools/graphrag/ingest/sources/source_loader.py +++ b/tools/graphrag/ingest/sources/source_loader.py @@ -67,17 +67,29 @@ def from_file(cls, path: Path) -> "SourceManifest": class SourceLoader: """Loads and validates one manifest without exposing dataset-specific records.""" - def __init__(self, allow_download: bool = False, batch_size: int = 256): + def __init__(self, allow_download: bool = False, batch_size: int = 256, + max_download_bytes: int = 1024 * 1024 * 1024, + max_uncompressed_bytes: int = 1024 * 1024 * 1024, + max_zip_entries: int = 16, timeout: float = 30.0): if batch_size < 1: raise ValueError("batch_size must be positive") + if min(max_download_bytes, max_uncompressed_bytes, max_zip_entries) < 1 or timeout <= 0: + raise ValueError("source limits must be positive") self.allow_download = allow_download self.batch_size = batch_size + self.max_download_bytes = max_download_bytes + self.max_uncompressed_bytes = max_uncompressed_bytes + self.max_zip_entries = max_zip_entries + self.timeout = timeout def load_batches(self, manifest: SourceManifest) -> Iterator[List[Dict[str, Any]]]: path = self._resolve_source(manifest) self._verify_checksum(path, manifest) batch: List[Dict[str, Any]] = [] - for number, raw in enumerate(self._records(path), start=1): + if path.stat().st_size > self.max_download_bytes: + raise SourceError("source exceeds configured size limit: %s" % path) + for number, raw in enumerate(self._records(path, self.max_uncompressed_bytes, + self.max_zip_entries), start=1): batch.append(self._canonical_record(manifest, number, raw)) if len(batch) == self.batch_size: yield batch @@ -107,7 +119,17 @@ def _resolve_source(self, manifest: SourceManifest) -> Path: os.close(descriptor) temporary = Path(temporary_name) try: - urllib.request.urlretrieve(manifest.source_uri, temporary) + request = urllib.request.Request(manifest.source_uri) + total = 0 + with urllib.request.urlopen(request, timeout=self.timeout) as response, temporary.open("wb") as output: + while True: + block = response.read(1024 * 1024) + if not block: + break + total += len(block) + if total > self.max_download_bytes: + raise SourceError("download exceeds configured size limit") + output.write(block) self._verify_checksum(temporary, manifest) temporary.replace(target) except (OSError, SourceError) as exc: @@ -138,27 +160,66 @@ def _verify_checksum(path: Path, manifest: SourceManifest) -> None: ) @staticmethod - def _records(path: Path) -> Iterable[Any]: + def _records(path: Path, max_uncompressed_bytes: int = 1024 * 1024 * 1024, + max_zip_entries: int = 16) -> Iterable[Any]: if path.suffix == ".zip": try: with zipfile.ZipFile(path) as archive: names = sorted(name for name in archive.namelist() if name.endswith((".json", ".jsonl", ".json.gz", ".jsonl.gz"))) + if len(archive.namelist()) > max_zip_entries: + raise SourceError("ZIP source %s has too many entries" % path) if len(names) != 1: raise SourceError("ZIP source %s must contain exactly one JSON data file" % path) descriptor, temporary_name = tempfile.mkstemp(suffix=Path(names[0]).suffix) os.close(descriptor) extracted = Path(temporary_name) try: + total = 0 with archive.open(names[0]) as source, extracted.open("wb") as target: - target.write(source.read()) - yield from SourceLoader._records(extracted) + while True: + block = source.read(1024 * 1024) + if not block: + break + total += len(block) + if total > max_uncompressed_bytes: + raise SourceError("ZIP entry exceeds configured size limit") + target.write(block) + yield from SourceLoader._records(extracted, max_uncompressed_bytes, max_zip_entries) finally: extracted.unlink(missing_ok=True) return except (OSError, zipfile.BadZipFile) as exc: raise SourceError("cannot parse ZIP source %s: %s" % (path, exc)) from exc - opener = gzip.open if path.suffix == ".gz" else open + if path.suffix == ".gz": + descriptor, temporary_name = tempfile.mkstemp(suffix=".jsonl") + os.close(descriptor) + extracted = Path(temporary_name) + try: + total = 0 + with gzip.open(path, "rb") as source, extracted.open("wb") as target: + while True: + block = source.read(1024 * 1024) + if not block: + break + total += len(block) + if total > max_uncompressed_bytes: + raise SourceError("gzip source exceeds configured size limit") + target.write(block) + yield from SourceLoader._records(extracted, max_uncompressed_bytes, max_zip_entries) + except SourceError: + raise + except (OSError, gzip.BadGzipFile) as exc: + raise SourceError("cannot parse gzip source %s: %s" % (path, exc)) from exc + finally: + extracted.unlink(missing_ok=True) + return + try: + if path.stat().st_size > max_uncompressed_bytes: + raise SourceError("source exceeds configured uncompressed size limit: %s" % path) + except OSError as exc: + raise SourceError("cannot stat source %s: %s" % (path, exc)) from exc + opener = open try: with opener(path, "rt", encoding="utf-8") as stream: first = stream.read(1) diff --git a/tools/graphrag/tests/test_source_loader.py b/tools/graphrag/tests/test_source_loader.py index 3be423389..6c28a7683 100644 --- a/tools/graphrag/tests/test_source_loader.py +++ b/tools/graphrag/tests/test_source_loader.py @@ -1,4 +1,5 @@ import hashlib +import gzip import json import tempfile import unittest @@ -62,6 +63,27 @@ def test_rejects_trailing_comma_and_trailing_data(self): with self.assertRaises(SourceError): list(SourceLoader().load(manifest)) + def test_limits_uncompressed_plain_source(self): + with tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "records.jsonl" + source.write_text(json.dumps({"id": "a", "question": "q", "context": []}), + encoding="utf-8") + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + manifest = SourceManifest("d", "r", "s", None, str(source), checksum, "p") + with self.assertRaisesRegex(SourceError, "uncompressed size limit"): + list(SourceLoader(max_uncompressed_bytes=1).load(manifest)) + + def test_limits_uncompressed_gzip_source(self): + with tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "records.jsonl.gz" + contents = json.dumps({"id": "a", "question": "q", "context": []}) + with gzip.open(source, "wt", encoding="utf-8") as stream: + stream.write(contents) + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + manifest = SourceManifest("d", "r", "s", None, str(source), checksum, "p") + with self.assertRaisesRegex(SourceError, "gzip source exceeds"): + list(SourceLoader(max_uncompressed_bytes=len(contents) - 1).load(manifest)) + if __name__ == "__main__": unittest.main() From d9394f874b470af2cf82f1cd3c96c6d3222e70b9 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Thu, 8 Oct 2026 21:03:22 +0800 Subject: [PATCH 22/30] feat(ai): add budgeted hybrid recall with graph channel --- .../retrieval/channel/bm25/Bm25Retriever.java | 38 + .../channel/graph/EntityAnchorResolver.java | 103 +++ .../channel/graph/GraphNeighborProvider.java | 30 + .../channel/graph/GraphRetriever.java | 286 +++++++ .../graph/InMemoryGraphNeighborProvider.java | 60 ++ .../channel/vector/VectorRetriever.java | 38 + .../retrieval/config/RetrievalProperties.java | 8 +- .../ai/retrieval/execution/RecallPlan.java | 95 +++ .../execution/RecallStageStatus.java | 24 + .../retrieval/execution/RecallStopReason.java | 25 + .../retrieval/fusion/WeightedRrfConfig.java | 64 ++ .../retrieval/fusion/WeightedRrfFusion.java | 70 ++ .../retrieval/index/ChannelSearchResult.java | 57 ++ .../ai/retrieval/index/bm25/Bm25Hit.java | 58 ++ .../retrieval/index/bm25/Bm25IndexReader.java | 211 +++++ .../ai/retrieval/index/vector/VectorHit.java | 51 ++ .../index/vector/VectorIndexReader.java | 232 ++++++ .../ai/retrieval/model/evidence/Evidence.java | 23 +- .../model/evidence/GraphEvidenceTrace.java | 122 +++ .../ai/retrieval/service/RecallService.java | 496 ++++++++++++ .../service/RetrievalFixtureRegistry.java | 408 ++++++++++ .../graph/EntityAnchorResolverTest.java | 69 ++ .../channel/graph/GraphRetrieverTest.java | 122 +++ .../fusion/WeightedRrfFusionTest.java | 97 +++ .../service/GraphRagMvpEndToEndTest.java | 752 ++++++++++++++++++ 25 files changed, 3534 insertions(+), 5 deletions(-) create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/bm25/Bm25Retriever.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/EntityAnchorResolver.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphNeighborProvider.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphRetriever.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/InMemoryGraphNeighborProvider.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/vector/VectorRetriever.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallPlan.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallStageStatus.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallStopReason.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfConfig.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfFusion.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ChannelSearchResult.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/Bm25Hit.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/Bm25IndexReader.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/VectorHit.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/VectorIndexReader.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/GraphEvidenceTrace.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RecallService.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalFixtureRegistry.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/channel/graph/EntityAnchorResolverTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphRetrieverTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfFusionTest.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/GraphRagMvpEndToEndTest.java diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/bm25/Bm25Retriever.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/bm25/Bm25Retriever.java new file mode 100644 index 000000000..62b595cc4 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/bm25/Bm25Retriever.java @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.channel.bm25; + +import java.util.List; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.index.bm25.Bm25Hit; +import org.apache.geaflow.ai.retrieval.index.bm25.Bm25IndexReader; + +/** Session-free BM25 channel backed by a reusable artifact reader. */ +public final class Bm25Retriever { + private final Bm25IndexReader reader; + + public Bm25Retriever(Bm25IndexReader reader) { + this.reader = Objects.requireNonNull(reader, "reader"); + } + + public List retrieve(String query, int topK, int candidates, long deadlineNanos) { + return reader.search(query, candidates, topK, deadlineNanos); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/EntityAnchorResolver.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/EntityAnchorResolver.java new file mode 100644 index 000000000..13e85d2c7 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/EntityAnchorResolver.java @@ -0,0 +1,103 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.channel.graph; + +import java.text.Normalizer; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Locale; +import java.util.Objects; +import java.util.function.LongSupplier; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; + +/** Deterministic canonical-name and alias resolver used as the graph expansion boundary. */ +public final class EntityAnchorResolver { + public List resolve(String query, List entities) { + return resolve(query, entities, 0L, System::nanoTime); + } + + public List resolve(String query, List entities, long deadlineNanos, LongSupplier clock) { + if (query == null || entities == null) { + return Collections.emptyList(); + } + String normalized = normalize(query); + if (normalized.isEmpty()) { + return Collections.emptyList(); + } + List result = new ArrayList<>(); + for (EntityRef entity : entities) { + if (deadlineNanos != 0L && clock.getAsLong() - deadlineNanos >= 0L) { + break; + } + if (entity == null) { + continue; + } + String canonical = normalize(entity.getCanonicalName()); + if (normalized.equals(canonical)) { + result.add(new Anchor(entity, "CANONICAL", 1.0)); + } else { + for (String alias : entity.getAliases()) { + if (normalized.equals(normalize(alias))) { + result.add(new Anchor(entity, "ALIAS", 0.9)); + break; + } + } + } + } + result.sort(Comparator.comparing(Anchor::getConfidence).reversed() + .thenComparing(a -> a.getEntity().getEntityId())); + return result; + } + + public static String normalize(String value) { + if (value == null) { + return ""; + } + String normalized = Normalizer.normalize(value, Normalizer.Form.NFKC); + return normalized.trim().toLowerCase(Locale.ROOT).replaceAll("\\s+", " "); + } + + /** An entity matched to a query with deterministic confidence metadata. */ + public static final class Anchor { + private final EntityRef entity; + private final String matchType; + private final double confidence; + + public Anchor(EntityRef entity, String matchType, double confidence) { + this.entity = Objects.requireNonNull(entity, "entity"); + this.matchType = Objects.requireNonNull(matchType, "matchType"); + this.confidence = confidence; + } + + public EntityRef getEntity() { + return entity; + } + + public String getMatchType() { + return matchType; + } + + public double getConfidence() { + return confidence; + } + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphNeighborProvider.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphNeighborProvider.java new file mode 100644 index 000000000..08ddb0377 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphNeighborProvider.java @@ -0,0 +1,30 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.channel.graph; + +import java.util.List; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; + +/** Versioned graph boundary used by the bounded graph retriever. */ +public interface GraphNeighborProvider { + + /** Returns non-null neighbors ordered by edge ID, source, target, and label. */ + List neighbors(String entityId); +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphRetriever.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphRetriever.java new file mode 100644 index 000000000..9ca89b59b --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphRetriever.java @@ -0,0 +1,286 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.channel.graph; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.LongSupplier; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphPathRef; + +/** Bounded deterministic one-hop graph expansion over fixture graph records. */ +public final class GraphRetriever { + private final EntityAnchorResolver resolver; + private GraphStats lastStats = new GraphStats(); + + public GraphRetriever() { + this(new EntityAnchorResolver()); + } + + public GraphRetriever(EntityAnchorResolver resolver) { + this.resolver = resolver; + } + + public List retrieve(String query, List entities, List edges, + Map chunks, int topK, int edgeCap, + long deadlineNanos) { + return retrieve(query, entities, new InMemoryGraphNeighborProvider(edges), chunks, + topK, edgeCap, deadlineNanos, System::nanoTime); + } + + public List retrieve(String query, List entities, List edges, + Map chunks, int candidateCap, int edgeCap, + long deadlineNanos, LongSupplier clock) { + return retrieve(query, entities, new InMemoryGraphNeighborProvider(edges), chunks, + candidateCap, edgeCap, deadlineNanos, clock); + } + + public List retrieve(String query, List entities, GraphNeighborProvider provider, + Map chunks, int candidateCap, int edgeCap, + long deadlineNanos, LongSupplier clock) { + lastStats = new GraphStats(); + if (candidateCap <= 0 || edgeCap <= 0 || entities == null || provider == null || chunks == null) { + return Collections.emptyList(); + } + if (expired(deadlineNanos, clock)) { + return Collections.emptyList(); + } + List anchors = resolver.resolve(query, entities, deadlineNanos, clock); + if (expired(deadlineNanos, clock)) { + return Collections.emptyList(); + } + lastStats.anchorsConsidered = anchors.size(); + if (anchors.isEmpty()) { + lastStats.stopReason = RecallStopReason.NO_RELIABLE_ANCHOR; + return Collections.emptyList(); + } + Map entityMap = new HashMap<>(); + for (EntityRef entity : entities) { + if (entity != null) { + entityMap.put(entity.getEntityId(), entity); + } + } + List result = new ArrayList<>(); + Set seenEdges = new HashSet<>(); + Set producedChunks = new HashSet<>(); + int provenanceIdsExamined = 0; + for (EntityAnchorResolver.Anchor anchor : anchors) { + if (expired(deadlineNanos, clock)) { + return result; + } + List neighbors = provider.neighbors(anchor.getEntity().getEntityId()); + if (neighbors == null) { + continue; + } + java.util.Iterator iterator = neighbors.iterator(); + while (iterator.hasNext()) { + if (expired(deadlineNanos, clock)) { + return result; + } + if (lastStats.edgesExamined >= edgeCap) { + lastStats.stopReason = RecallStopReason.EDGE_SCAN_LIMIT; + return result; + } + GraphEdgeRef edge = iterator.next(); + lastStats.neighborsSampled++; + lastStats.edgesExamined++; + if (edge == null) { + continue; + } + if (!anchor.getEntity().getEntityId().equals(edge.getSourceEntityId()) + && !anchor.getEntity().getEntityId().equals(edge.getTargetEntityId())) { + continue; + } + if (!seenEdges.add(edge.getEdgeId())) { + continue; + } + String target = anchor.getEntity().getEntityId().equals(edge.getSourceEntityId()) + ? edge.getTargetEntityId() : edge.getSourceEntityId(); + if (!entityMap.containsKey(target)) { + continue; + } + List sourceIds = edge.getSourceChunkIds(); + if (sourceIds.isEmpty()) { + sourceIds = anchor.getEntity().getSourceChunkIds(); + } + if (sourceIds.isEmpty()) { + lastStats.validationErrors.add("edge " + edge.getEdgeId() + " references an unknown chunk"); + continue; + } + List boundedSources = new ArrayList<>(); + Set edgeSources = new HashSet<>(); + Set newCandidates = new HashSet<>(); + boolean invalidSource = false; + boolean provenanceLimitReached = false; + for (String id : sourceIds) { + if (expired(deadlineNanos, clock)) { + return result; + } + if (provenanceIdsExamined >= candidateCap) { + provenanceLimitReached = true; + break; + } + provenanceIdsExamined++; + if (!chunks.containsKey(id)) { + invalidSource = true; + break; + } + if (edgeSources.add(id)) { + boundedSources.add(id); + } + if (!producedChunks.contains(id)) { + newCandidates.add(id); + } + if (producedChunks.size() + newCandidates.size() >= candidateCap) { + break; + } + } + if (invalidSource) { + lastStats.validationErrors.add("edge " + edge.getEdgeId() + " references an unknown chunk"); + continue; + } + if (boundedSources.isEmpty()) { + if (provenanceLimitReached) { + lastStats.stopReason = RecallStopReason.CANDIDATE_LIMIT; + return result; + } + continue; + } + producedChunks.addAll(newCandidates); + GraphPathRef path = new GraphPathRef( + java.util.Arrays.asList(anchor.getEntity().getEntityId(), target), + Collections.singletonList(edge.getEdgeId()), 1, false); + result.add(new GraphHit(anchor, edge, path, boundedSources, + 1.0 / (1.0 + result.size()))); + lastStats.verticesReached++; + lastStats.candidatesProduced = producedChunks.size(); + if (producedChunks.size() >= candidateCap || provenanceLimitReached) { + lastStats.stopReason = RecallStopReason.CANDIDATE_LIMIT; + return result; + } + } + } + if (!lastStats.validationErrors.isEmpty() && lastStats.stopReason == RecallStopReason.COMPLETED) { + lastStats.stopReason = RecallStopReason.TRACE_VALIDATION_ERROR; + } + return result; + } + + public GraphStats getLastStats() { + return lastStats; + } + + private boolean expired(long deadline, LongSupplier clock) { + if (deadline != 0L && clock.getAsLong() - deadline >= 0L) { + lastStats.stopReason = RecallStopReason.DEADLINE; + return true; + } + return false; + } + + /** Graph candidate with path and source provenance. */ + public static final class GraphHit { + private final EntityAnchorResolver.Anchor anchor; + private final GraphEdgeRef edge; + private final GraphPathRef path; + private final List sourceChunkIds; + private final double score; + + private GraphHit(EntityAnchorResolver.Anchor anchor, GraphEdgeRef edge, GraphPathRef path, + List sourceChunkIds, double score) { + this.anchor = anchor; + this.edge = edge; + this.path = path; + this.sourceChunkIds = new ArrayList<>(sourceChunkIds); + this.score = score; + } + + public EntityAnchorResolver.Anchor getAnchor() { + return anchor; + } + + public GraphEdgeRef getEdge() { + return edge; + } + + public GraphPathRef getPath() { + return path; + } + + public List getSourceChunkIds() { + return Collections.unmodifiableList(sourceChunkIds); + } + + public double getScore() { + return score; + } + } + + /** Execution counters for trace construction. */ + public static final class GraphStats { + private int anchorsConsidered; + private int edgesExamined; + private int neighborsSampled; + private int verticesReached; + private int candidatesProduced; + private RecallStopReason stopReason = RecallStopReason.COMPLETED; + private final List validationErrors = new ArrayList<>(); + + public int getAnchorsConsidered() { + return anchorsConsidered; + } + + public int getEdgesExamined() { + return edgesExamined; + } + + public int getNeighborsSampled() { + return neighborsSampled; + } + + public int getVerticesReached() { + return verticesReached; + } + + public int getCandidatesProduced() { + return candidatesProduced; + } + + public RecallStopReason getStopReason() { + return stopReason; + } + + public List getValidationErrors() { + return Collections.unmodifiableList(validationErrors); + } + + public boolean isDeadlineReached() { + return stopReason == RecallStopReason.DEADLINE; + } + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/InMemoryGraphNeighborProvider.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/InMemoryGraphNeighborProvider.java new file mode 100644 index 000000000..ccb0d83ab --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/graph/InMemoryGraphNeighborProvider.java @@ -0,0 +1,60 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.channel.graph; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; + +/** Deterministic in-memory graph provider used by fixtures and artifact readers. */ +public final class InMemoryGraphNeighborProvider implements GraphNeighborProvider { + private final Map> byEntity = new HashMap<>(); + + public InMemoryGraphNeighborProvider(List edges) { + if (edges != null) { + for (GraphEdgeRef edge : edges) { + if (edge != null) { + add(edge.getSourceEntityId(), edge); + add(edge.getTargetEntityId(), edge); + } + } + } + for (List values : byEntity.values()) { + values.sort(Comparator.comparing(GraphEdgeRef::getEdgeId) + .thenComparing(GraphEdgeRef::getSourceEntityId) + .thenComparing(GraphEdgeRef::getTargetEntityId) + .thenComparing(GraphEdgeRef::getLabel)); + } + } + + private void add(String entityId, GraphEdgeRef edge) { + byEntity.computeIfAbsent(entityId, ignored -> new ArrayList<>()).add(edge); + } + + @Override + public List neighbors(String entityId) { + List result = byEntity.get(entityId); + return result == null ? Collections.emptyList() : Collections.unmodifiableList(result); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/vector/VectorRetriever.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/vector/VectorRetriever.java new file mode 100644 index 000000000..f2afe7e70 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/channel/vector/VectorRetriever.java @@ -0,0 +1,38 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.channel.vector; + +import java.util.List; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.index.vector.VectorHit; +import org.apache.geaflow.ai.retrieval.index.vector.VectorIndexReader; + +/** Session-free vector channel using caller supplied query vectors. */ +public final class VectorRetriever { + private final VectorIndexReader reader; + + public VectorRetriever(VectorIndexReader reader) { + this.reader = Objects.requireNonNull(reader, "reader"); + } + + public List retrieve(float[] query, int topK, int candidates, long deadlineNanos) { + return reader.search(query, candidates, topK, deadlineNanos); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java index c2fb3fbda..f00b38079 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java @@ -61,8 +61,12 @@ public void validateConfiguration() { org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.INVALID_REQUEST, "retrieval config-version must be v1 and ready-graph-name is required"); } - if (!RetrievalMode.KEYWORD.name().equals(defaultMode)) { - throw invalid("default-mode must be KEYWORD"); + try { + if (RetrievalMode.valueOf(defaultMode) != RetrievalMode.KEYWORD) { + throw invalid("default-mode must be KEYWORD"); + } + } catch (IllegalArgumentException | NullPointerException e) { + throw invalid("default-mode is unsupported: " + defaultMode); } if (!ExecutionMode.SEQUENTIAL.name().equals(defaultExecutionMode)) { throw invalid("default-execution-mode must be SEQUENTIAL"); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallPlan.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallPlan.java new file mode 100644 index 000000000..52a41c6ea --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallPlan.java @@ -0,0 +1,95 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.execution; + +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalMode; + +/** Immutable sequential plan with fixed channel allocations and one absolute deadline. */ +public final class RecallPlan { + private final RetrievalMode mode; + private final int topK; + private final long deadlineNanos; + private final Map channelBudgets; + private final String graphVersion; + private final String indexVersion; + private final String vectorVersion; + + public RecallPlan(RetrievalMode mode, int topK, int candidates, long deadlineNanos) { + this(mode, topK, candidates, deadlineNanos, null, null, null); + } + + public RecallPlan(RetrievalMode mode, int topK, int candidates, long deadlineNanos, + String graphVersion, String indexVersion, String vectorVersion) { + this.graphVersion = graphVersion; + this.indexVersion = indexVersion; + this.vectorVersion = vectorVersion; + if (mode == null || topK < 1 || topK > candidates + || (mode.canonical() == RetrievalMode.HYBRID && candidates < 3)) { + throw new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, "invalid recall budget or mode"); + } + this.mode = mode.canonical(); + this.topK = topK; + this.deadlineNanos = deadlineNanos; + Map budgets = new LinkedHashMap<>(); + if (this.mode == RetrievalMode.HYBRID) { + int base = candidates / 3; + int remainder = candidates % 3; + budgets.put("BM25", base + (remainder >= 1 ? 1 : 0)); + budgets.put("VECTOR", base + (remainder >= 2 ? 1 : 0)); + budgets.put("GRAPH", base); + } else { + budgets.put(this.mode == RetrievalMode.BM25_ONLY ? "BM25" + : this.mode == RetrievalMode.VECTOR_ONLY ? "VECTOR" : "GRAPH", candidates); + } + channelBudgets = Collections.unmodifiableMap(budgets); + } + + public String getGraphVersion() { + return graphVersion; + } + + public String getIndexVersion() { + return indexVersion; + } + + public String getVectorVersion() { + return vectorVersion; + } + + public RetrievalMode getMode() { + return mode; + } + + public int getTopK() { + return topK; + } + + public long getDeadlineNanos() { + return deadlineNanos; + } + + public Map getChannelBudgets() { + return channelBudgets; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallStageStatus.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallStageStatus.java new file mode 100644 index 000000000..eaf014ec3 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallStageStatus.java @@ -0,0 +1,24 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.execution; + +/** Observable execution status of a selected recall channel. */ +public enum RecallStageStatus { + SUCCESS, DEGRADED, NOT_RUN +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallStopReason.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallStopReason.java new file mode 100644 index 000000000..bade092cb --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/execution/RecallStopReason.java @@ -0,0 +1,25 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.execution; + +/** Typed terminal reasons for a bounded recall stage. */ +public enum RecallStopReason { + COMPLETED, CANDIDATE_BUDGET, TOP_K, NO_RESULT, DEADLINE, INDEX_NOT_READY, + NO_RELIABLE_ANCHOR, EDGE_SCAN_LIMIT, CANDIDATE_LIMIT, MAX_HOP_REACHED, TRACE_VALIDATION_ERROR +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfConfig.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfConfig.java new file mode 100644 index 000000000..f5077a2e4 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfConfig.java @@ -0,0 +1,64 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.fusion; + +/** Validated fixed weights for the v1 weighted reciprocal rank fusion. */ +public final class WeightedRrfConfig { + public static final WeightedRrfConfig DEFAULT = new WeightedRrfConfig(1.0, 1.0, 1.0, 60); + private final double bm25Weight; + private final double vectorWeight; + private final double graphWeight; + private final int rankConstant; + + public WeightedRrfConfig(double bm25Weight, double vectorWeight, int rankConstant) { + this(bm25Weight, vectorWeight, 0.0, rankConstant); + } + + public WeightedRrfConfig(double bm25Weight, double vectorWeight, double graphWeight, + int rankConstant) { + if (!Double.isFinite(bm25Weight) || !Double.isFinite(vectorWeight) + || !Double.isFinite(graphWeight) + || bm25Weight < 0.0 || vectorWeight < 0.0 + || graphWeight < 0.0 + || (bm25Weight == 0.0 && vectorWeight == 0.0 && graphWeight == 0.0) + || rankConstant < 1) { + throw new IllegalArgumentException("invalid weighted RRF configuration"); + } + this.bm25Weight = bm25Weight; + this.vectorWeight = vectorWeight; + this.graphWeight = graphWeight; + this.rankConstant = rankConstant; + } + + public double getBm25Weight() { + return bm25Weight; + } + + public double getVectorWeight() { + return vectorWeight; + } + + public double getGraphWeight() { + return graphWeight; + } + + public int getRankConstant() { + return rankConstant; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfFusion.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfFusion.java new file mode 100644 index 000000000..b6d18db92 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfFusion.java @@ -0,0 +1,70 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.fusion; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import org.apache.geaflow.ai.retrieval.model.evidence.ChannelScore; +import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; + +/** Deterministic weighted RRF ordering for BM25, vector, and graph evidence. */ +public final class WeightedRrfFusion { + private final WeightedRrfConfig config; + + public WeightedRrfFusion(WeightedRrfConfig config) { + this.config = java.util.Objects.requireNonNull(config, "config"); + } + + public WeightedRrfConfig getConfig() { + return config; + } + + public double score(Evidence evidence) { + double result = 0.0; + for (ChannelScore score : evidence.getStageScores().values()) { + double weight = "BM25".equals(score.getChannel()) ? config.getBm25Weight() + : "VECTOR".equals(score.getChannel()) ? config.getVectorWeight() + : "GRAPH".equals(score.getChannel()) ? config.getGraphWeight() : 0.0; + if (weight > 0.0) { + result += weight / ((double) config.getRankConstant() + score.getRank()); + } + } + return result; + } + + public List sort(List evidence) { + List result = new ArrayList<>(); + for (Evidence value : evidence) { + result.add(new Evidence(value.getEvidenceId(), value.getKind(), value.getText(), + value.getChunks(), value.getEntities(), value.getPaths(), value.getSources(), + value.getStageScores(), score(value), value.getFinalScore(), value.getRank(), + value.getGraphTraces())); + } + result.sort(Comparator.comparingDouble(Evidence::getFusedScore).reversed() + .thenComparing(Evidence::getEvidenceId)); + List ranked = new ArrayList<>(); + for (Evidence value : result) { + ranked.add(new Evidence(value.getEvidenceId(), value.getKind(), value.getText(), value.getChunks(), + value.getEntities(), value.getPaths(), value.getSources(), value.getStageScores(), value.getFusedScore(), + value.getFinalScore(), ranked.size() + 1, value.getGraphTraces())); + } + return ranked; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ChannelSearchResult.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ChannelSearchResult.java new file mode 100644 index 000000000..c84b43df7 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ChannelSearchResult.java @@ -0,0 +1,57 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.index; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; + +/** Search hits and execution counters produced by one recall channel. */ +public final class ChannelSearchResult { + private final List hits; + private final int candidatesEvaluated; + private final RecallStopReason stopReason; + + public ChannelSearchResult(List hits, int candidatesEvaluated, boolean deadlineReached) { + this(hits, candidatesEvaluated, deadlineReached ? RecallStopReason.DEADLINE : RecallStopReason.COMPLETED); + } + + public ChannelSearchResult(List hits, int candidatesEvaluated, RecallStopReason stopReason) { + this.hits = Collections.unmodifiableList(new ArrayList<>(hits)); + this.candidatesEvaluated = candidatesEvaluated; + this.stopReason = stopReason; + } + + public List getHits() { + return hits; + } + + public int getCandidatesEvaluated() { + return candidatesEvaluated; + } + + public RecallStopReason getStopReason() { + return stopReason; + } + + public boolean isDeadlineReached() { + return stopReason == RecallStopReason.DEADLINE; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/Bm25Hit.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/Bm25Hit.java new file mode 100644 index 000000000..4e92d5d48 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/Bm25Hit.java @@ -0,0 +1,58 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.index.bm25; + +/** + * A deterministic BM25 match returned by {@link Bm25IndexReader}. + */ +public final class Bm25Hit { + private final String chunkId; + private final String documentId; + private final float score; + private final int rank; + private final String text; + + public Bm25Hit(String chunkId, String documentId, float score, int rank, String text) { + this.chunkId = chunkId; + this.documentId = documentId; + this.score = score; + this.rank = rank; + this.text = text; + } + + public String getChunkId() { + return chunkId; + } + + public String getDocumentId() { + return documentId; + } + + public float getScore() { + return score; + } + + public int getRank() { + return rank; + } + + public String getText() { + return text; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/Bm25IndexReader.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/Bm25IndexReader.java new file mode 100644 index 000000000..efc2a3d39 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/Bm25IndexReader.java @@ -0,0 +1,211 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.index.bm25; + +import java.io.Closeable; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.function.LongSupplier; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; +import org.apache.geaflow.ai.retrieval.index.ArtifactIdentity; +import org.apache.geaflow.ai.retrieval.index.ChannelSearchResult; +import org.apache.lucene.analysis.standard.StandardAnalyzer; +import org.apache.lucene.document.Document; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.queryparser.classic.QueryParser; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.Scorable; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.SimpleCollector; +import org.apache.lucene.store.Directory; +import org.apache.lucene.store.FSDirectory; + +/** Reusable bounded reader for a published Lucene BM25 artifact. */ +public final class Bm25IndexReader implements Closeable { + private final Directory directory; + private final DirectoryReader reader; + private final IndexSearcher searcher; + private final ArtifactIdentity identity; + + public Bm25IndexReader(Path artifact) { + this(artifact, null, null); + } + + public Bm25IndexReader(Path artifact, String expectedGraphVersion, String expectedIndexVersion) { + Directory openedDirectory = null; + DirectoryReader openedReader = null; + try { + if (artifact == null || !Files.isDirectory(artifact)) { + throw new IOException("missing BM25 artifact: " + artifact); + } + openedDirectory = FSDirectory.open(artifact); + openedReader = DirectoryReader.open(openedDirectory); + identity = ArtifactIdentity.fromMap(openedReader.getIndexCommit().getUserData()); + identity.validateVersions(expectedGraphVersion, expectedIndexVersion); + directory = openedDirectory; + reader = openedReader; + searcher = new IndexSearcher(reader); + } catch (IOException | RuntimeException error) { + closeOnFailure(openedReader, error); + closeOnFailure(openedDirectory, error); + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, "unable to open BM25 artifact", error); + } + } + + private static void closeOnFailure(Closeable resource, Exception failure) { + if (resource != null) { + try { + resource.close(); + } catch (IOException error) { + failure.addSuppressed(error); + } + } + } + + public ArtifactIdentity getIdentity() { + return identity; + } + + public List search(String queryText, int candidateLimit, int topK) { + return search(queryText, candidateLimit, topK, 0L); + } + + public List search(String queryText, int candidateLimit, int topK, long deadlineNanos) { + ChannelSearchResult found = searchWithStats(queryText, candidateLimit, topK, deadlineNanos); + return new ArrayList<>(found.getHits().subList(0, Math.min(topK, found.getHits().size()))); + } + + public ChannelSearchResult searchWithStats(String queryText, int candidateLimit, + int topK, long deadlineNanos) { + return searchWithStats(queryText, candidateLimit, topK, deadlineNanos, System::nanoTime); + } + + /** Counts documents scored, rather than the number of hits returned by Lucene. */ + public ChannelSearchResult searchWithStats(String queryText, int candidateLimit, + int topK, long deadlineNanos, LongSupplier clock) { + if (queryText == null || queryText.trim().isEmpty() || candidateLimit < 1 || topK < 1) { + throw new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, "invalid BM25 query or budget"); + } + BoundedCollector collector = new BoundedCollector(candidateLimit, deadlineNanos, clock); + try (StandardAnalyzer analyzer = new StandardAnalyzer()) { + collector.checkDeadline(); + Query query = new QueryParser("text", analyzer).parse(queryText); + collector.checkDeadline(); + searcher.search(query, collector); + collector.checkDeadline(); + } catch (SearchStopped stopped) { + // The collector preserves all fully evaluated hits and the typed stopping reason. + } catch (org.apache.lucene.queryparser.classic.ParseException error) { + throw new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, "invalid BM25 query", error); + } catch (IOException | RuntimeException error) { + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, "BM25 artifact search failed", error); + } + collector.hits.sort(Comparator.comparing(Bm25Hit::getScore).reversed() + .thenComparing(Bm25Hit::getDocumentId).thenComparing(Bm25Hit::getChunkId)); + List result = new ArrayList<>(); + for (Bm25Hit hit : collector.hits) { + result.add(new Bm25Hit(hit.getChunkId(), hit.getDocumentId(), hit.getScore(), result.size() + 1, hit.getText())); + } + return new ChannelSearchResult<>(result, collector.evaluated, collector.reason); + } + + @Override + public void close() throws IOException { + try { + reader.close(); + } finally { + directory.close(); + } + } + + private final class BoundedCollector extends SimpleCollector { + private final int limit; + private final long deadline; + private final LongSupplier clock; + private final List hits = new ArrayList<>(); + private final Set chunkIds = new HashSet<>(); + private Scorable scorer; + private int docBase; + private int evaluated; + private RecallStopReason reason = RecallStopReason.COMPLETED; + + private BoundedCollector(int limit, long deadline, LongSupplier clock) { + this.limit = limit; + this.deadline = deadline; + this.clock = clock; + } + + private void checkDeadline() { + if (deadline != 0L && clock.getAsLong() - deadline >= 0L) { + reason = RecallStopReason.DEADLINE; + throw new SearchStopped(); + } + } + + @Override + protected void doSetNextReader(LeafReaderContext context) { + checkDeadline(); + docBase = context.docBase; + } + + @Override + public void setScorer(Scorable scorer) { + this.scorer = scorer; + } + + @Override + public void collect(int doc) throws IOException { + checkDeadline(); + if (evaluated >= limit) { + reason = RecallStopReason.CANDIDATE_BUDGET; + throw new SearchStopped(); + } + evaluated++; + float score = scorer.score(); + Document document = searcher.doc(docBase + doc); + String chunkId = document.get("chunkId"); + String documentId = document.get("documentId"); + String text = document.get("text"); + if (chunkId == null || documentId == null || text == null || !Float.isFinite(score) + || !chunkIds.add(chunkId)) { + throw new IOException("invalid or duplicate BM25 document"); + } + hits.add(new Bm25Hit(chunkId, documentId, score, 0, text)); + } + + @Override + public ScoreMode scoreMode() { + return ScoreMode.COMPLETE; + } + } + + private static final class SearchStopped extends RuntimeException { + private static final long serialVersionUID = 1L; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/VectorHit.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/VectorHit.java new file mode 100644 index 000000000..e4c8ff1fc --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/VectorHit.java @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index.vector; + +/** A deterministic cosine-similarity match. */ +public final class VectorHit { + private final String chunkId; + private final String documentId; + private final double similarity; + private final int rank; + + public VectorHit(String chunkId, String documentId, double similarity, int rank) { + this.chunkId = chunkId; + this.documentId = documentId; + this.similarity = similarity; + this.rank = rank; + } + + public String getChunkId() { + return chunkId; + } + + public String getDocumentId() { + return documentId; + } + + public double getSimilarity() { + return similarity; + } + + public int getRank() { + return rank; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/VectorIndexReader.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/VectorIndexReader.java new file mode 100644 index 000000000..f52146303 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/VectorIndexReader.java @@ -0,0 +1,232 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.index.vector; + +import java.io.Closeable; +import java.io.DataInputStream; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.function.LongSupplier; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; +import org.apache.geaflow.ai.retrieval.index.ArtifactIdentity; +import org.apache.geaflow.ai.retrieval.index.ChannelSearchResult; + +/** + * Reader for the binary format emitted by {@link OfflineVectorIndexBuilder}. + */ +public final class VectorIndexReader implements Closeable { + private static final String MAGIC = "GEAFLOW-VECTOR-2"; + private final ArtifactIdentity identity; + private final List entries; + private final int dimensions; + private final String vectorSource; + private final String vectorVersion; + + public VectorIndexReader(Path artifact) { + this(artifact, null, null, null, null); + } + + public VectorIndexReader(Path artifact, String expectedGraphVersion, String expectedIndexVersion, + String expectedSource, String expectedVectorVersion) { + if (artifact == null || !Files.isRegularFile(artifact)) { + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, + "missing vector artifact: " + artifact); + } + try (DataInputStream input = new DataInputStream(Files.newInputStream(artifact))) { + if (!MAGIC.equals(input.readUTF())) { + throw new IOException("invalid vector artifact magic"); + } + identity = new ArtifactIdentity(input.readUTF(), input.readUTF(), input.readUTF(), input.readUTF()); + identity.validateVersions(expectedGraphVersion, expectedIndexVersion); + vectorSource = input.readUTF(); + vectorVersion = input.readUTF(); + if (expectedSource != null && !expectedSource.equals(vectorSource) + || expectedVectorVersion != null && !expectedVectorVersion.equals(vectorVersion)) { + throw new IOException("vector source/version mismatch"); + } + dimensions = input.readInt(); + int count = input.readInt(); + if (dimensions <= 0 || count < 0 || (long) dimensions * 4 > Files.size(artifact) + || (long) count * ((long) dimensions * 4 + 4) > Files.size(artifact)) { + throw new IOException("invalid vector dimensions/count"); + } + entries = new ArrayList<>(count); + Set chunkIds = new HashSet<>(); + for (int i = 0; i < count; i++) { + String chunkId = input.readUTF(); + String documentId = input.readUTF(); + if (chunkId.isEmpty() || documentId.isEmpty() || !chunkIds.add(chunkId)) { + throw new IOException("invalid or duplicate vector chunk ID"); + } + float[] vector = new float[dimensions]; + for (int d = 0; d < dimensions; d++) { + vector[d] = input.readFloat(); + if (!Float.isFinite(vector[d])) { + throw new IOException("non-finite vector value"); + } + } + if (norm(vector) == 0.0) { + throw new IOException("zero-norm artifact vector"); + } + entries.add(new Entry(chunkId, documentId, vector)); + } + entries.sort(Comparator.comparing((Entry entry) -> entry.chunkId)); + if (input.read() != -1) { + throw new IOException("trailing vector artifact data"); + } + } catch (Exception error) { + if (error instanceof RetrievalException) { + throw (RetrievalException) error; + } + throw new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, + "unable to read vector artifact", error); + } + } + + public ArtifactIdentity getIdentity() { + return identity; + } + + public int getDimensions() { + return dimensions; + } + + public String getVectorSource() { + return vectorSource; + } + + public String getVectorVersion() { + return vectorVersion; + } + + public List search(float[] query, int candidateLimit, int topK) { + return search(query, candidateLimit, topK, 0L); + } + + public List search(float[] query, int candidateLimit, int topK, long deadlineNanos) { + List hits = searchWithStats(query, candidateLimit, topK, deadlineNanos).getHits(); + return new ArrayList<>(hits.subList(0, Math.min(topK, hits.size()))); + } + + public ChannelSearchResult searchWithStats(float[] query, int candidateLimit, + int topK, long deadlineNanos) { + return searchWithStats(query, candidateLimit, topK, deadlineNanos, System::nanoTime); + } + + public ChannelSearchResult searchWithStats(float[] query, int candidateLimit, + int topK, long deadlineNanos, LongSupplier clock) { + validateQuery(query); + if (candidateLimit < 1 || topK < 1) { + throw new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, "invalid vector budget"); + } + double queryNorm = norm(query); + List hits = new ArrayList<>(); + int evaluated = 0; + RecallStopReason reason = RecallStopReason.COMPLETED; + for (Entry entry : entries) { + if (deadlineNanos != 0L && clock.getAsLong() - deadlineNanos >= 0L) { + reason = RecallStopReason.DEADLINE; + break; + } + if (evaluated >= candidateLimit) { + reason = RecallStopReason.CANDIDATE_BUDGET; + break; + } + evaluated++; + double score = dot(query, entry.vector) / (queryNorm * norm(entry.vector)); + if (Double.isFinite(score)) { + hits.add(new VectorHit(entry.chunkId, entry.documentId, score, 0)); + } + } + hits.sort(Comparator.comparing(VectorHit::getSimilarity).reversed() + .thenComparing(VectorHit::getChunkId)); + List result = new ArrayList<>(); + int rank = 1; + for (VectorHit hit : hits) { + if (rank > candidateLimit) { + break; + } + result.add(new VectorHit(hit.getChunkId(), hit.getDocumentId(), hit.getSimilarity(), rank++)); + } + if (deadlineNanos != 0L && clock.getAsLong() - deadlineNanos >= 0L) { + reason = RecallStopReason.DEADLINE; + } + return new ChannelSearchResult<>(result, evaluated, reason); + } + + public void validateQuery(float[] query) { + validateValues(query); + if (query.length != dimensions) { + throw new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, "query vector dimensions do not match artifact"); + } + } + + public static void validateValues(float[] query) { + if (query == null || query.length == 0) { + throw new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, "queryVector is required"); + } + for (float value : query) { + if (!Float.isFinite(value)) { + throw new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, "query vector must contain finite values"); + } + } + double queryNorm = norm(query); + if (!Double.isFinite(queryNorm) || queryNorm == 0.0) { + throw new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, + "query vector must have a finite, non-zero norm"); + } + } + + private static double dot(float[] first, float[] second) { + double value = 0.0; + for (int i = 0; i < first.length; i++) { + value += (double) first[i] * second[i]; + } + return value; + } + + private static double norm(float[] vector) { + return Math.sqrt(dot(vector, vector)); + } + + @Override + public void close() { + // The artifact is loaded into immutable memory; no resources remain open. + } + + private static final class Entry { + private final String chunkId; + private final String documentId; + private final float[] vector; + + private Entry(String chunkId, String documentId, float[] vector) { + this.chunkId = chunkId; + this.documentId = documentId; + this.vector = vector; + } + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/Evidence.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/Evidence.java index a2e7a94fb..f0e1eefd6 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/Evidence.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/Evidence.java @@ -54,6 +54,8 @@ public final class Evidence { private final List entities; @SerializedName("paths") private final List paths; + @SerializedName("graphTraces") + private final List graphTraces; @SerializedName("sources") private final List sources; @SerializedName("stageScores") @@ -72,6 +74,7 @@ private Evidence() { this.chunks = Collections.emptyList(); this.entities = Collections.emptyList(); this.paths = Collections.emptyList(); + this.graphTraces = Collections.emptyList(); this.sources = Collections.emptyList(); this.stageScores = Collections.emptyMap(); this.fusedScore = null; @@ -83,19 +86,28 @@ public Evidence(String evidenceId, EvidenceKind kind, String text, List entities, List paths, List sources, Map stageScores, Double fusedScore, Integer rank) { this(evidenceId, kind, text, chunks, entities, paths, sources, stageScores, fusedScore, - null, rank); + null, rank, Collections.emptyList()); } public Evidence(String evidenceId, EvidenceKind kind, String text, List chunks, List entities, List paths, List sources, Map stageScores, Double fusedScore, Double finalScore, Integer rank) { + this(evidenceId, kind, text, chunks, entities, paths, sources, stageScores, fusedScore, + finalScore, rank, Collections.emptyList()); + } + + public Evidence(String evidenceId, EvidenceKind kind, String text, List chunks, + List entities, List paths, List sources, + Map stageScores, Double fusedScore, Double finalScore, + Integer rank, List graphTraces) { this.evidenceId = ModelValidation.optionalNonBlank(evidenceId, "evidenceId"); this.kind = Objects.requireNonNull(kind, "kind"); this.text = ModelValidation.optional(text); this.chunks = ModelValidation.immutableList(chunks, "chunks"); this.entities = ModelValidation.immutableList(entities, "entities"); this.paths = ModelValidation.immutableList(paths, "paths"); + this.graphTraces = ModelValidation.immutableList(graphTraces, "graphTraces"); this.sources = ModelValidation.immutableList(sources, "sources"); this.stageScores = ModelValidation.sortedMap(stageScores); for (Map.Entry entry : this.stageScores.entrySet()) { @@ -140,6 +152,11 @@ public List getPaths() { return Collections.unmodifiableList(paths == null ? Collections.emptyList() : paths); } + public List getGraphTraces() { + return Collections.unmodifiableList(graphTraces == null + ? Collections.emptyList() : graphTraces); + } + public List getSources() { return Collections.unmodifiableList(sources == null ? Collections.emptyList() : sources); } @@ -219,7 +236,7 @@ public boolean equals(Object object) { return Objects.equals(evidenceId, that.evidenceId) && kind == that.kind && Objects.equals(text, that.text) && Objects.equals(chunks, that.chunks) && Objects.equals(entities, that.entities) && Objects.equals(paths, that.paths) - && Objects.equals(sources, that.sources) + && Objects.equals(sources, that.sources) && Objects.equals(graphTraces, that.graphTraces) && Objects.equals(stageScores, that.stageScores) && Objects.equals(fusedScore, that.fusedScore) && Objects.equals(finalScore, that.finalScore) @@ -228,7 +245,7 @@ public boolean equals(Object object) { @Override public int hashCode() { - return Objects.hash(evidenceId, kind, text, chunks, entities, paths, sources, + return Objects.hash(evidenceId, kind, text, chunks, entities, paths, sources, graphTraces, stageScores, fusedScore, finalScore, rank); } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/GraphEvidenceTrace.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/GraphEvidenceTrace.java new file mode 100644 index 000000000..b1a758662 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/evidence/GraphEvidenceTrace.java @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.model.evidence; + +import com.google.gson.annotations.SerializedName; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import org.apache.geaflow.ai.retrieval.model.graph.GraphPathRef; +import org.apache.geaflow.ai.retrieval.validation.ModelValidation; + +/** Per-candidate graph provenance retained alongside a fused evidence item. */ +public final class GraphEvidenceTrace { + @SerializedName("queryId") + private final String queryId; + @SerializedName("graphVersion") + private final String graphVersion; + @SerializedName("anchorEntityId") + private final String anchorEntityId; + @SerializedName("anchorMatchType") + private final String anchorMatchType; + @SerializedName("anchorConfidence") + private final double anchorConfidence; + @SerializedName("path") + private final GraphPathRef path; + @SerializedName("supportingChunkIds") + private final List supportingChunkIds; + + private GraphEvidenceTrace() { + queryId = null; + graphVersion = null; + anchorEntityId = null; + anchorMatchType = null; + anchorConfidence = 0.0; + path = null; + supportingChunkIds = Collections.emptyList(); + } + + public GraphEvidenceTrace(String queryId, String graphVersion, String anchorEntityId, + String anchorMatchType, double anchorConfidence, GraphPathRef path, + List supportingChunkIds) { + this.queryId = ModelValidation.required(queryId, "queryId"); + this.graphVersion = ModelValidation.required(graphVersion, "graphVersion"); + this.anchorEntityId = ModelValidation.required(anchorEntityId, "anchorEntityId"); + this.anchorMatchType = ModelValidation.required(anchorMatchType, "anchorMatchType"); + this.anchorConfidence = ModelValidation.finite(anchorConfidence, "anchorConfidence"); + if (anchorConfidence < 0.0 || anchorConfidence > 1.0) { + throw new IllegalArgumentException("anchorConfidence must be in [0, 1]"); + } + this.path = Objects.requireNonNull(path, "path"); + this.supportingChunkIds = ModelValidation.sortedStrings(supportingChunkIds, "supportingChunkId"); + } + + public String getQueryId() { + return queryId; + } + + public String getGraphVersion() { + return graphVersion; + } + + public String getAnchorEntityId() { + return anchorEntityId; + } + + public String getAnchorMatchType() { + return anchorMatchType; + } + + public double getAnchorConfidence() { + return anchorConfidence; + } + + public GraphPathRef getPath() { + return path; + } + + public List getSupportingChunkIds() { + return Collections.unmodifiableList(supportingChunkIds); + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof GraphEvidenceTrace)) { + return false; + } + GraphEvidenceTrace that = (GraphEvidenceTrace) object; + return Double.compare(anchorConfidence, that.anchorConfidence) == 0 + && Objects.equals(queryId, that.queryId) + && Objects.equals(graphVersion, that.graphVersion) + && Objects.equals(anchorEntityId, that.anchorEntityId) + && Objects.equals(anchorMatchType, that.anchorMatchType) + && Objects.equals(path, that.path) + && Objects.equals(supportingChunkIds, that.supportingChunkIds); + } + + @Override + public int hashCode() { + return Objects.hash(queryId, graphVersion, anchorEntityId, anchorMatchType, + anchorConfidence, path, supportingChunkIds); + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RecallService.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RecallService.java new file mode 100644 index 000000000..8f9b81e21 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RecallService.java @@ -0,0 +1,496 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.service; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Objects; +import java.util.UUID; +import java.util.function.LongSupplier; +import org.apache.geaflow.ai.retrieval.api.model.ExecutionMode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalCommand; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalMode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalRequest; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalTrace; +import org.apache.geaflow.ai.retrieval.api.model.TraceStage; +import org.apache.geaflow.ai.retrieval.channel.graph.GraphNeighborProvider; +import org.apache.geaflow.ai.retrieval.channel.graph.GraphRetriever; +import org.apache.geaflow.ai.retrieval.channel.graph.InMemoryGraphNeighborProvider; +import org.apache.geaflow.ai.retrieval.config.RetrievalProperties; +import org.apache.geaflow.ai.retrieval.execution.RecallPlan; +import org.apache.geaflow.ai.retrieval.execution.RecallStageStatus; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; +import org.apache.geaflow.ai.retrieval.fusion.WeightedRrfConfig; +import org.apache.geaflow.ai.retrieval.fusion.WeightedRrfFusion; +import org.apache.geaflow.ai.retrieval.index.ChannelSearchResult; +import org.apache.geaflow.ai.retrieval.index.bm25.Bm25Hit; +import org.apache.geaflow.ai.retrieval.index.bm25.Bm25IndexReader; +import org.apache.geaflow.ai.retrieval.index.vector.VectorHit; +import org.apache.geaflow.ai.retrieval.index.vector.VectorIndexReader; +import org.apache.geaflow.ai.retrieval.model.document.SourceRef; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.evidence.ChannelScore; +import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; +import org.apache.geaflow.ai.retrieval.model.evidence.EvidenceKind; +import org.apache.geaflow.ai.retrieval.model.evidence.GraphEvidenceTrace; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphPathRef; + +/** Session-free sequential recall; all execution state belongs to one request. */ +public final class RecallService { + private final Bm25IndexReader bm25; + private final VectorIndexReader vector; + private final List chunks; + private final Map chunksById; + private final List entities; + private final List edges; + private final GraphNeighborProvider graphProvider; + private final RetrievalFixtureRegistry registry; + private final WeightedRrfFusion fusion; + private final LongSupplier clock; + + public RecallService(Bm25IndexReader bm25, VectorIndexReader vector, List chunks, + List entities, List edges) { + this(bm25, vector, chunks, entities, edges, WeightedRrfConfig.DEFAULT); + } + + public RecallService(Bm25IndexReader bm25, VectorIndexReader vector, List chunks, + List entities, List edges, WeightedRrfConfig config) { + this(bm25, vector, chunks, entities, edges, null, new InMemoryGraphNeighborProvider(edges), config, + System::nanoTime); + } + + public RecallService(RetrievalFixtureRegistry registry) { + this(registry, WeightedRrfConfig.DEFAULT); + } + + public RecallService(RetrievalFixtureRegistry registry, WeightedRrfConfig config) { + this(registry, config, System::nanoTime); + } + + /** Accepts a monotonic clock for reproducible deadline verification. */ + public RecallService(RetrievalFixtureRegistry registry, WeightedRrfConfig config, LongSupplier clock) { + this(null, null, Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), + Objects.requireNonNull(registry, "registry"), null, config, clock); + } + + private RecallService(Bm25IndexReader bm25, VectorIndexReader vector, List chunks, + List entities, List edges, RetrievalFixtureRegistry registry, + GraphNeighborProvider graphProvider, WeightedRrfConfig config, LongSupplier clock) { + this.bm25 = bm25; + this.vector = vector; + this.chunks = immutable(chunks); + this.chunksById = indexChunks(this.chunks); + this.entities = immutable(entities); + this.edges = immutable(edges); + this.graphProvider = graphProvider; + this.registry = registry; + this.fusion = new WeightedRrfFusion(config); + this.clock = Objects.requireNonNull(clock, "clock"); + } + + private static List immutable(List values) { + return values == null ? Collections.emptyList() : Collections.unmodifiableList(new ArrayList<>(values)); + } + + private static Map indexChunks(List chunks) { + Map indexed = new LinkedHashMap<>(); + for (TextChunk chunk : chunks) { + indexed.put(chunk.getChunkId(), chunk); + } + return Collections.unmodifiableMap(indexed); + } + + private RecallService(RetrievalFixtureRegistry.Fixture fixture, WeightedRrfConfig config, + LongSupplier clock) { + this.bm25 = fixture.getBm25(); + this.vector = fixture.getVector(); + this.chunks = fixture.getChunks(); + this.chunksById = fixture.getChunksById(); + this.entities = fixture.getEntities(); + this.edges = fixture.getEdges(); + this.graphProvider = fixture.getGraphProvider(); + this.registry = null; + this.fusion = new WeightedRrfFusion(config); + this.clock = Objects.requireNonNull(clock, "clock"); + } + + public RetrievalResponse retrieve(RetrievalRequest request) { + final long started = clock.getAsLong(); + RetrievalCommand command = new RetrievalRequestValidator(new RetrievalProperties()).validate(request); + requireStructuredVersions(request); + if (registry == null) { + throw notReady("fixture registry is required for structured retrieval"); + } + RetrievalBudget budget = command.getBudget(); + long deadline = started + budget.getTimeoutMs() * 1_000_000L; + try (RetrievalFixtureRegistry.Fixture.Lease lease = registry.acquire(command.getGraphName(), + request.getGraphVersion(), request.getIndexVersion())) { + RetrievalFixtureRegistry.Fixture fixture = lease.getFixture(); + if (fixture.getVector() != null + && ((request.getVectorSource() != null + && !request.getVectorSource().equals(fixture.getVector().getVectorSource())) + || (request.getVectorVersion() != null + && !request.getVectorVersion().equals(fixture.getVector().getVectorVersion())))) { + throw notReady("vector source/version mismatch"); + } + RecallService delegate = new RecallService(fixture, fusion.getConfig(), clock); + RetrievalResponse response = delegate.execute(command, toFloat(command.getQueryVector()), deadline, + request.isAllowPartialResults(), started, fixture); + response.setGraphName(fixture.getGraphName()); + response.setGraphVersion(fixture.getGraphVersion()); + response.getTrace().setGraphVersion(fixture.getGraphVersion()); + response.getTrace().setIndexVersion(fixture.getIndexVersion()); + if (fixture.getVector() != null) { + response.getTrace().setVectorVersion(fixture.getVector().getVectorVersion()); + } + if (request.getRequestId() != null) { + response.setRequestId(request.getRequestId()); + response.getTrace().setRequestId(request.getRequestId()); + } + response.getTrace().setElapsedNanos(clock.getAsLong() - started); + return response; + } + } + + private static float[] toFloat(List values) { + float[] result = new float[values.size()]; + for (int index = 0; index < values.size(); index++) { + result[index] = values.get(index).floatValue(); + } + return result; + } + + public List retrieve(String query, float[] queryVector, String mode, int topK, + int maxCandidates, long deadlineNanos) { + return retrieveWithTrace(query, queryVector, mode, topK, maxCandidates, deadlineNanos).getEvidence(); + } + + public RetrievalResponse retrieveWithTrace(String query, float[] queryVector, String mode, + int topK, int maxCandidates, long deadlineNanos) { + return retrieveWithTrace(query, queryVector, mode, topK, maxCandidates, deadlineNanos, false); + } + + public RetrievalResponse retrieveWithTrace(String query, float[] queryVector, String mode, + int topK, int maxCandidates, long deadlineNanos, + boolean allowPartialResults) { + final long started = clock.getAsLong(); + RetrievalRequest request = new RetrievalRequest(); + request.setGraphName("fixture"); + request.setQuery(query); + request.setMode(mode == null ? "BM25_ONLY" : mode.trim().toUpperCase(Locale.ROOT)); + request.setBudget(new RetrievalBudget(topK, 3000, maxCandidates, null)); + List values = new ArrayList<>(); + if (queryVector != null) { + for (float value : queryVector) { + values.add((double) value); + } + } + request.setQueryVector(values); + RetrievalCommand command = new RetrievalRequestValidator(new RetrievalProperties()).validate(request); + if (registry != null) { + throw notReady("registry retrieval requires a structured request with graph/index versions"); + } + return execute(command, queryVector, deadlineNanos, allowPartialResults, started, null); + } + + private RetrievalResponse execute(RetrievalCommand command, float[] queryVector, long deadline, + boolean partial, long started, RetrievalFixtureRegistry.Fixture fixture) { + RetrievalBudget budget = command.getBudget(); + RecallPlan plan = new RecallPlan(command.getMode(), budget.getTopK(), budget.getMaxCandidates(), deadline, + fixture == null ? null : fixture.getGraphVersion(), fixture == null ? null : fixture.getIndexVersion(), + vector == null ? null : vector.getVectorVersion()); + Map failures = fixture == null ? Collections.emptyMap() : fixture.getFailures(); + if (plan.getChannelBudgets().containsKey("VECTOR")) { + VectorIndexReader.validateValues(queryVector); + if (vector != null) { + vector.validateQuery(queryVector); + } + } + Map byId = chunksById; + Map merged = new LinkedHashMap<>(); + RetrievalTrace trace = new RetrievalTrace(); + trace.setTraceVersion("v1"); + trace.setGraphVersion(plan.getGraphVersion()); + trace.setIndexVersion(plan.getIndexVersion()); + trace.setVectorVersion(plan.getVectorVersion()); + trace.setOriginalQuery(command.getQuery()); + trace.setSelectedMode(plan.getMode()); + trace.setExecutionMode(ExecutionMode.SEQUENTIAL); + trace.setSelectedChannels(new ArrayList<>(plan.getChannelBudgets().keySet())); + trace.setChannelBudgets(plan.getChannelBudgets()); + trace.setEffectiveCandidateBudget(budget.getMaxCandidates()); + trace.setEffectiveTopK(plan.getTopK()); + trace.setBm25Weight(fusion.getConfig().getBm25Weight()); + trace.setVectorWeight(fusion.getConfig().getVectorWeight()); + trace.setGraphWeight(fusion.getConfig().getGraphWeight()); + trace.setRrfRankConstant(fusion.getConfig().getRankConstant()); + Map channelElapsedNanos = new LinkedHashMap<>(); + for (Map.Entry channel : plan.getChannelBudgets().entrySet()) { + String name = channel.getKey(); + long channelStarted = clock.getAsLong(); + trace.getEvaluatedCounts().put(name, 0); + trace.getCandidateCounts().put(name, 0); + if (expired(deadline)) { + timedOut(partial); + record(trace, name, RecallStageStatus.NOT_RUN, RecallStopReason.DEADLINE, "deadline reached before channel execution"); + channelElapsedNanos.put(name, clock.getAsLong() - channelStarted); + continue; + } + try { + RetrievalException loadFailure = failures.get(name); + if (loadFailure != null) { + throw loadFailure; + } + RecallStopReason reason; + if ("BM25".equals(name)) { + if (bm25 == null) { + throw notReady("artifact unavailable"); + } + ChannelSearchResult search = bm25.searchWithStats(command.getQuery(), channel.getValue(), + plan.getTopK(), deadline, clock); + trace.getEvaluatedCounts().put(name, search.getCandidatesEvaluated()); + trace.getCandidateCounts().put(name, search.getHits().size()); + for (Bm25Hit hit : search.getHits()) { + EvidenceBuilder builder = builder(merged, byId, hit.getChunkId(), hit.getDocumentId(), EvidenceKind.CHUNK); + if (!builder.chunk.getText().equals(hit.getText())) { + throw notReady("BM25 text does not match fixture"); + } + builder.scores.put(name, new ChannelScore(name, hit.getScore(), null, hit.getRank())); + } + reason = search.getStopReason(); + } else if ("VECTOR".equals(name)) { + if (vector == null) { + throw notReady("artifact unavailable"); + } + ChannelSearchResult search = vector.searchWithStats(queryVector, channel.getValue(), + plan.getTopK(), deadline, clock); + trace.getEvaluatedCounts().put(name, search.getCandidatesEvaluated()); + trace.getCandidateCounts().put(name, search.getHits().size()); + for (VectorHit hit : search.getHits()) { + EvidenceBuilder builder = builder(merged, byId, hit.getChunkId(), hit.getDocumentId(), EvidenceKind.CHUNK); + builder.scores.put(name, new ChannelScore(name, hit.getSimilarity(), null, hit.getRank())); + } + reason = search.getStopReason(); + } else { + GraphRetriever retriever = new GraphRetriever(); + final List hits = retriever.retrieve(command.getQuery(), entities, graphProvider, byId, + channel.getValue(), channel.getValue(), deadline, clock); + GraphRetriever.GraphStats stats = retriever.getLastStats(); + trace.setAnchorsConsidered(stats.getAnchorsConsidered()); + trace.setEdgesExamined(stats.getEdgesExamined()); + trace.setNeighborsSampled(stats.getNeighborsSampled()); + trace.setVerticesReached(stats.getVerticesReached()); + trace.setGraphCandidatesProduced(stats.getCandidatesProduced()); + trace.getValidationErrors().addAll(stats.getValidationErrors()); + trace.getEvaluatedCounts().put(name, stats.getEdgesExamined()); + trace.getCandidateCounts().put(name, stats.getCandidatesProduced()); + int rank = 0; + for (GraphRetriever.GraphHit hit : hits) { + for (String id : hit.getSourceChunkIds()) { + EvidenceBuilder builder = builder(merged, byId, id, byId.get(id).getDocumentId(), EvidenceKind.GRAPH_PATH); + if (!builder.entities.contains(hit.getAnchor().getEntity())) { + builder.entities.add(hit.getAnchor().getEntity()); + } + if (!builder.paths.contains(hit.getPath())) { + builder.paths.add(hit.getPath()); + } + builder.graphTraces.add(new GraphEvidenceTrace(command.getQuery(), + plan.getGraphVersion() == null ? "fixture" : plan.getGraphVersion(), + hit.getAnchor().getEntity().getEntityId(), hit.getAnchor().getMatchType(), + hit.getAnchor().getConfidence(), hit.getPath(), + hit.getSourceChunkIds())); + if (!builder.scores.containsKey(name)) { + builder.scores.put(name, new ChannelScore(name, hit.getScore(), null, ++rank)); + } + } + } + reason = stats.getStopReason(); + } + if (reason == RecallStopReason.DEADLINE) { + timedOut(partial); + } + record(trace, name, reason == RecallStopReason.DEADLINE + || reason == RecallStopReason.TRACE_VALIDATION_ERROR + ? RecallStageStatus.DEGRADED : RecallStageStatus.SUCCESS, + reason, null); + channelElapsedNanos.put(name, clock.getAsLong() - channelStarted); + } catch (RetrievalException error) { + if (!partial || error.getCode() != RetrievalErrorCode.INDEX_NOT_READY) { + throw error; + } + for (EvidenceBuilder value : merged.values()) { + value.scores.remove(name); + } + merged.values().removeIf(value -> value.scores.isEmpty()); + record(trace, name, RecallStageStatus.DEGRADED, RecallStopReason.INDEX_NOT_READY, error.getMessage()); + channelElapsedNanos.put(name, clock.getAsLong() - channelStarted); + } + } + List evidence = new ArrayList<>(); + for (EvidenceBuilder value : merged.values()) { + evidence.add(value.build()); + } + if (plan.getMode() == RetrievalMode.HYBRID) { + evidence = fusion.sort(evidence); + } else { + evidence.sort(Comparator.comparingDouble(Evidence::getFusedScore).reversed().thenComparing(Evidence::getEvidenceId)); + } + boolean truncated = evidence.size() > plan.getTopK(); + List ranked = new ArrayList<>(); + for (Evidence value : evidence.subList(0, Math.min(plan.getTopK(), evidence.size()))) { + ranked.add(new Evidence(value.getEvidenceId(), value.getKind(), value.getText(), value.getChunks(), + value.getEntities(), value.getPaths(), value.getSources(), value.getStageScores(), value.getFusedScore(), + value.getFinalScore(), ranked.size() + 1, value.getGraphTraces())); + } + boolean deadlineReached = expired(deadline) || trace.getChannelStopReasons().containsValue(RecallStopReason.DEADLINE); + if (deadlineReached) { + timedOut(partial); + } + RecallStopReason stop = deadlineReached ? RecallStopReason.DEADLINE + : trace.getChannelStopReasons().containsValue(RecallStopReason.TRACE_VALIDATION_ERROR) + ? RecallStopReason.TRACE_VALIDATION_ERROR + : trace.getChannelStopReasons().containsValue(RecallStopReason.EDGE_SCAN_LIMIT) + ? RecallStopReason.EDGE_SCAN_LIMIT + : trace.getChannelStopReasons().containsValue(RecallStopReason.CANDIDATE_LIMIT) + ? RecallStopReason.CANDIDATE_LIMIT + : trace.getChannelStopReasons().containsValue(RecallStopReason.CANDIDATE_BUDGET) + ? RecallStopReason.CANDIDATE_BUDGET + : truncated ? RecallStopReason.TOP_K : ranked.isEmpty() ? RecallStopReason.NO_RESULT : RecallStopReason.COMPLETED; + trace.setStopReason(stop.name()); + List stages = new ArrayList<>(); + int evaluated = 0; + for (String name : trace.getSelectedChannels()) { + evaluated += trace.getEvaluatedCounts().get(name); + stages.add(new TraceStage(name, trace.getChannelStatuses().get(name).name(), + "reason=" + trace.getChannelStopReasons().get(name) + ";evaluated=" + trace.getEvaluatedCounts().get(name) + + ";candidates=" + trace.getCandidateCounts().get(name) + + ";elapsedNanos=" + channelElapsedNanos.get(name))); + } + stages.add(new TraceStage("FUSION", deadlineReached ? "DEGRADED" : "SUCCESS", + "weighted_rrf_k=" + trace.getRrfRankConstant() + ";weights=bm25:" + trace.getBm25Weight() + + ",vector:" + trace.getVectorWeight() + ",graph:" + trace.getGraphWeight())); + trace.setStages(stages); + trace.setTotalCandidatesEvaluated(evaluated); + RetrievalResponse response = new RetrievalResponse(); + response.setEvidence(ranked); + List sources = new ArrayList<>(); + List paths = new ArrayList<>(); + for (Evidence value : ranked) { + for (SourceRef source : value.getSources()) { + if (!sources.contains(source)) { + sources.add(source); + } + } + for (GraphPathRef path : value.getPaths()) { + if (!paths.contains(path)) { + paths.add(path); + } + } + } + response.setSources(sources); + response.setPaths(paths); + String responseId = UUID.randomUUID().toString(); + response.setRequestId(responseId); + trace.setRequestId(responseId); + trace.setElapsedNanos(clock.getAsLong() - started); + response.setTrace(trace); + response.setEffectiveBudget(budget); + response.setDegradedChannels(trace.getDegradedChannels()); + return response; + } + + private boolean expired(long deadline) { + return deadline != 0L && clock.getAsLong() - deadline >= 0L; + } + + private static void timedOut(boolean partial) { + if (!partial) { + throw new RetrievalException(RetrievalErrorCode.RETRIEVAL_TIMEOUT, "retrieval deadline exceeded"); + } + } + + private static void record(RetrievalTrace trace, String name, RecallStageStatus status, RecallStopReason reason, String message) { + trace.getChannelStatuses().put(name, status); + trace.getChannelStopReasons().put(name, reason); + if (status != RecallStageStatus.SUCCESS) { + trace.getDegradedChannels().add(name); + trace.getDegradationReasons().put(name, message == null ? reason.name() : message); + } + } + + private static RetrievalException notReady(String message) { + return new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, message); + } + + + private static void requireStructuredVersions(RetrievalRequest request) { + if (request.getGraphVersion() == null || request.getGraphVersion().trim().isEmpty() + || request.getIndexVersion() == null || request.getIndexVersion().trim().isEmpty()) { + throw new RetrievalException(RetrievalErrorCode.INVALID_REQUEST, + "graphVersion and indexVersion are required for structured retrieval"); + } + } + + private static EvidenceBuilder builder(Map merged, Map chunks, + String chunkId, String documentId, EvidenceKind kind) { + TextChunk chunk = chunks.get(chunkId); + if (chunk == null || !chunk.getDocumentId().equals(documentId)) { + throw notReady("channel chunk does not match fixture"); + } + String id = documentId.length() + ":" + documentId + chunkId; + return merged.computeIfAbsent(id, key -> new EvidenceBuilder(id, chunk, kind)); + } + + private static final class EvidenceBuilder { + private final String id; + private final TextChunk chunk; + private final EvidenceKind kind; + private final Map scores = new LinkedHashMap<>(); + private final List entities = new ArrayList<>(); + private final List paths = new ArrayList<>(); + private final List graphTraces = new ArrayList<>(); + + private EvidenceBuilder(String id, TextChunk chunk, EvidenceKind kind) { + this.id = id; + this.chunk = chunk; + this.kind = kind; + } + + private Evidence build() { + double fused = 0.0; + for (ChannelScore score : scores.values()) { + fused += 1.0 / (60.0 + score.getRank()); + } + SourceRef source = new SourceRef(chunk.getDocumentId(), chunk.getDocumentId(), chunk.getStartOffset(), chunk.getEndOffset()); + return new Evidence(id, kind, chunk.getText(), Collections.singletonList(chunk), entities, paths, + Collections.singletonList(source), scores, fused, null, 1, graphTraces); + } + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalFixtureRegistry.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalFixtureRegistry.java new file mode 100644 index 000000000..95811706d --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalFixtureRegistry.java @@ -0,0 +1,408 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.service; + +import java.io.IOException; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.channel.graph.GraphNeighborProvider; +import org.apache.geaflow.ai.retrieval.channel.graph.InMemoryGraphNeighborProvider; +import org.apache.geaflow.ai.retrieval.index.bm25.Bm25IndexReader; +import org.apache.geaflow.ai.retrieval.index.graph.GraphArtifactReader; +import org.apache.geaflow.ai.retrieval.index.vector.VectorIndexReader; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; + +/** Session-free registry for an atomically published GraphRAG fixture. */ +public final class RetrievalFixtureRegistry implements AutoCloseable { + private final Map fixtures = new HashMap<>(); + private boolean closed; + + public synchronized void register(String graphName, String graphVersion, String indexVersion, + Bm25IndexReader bm25, VectorIndexReader vector, + List chunks, List entities, + List edges, boolean ready) { + register(graphName, graphVersion, indexVersion, bm25, vector, chunks, entities, edges, ready, + new InMemoryGraphNeighborProvider(edges), Collections.emptyMap()); + } + + private synchronized void register(String graphName, String graphVersion, String indexVersion, + Bm25IndexReader bm25, VectorIndexReader vector, + List chunks, List entities, List edges, + boolean ready, GraphNeighborProvider graphProvider, + Map failures) { + boolean registered = false; + try { + requireVersion(graphName, graphVersion, indexVersion); + if (chunks == null || entities == null || edges == null + || graphProvider == null || (bm25 == null && vector == null && graphProvider == null)) { + throw notReady("fixture contains no usable channel or incomplete graph fixture"); + } + java.util.Set chunkIds = new java.util.HashSet<>(); + for (TextChunk chunk : chunks) { + if (chunk == null || !chunkIds.add(chunk.getChunkId())) { + throw notReady("invalid or duplicate fixture chunk"); + } + } + if (bm25 != null) { + bm25.getIdentity().validate(graphName, graphVersion, indexVersion, chunks); + } + if (vector != null) { + vector.getIdentity().validate(graphName, graphVersion, indexVersion, chunks); + } + ensureOpen(); + String fixtureKey = key(graphName, graphVersion, indexVersion); + Fixture replacement = new Fixture(graphName, graphVersion, indexVersion, bm25, vector, + chunks, entities, edges, graphProvider, ready, failures); + Fixture previous = fixtures.put(fixtureKey, replacement); + registered = true; + if (previous != null) { + previous.retire(); + } + } catch (RuntimeException error) { + // Registration takes ownership of readers. Release them if validation or publication fails. + if (!registered) { + closeReaders(bm25, vector, error); + } + throw error; + } + } + + /** Loads channels independently, preserving typed failures for explicit partial retrieval. */ + public void load(String graphName, String graphVersion, String indexVersion, + Path bm25Artifact, Path vectorArtifact, List chunks, + List entities, List edges, + String vectorSource, String vectorVersion) { + requireVersion(graphName, graphVersion, indexVersion); + Map failures = new LinkedHashMap<>(); + Bm25IndexReader bm25 = null; + VectorIndexReader vector = null; + try { + bm25 = new Bm25IndexReader(bm25Artifact, graphVersion, indexVersion); + bm25.getIdentity().validate(graphName, graphVersion, indexVersion, chunks); + } catch (RetrievalException error) { + if (bm25 != null) { + try { + bm25.close(); + } catch (IOException closeError) { + error.addSuppressed(closeError); + } + bm25 = null; + } + failures.put("BM25", error); + } + try { + vector = new VectorIndexReader(vectorArtifact, graphVersion, indexVersion, vectorSource, vectorVersion); + vector.getIdentity().validate(graphName, graphVersion, indexVersion, chunks); + } catch (RetrievalException error) { + if (vector != null) { + vector.close(); + } + vector = null; + failures.put("VECTOR", error); + } + if (bm25 == null && vector == null) { + throw notReady("no usable BM25 or vector artifact"); + } + register(graphName, graphVersion, indexVersion, bm25, vector, chunks, entities, edges, true, + new InMemoryGraphNeighborProvider(edges), failures); + } + + public void loadGraph(String graphName, String graphVersion, String indexVersion, Path graphArtifact, + List chunks) { + GraphArtifactReader reader = new GraphArtifactReader(graphArtifact, graphVersion, indexVersion); + try { + reader.getIdentity().validate(graphName, graphVersion, indexVersion, chunks); + reader.validateReferences(chunks); + register(graphName, graphVersion, indexVersion, null, null, chunks, reader.getEntities(), + reader.getEdges(), true, reader.getNeighbors(), Collections.emptyMap()); + } catch (RuntimeException error) { + reader.close(); + throw error; + } + } + + public synchronized Fixture get(String graphName, String graphVersion, String indexVersion) { + requireVersion(graphName, graphVersion, indexVersion); + Fixture fixture = fixtures.get(key(graphName, graphVersion, indexVersion)); + if (fixture == null || !fixture.isReady()) { + throw notReady("fixture is missing or not READY"); + } + return fixture; + } + + /** Acquires a lease that keeps the selected fixture readers open until released. */ + public synchronized Fixture.Lease acquire(String graphName, String graphVersion, String indexVersion) { + return get(graphName, graphVersion, indexVersion).acquire(); + } + + public synchronized void clear() { + RuntimeException failure = null; + for (Fixture fixture : fixtures.values()) { + try { + fixture.retire(); + } catch (RuntimeException error) { + if (failure == null) { + failure = new IllegalStateException("failed to close fixture readers", error); + } else { + failure.addSuppressed(error); + } + } + } + fixtures.clear(); + if (failure != null) { + throw failure; + } + } + + /** Closes all published readers and prevents further registrations. */ + @Override + public synchronized void close() { + if (!closed) { + closed = true; + clear(); + } + } + + private void ensureOpen() { + if (closed) { + throw new IllegalStateException("fixture registry is closed"); + } + } + + private static void closeReaders(Bm25IndexReader bm25, VectorIndexReader vector, + Exception failure) { + Exception closeFailure = null; + if (bm25 != null) { + try { + bm25.close(); + } catch (IOException closeError) { + failure.addSuppressed(closeError); + closeFailure = closeError; + } + } + if (vector != null) { + try { + vector.close(); + } catch (RuntimeException closeError) { + failure.addSuppressed(closeError); + if (closeFailure != null) { + closeFailure.addSuppressed(closeError); + } + } + } + } + + private static void requireVersion(String graphName, String graphVersion, String indexVersion) { + if (blank(graphName) || blank(graphVersion) || blank(indexVersion) + || !graphVersion.equals(indexVersion)) { + throw notReady("graph/index versions are missing or inconsistent"); + } + } + + private static String key(String graphName, String graphVersion, String indexVersion) { + return graphName + "\u0000" + graphVersion + "\u0000" + indexVersion; + } + + private static boolean blank(String value) { + return value == null || value.trim().isEmpty(); + } + + private static RetrievalException notReady(String message) { + return new RetrievalException(RetrievalErrorCode.INDEX_NOT_READY, message); + } + + /** Immutable published fixture. */ + public static final class Fixture { + private final String graphName; + private final String graphVersion; + private final String indexVersion; + private final Bm25IndexReader bm25; + private final VectorIndexReader vector; + private final List chunks; + private final Map chunksById; + private final List entities; + private final List edges; + private final GraphNeighborProvider graphProvider; + private final boolean ready; + private final Map failures; + private boolean closed; + private boolean retired; + private int activeLeases; + + private Fixture(String graphName, String graphVersion, String indexVersion, + Bm25IndexReader bm25, VectorIndexReader vector, List chunks, + List entities, List edges, GraphNeighborProvider graphProvider, + boolean ready, + Map failures) { + this.graphName = graphName; + this.graphVersion = graphVersion; + this.indexVersion = indexVersion; + this.bm25 = bm25; + this.vector = vector; + this.chunks = immutable(chunks); + Map indexedChunks = new HashMap<>(); + for (TextChunk chunk : this.chunks) { + indexedChunks.put(chunk.getChunkId(), chunk); + } + this.chunksById = Collections.unmodifiableMap(indexedChunks); + this.entities = immutable(entities); + this.edges = immutable(edges); + this.graphProvider = graphProvider; + this.ready = ready; + this.failures = Collections.unmodifiableMap(new LinkedHashMap<>(failures)); + } + + private static List immutable(List values) { + return Collections.unmodifiableList(new ArrayList<>(values)); + } + + public String getGraphName() { + return graphName; + } + + public String getGraphVersion() { + return graphVersion; + } + + public String getIndexVersion() { + return indexVersion; + } + + public Bm25IndexReader getBm25() { + return bm25; + } + + public VectorIndexReader getVector() { + return vector; + } + + public List getChunks() { + return chunks; + } + + public Map getChunksById() { + return chunksById; + } + + public List getEntities() { + return entities; + } + + public List getEdges() { + return edges; + } + + public GraphNeighborProvider getGraphProvider() { + return graphProvider; + } + + public Map getFailures() { + return failures; + } + + public boolean isReady() { + return ready; + } + + private synchronized Lease acquire() { + if (retired || closed) { + throw new IllegalStateException("fixture is retired"); + } + activeLeases++; + return new Lease(this); + } + + /** Retires the fixture, deferring reader closure while leases remain active. */ + private synchronized void retire() { + retired = true; + closeIfUnused(); + } + + private synchronized void release() { + if (activeLeases <= 0) { + throw new IllegalStateException("fixture lease released more than once"); + } + activeLeases--; + closeIfUnused(); + } + + private void closeIfUnused() { + if (!retired || closed || activeLeases != 0) { + return; + } + closed = true; + Exception failure = null; + if (bm25 != null) { + try { + bm25.close(); + } catch (IOException | RuntimeException error) { + failure = error; + } + } + if (vector != null) { + try { + vector.close(); + } catch (RuntimeException error) { + if (failure == null) { + failure = error; + } else { + failure.addSuppressed(error); + } + } + } + if (failure != null) { + throw new IllegalStateException("failed to close fixture readers", failure); + } + } + + /** A scoped reader lease held for the duration of one retrieval. */ + public static final class Lease implements AutoCloseable { + private Fixture fixture; + + private Lease(Fixture fixture) { + this.fixture = fixture; + } + + public Fixture getFixture() { + if (fixture == null) { + throw new IllegalStateException("fixture lease is closed"); + } + return fixture; + } + + @Override + public void close() { + Fixture leased = fixture; + if (leased != null) { + fixture = null; + leased.release(); + } + } + } + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/channel/graph/EntityAnchorResolverTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/channel/graph/EntityAnchorResolverTest.java new file mode 100644 index 000000000..39546a09f --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/channel/graph/EntityAnchorResolverTest.java @@ -0,0 +1,69 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.channel.graph; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.junit.jupiter.api.Test; + +class EntityAnchorResolverTest { + private final EntityAnchorResolver resolver = new EntityAnchorResolver(); + + @Test + void resolvesCanonicalNamesAndAliasesWithStableConfidenceOrdering() { + EntityRef canonicalZ = new EntityRef("z", "Confucius", "person"); + EntityRef canonicalA = new EntityRef("a", "Ｃｏｎｆｕｃｉｕｓ", "person"); + EntityRef alias = new EntityRef("alias", "Master Kong", Collections.singletonList("kongzi"), + "person", Collections.emptyList()); + + List anchors = resolver.resolve(" CONFUCIUS ", + Arrays.asList(canonicalZ, alias, canonicalA)); + assertEquals(Arrays.asList("a", "z"), entityIds(anchors)); + assertEquals("CANONICAL", anchors.get(0).getMatchType()); + assertEquals(1.0, anchors.get(0).getConfidence()); + + anchors = resolver.resolve(" KONGZI ", Collections.singletonList(alias)); + assertEquals(Collections.singletonList("alias"), entityIds(anchors)); + assertEquals("ALIAS", anchors.get(0).getMatchType()); + assertEquals(0.9, anchors.get(0).getConfidence()); + } + + @Test + void rejectsPartialAndUnknownMatchesAndReturnsNoAnchorsAfterDeadline() { + EntityRef entity = new EntityRef("e1", "Confucius", Collections.singletonList("kongzi"), + "person", Collections.emptyList()); + assertTrue(resolver.resolve("Confucius and astronomy", Collections.singletonList(entity)).isEmpty()); + assertTrue(resolver.resolve("unknown", Collections.singletonList(entity)).isEmpty()); + assertTrue(resolver.resolve("Confucius", Collections.singletonList(entity), 1L, () -> 1L).isEmpty()); + } + + private static List entityIds(List anchors) { + java.util.ArrayList ids = new java.util.ArrayList<>(); + for (EntityAnchorResolver.Anchor anchor : anchors) { + ids.add(anchor.getEntity().getEntityId()); + } + return ids; + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphRetrieverTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphRetrieverTest.java new file mode 100644 index 000000000..b79c3b679 --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/channel/graph/GraphRetrieverTest.java @@ -0,0 +1,122 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.channel.graph; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.AbstractList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; +import org.junit.jupiter.api.Test; + +class GraphRetrieverTest { + + @Test + void edgeCapStopsNeighborIterationWithoutCopyingOrSortingTheWholeList() { + AtomicInteger reads = new AtomicInteger(); + List neighbors = new AbstractList() { + @Override + public GraphEdgeRef get(int index) { + reads.incrementAndGet(); + if (index >= 2) { + throw new AssertionError("retriever read beyond edge cap"); + } + return edge(index); + } + + @Override + public int size() { + return 100_000; + } + }; + GraphNeighborProvider provider = entityId -> neighbors; + Map chunks = new HashMap<>(); + chunks.put("c1", new TextChunk("c1", "doc-1", 0, 0, 1, 1, "Confucius")); + List entities = Arrays.asList( + new EntityRef("e1", "Confucius", "person"), + new EntityRef("e2", "Ethics", "topic")); + GraphRetriever retriever = new GraphRetriever(); + + List hits = retriever.retrieve("Confucius", entities, provider, + chunks, 10, 2, 0L, System::nanoTime); + + assertEquals(2, hits.size()); + assertEquals(2, reads.get()); + assertEquals(2, retriever.getLastStats().getEdgesExamined()); + } + + @Test + void expiredDeadlineDoesNotReadNeighbors() { + AtomicInteger reads = new AtomicInteger(); + GraphNeighborProvider provider = entityId -> new AbstractList() { + @Override + public GraphEdgeRef get(int index) { + reads.incrementAndGet(); + return edge(index); + } + + @Override + public int size() { + return 100_000; + } + }; + GraphRetriever retriever = new GraphRetriever(); + List entities = Collections.singletonList(new EntityRef("e1", "Confucius", "person")); + + retriever.retrieve("Confucius", entities, provider, Collections.emptyMap(), + 10, 10, 1L, () -> 1L); + + assertEquals(0, reads.get()); + } + + @Test + void duplicateProvenanceScanStopsAtCandidateBudget() { + List provenance = new java.util.ArrayList<>(); + for (int index = 0; index < 1000; index++) { + provenance.add("c1"); + } + GraphEdgeRef edge = new GraphEdgeRef("edge-1", "supports", "e1", "e2", provenance); + Map chunks = new HashMap<>(); + chunks.put("c1", new TextChunk("c1", "doc-1", 0, 0, 1, 1, "Confucius")); + List entities = Arrays.asList( + new EntityRef("e1", "Confucius", "person"), + new EntityRef("e2", "Ethics", "topic")); + GraphRetriever retriever = new GraphRetriever(); + + List hits = retriever.retrieve("Confucius", entities, + Collections.singletonList(edge), chunks, 4, 10, 0L); + + assertEquals(1, hits.size()); + assertEquals(1, retriever.getLastStats().getCandidatesProduced()); + assertEquals(RecallStopReason.CANDIDATE_LIMIT, retriever.getLastStats().getStopReason()); + } + + private static GraphEdgeRef edge(int index) { + return new GraphEdgeRef(String.format("edge-%06d", index), "supports", "e1", "e2", + Collections.singletonList("c1")); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfFusionTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfFusionTest.java new file mode 100644 index 000000000..4a5cef1fb --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/fusion/WeightedRrfFusionTest.java @@ -0,0 +1,97 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ +package org.apache.geaflow.ai.retrieval.fusion; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import org.apache.geaflow.ai.retrieval.model.evidence.ChannelScore; +import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; +import org.apache.geaflow.ai.retrieval.model.evidence.EvidenceKind; +import org.junit.jupiter.api.Test; + +class WeightedRrfFusionTest { + @Test + void computesWeightedRrfAndUsesStableTieBreak() { + WeightedRrfConfig config = new WeightedRrfConfig(1.0, 1.0, 60); + WeightedRrfFusion fusion = new WeightedRrfFusion(config); + Evidence both = evidence("c", 1, 2); + Evidence bm25 = evidence("b", 2, 0); + Evidence vector = evidence("a", 0, 1); + + assertEquals(1.0 / 61.0 + 1.0 / 62.0, fusion.score(both), 1.0e-12); + assertEquals(Arrays.asList("c", "a", "b"), ids(fusion.sort( + Arrays.asList(bm25, vector, both)))); + } + + @Test + void zeroWeightDoesNotAffectOrderingAndConfigurationMustBeValid() { + WeightedRrfFusion fusion = new WeightedRrfFusion(new WeightedRrfConfig(1.0, 0.0, 60)); + Evidence vectorOnly = evidence("a", 0, 1); + Evidence bm25 = evidence("z", 3, 1); + assertEquals(0.0, fusion.score(vectorOnly)); + assertEquals("z", fusion.sort(Arrays.asList(vectorOnly, bm25)).get(0).getEvidenceId()); + assertThrows(IllegalArgumentException.class, + () -> new WeightedRrfConfig(0.0, 0.0, 60)); + assertThrows(IllegalArgumentException.class, + () -> new WeightedRrfConfig(-1.0, 1.0, 60)); + } + + @Test + void unequalWeightsActualTiesEmptyResultsAndLargeConstant() { + WeightedRrfFusion fusion = new WeightedRrfFusion(new WeightedRrfConfig(2.0, 0.5, 10)); + Evidence both = evidence("both", 1, 2); + assertEquals(2.0 / 11 + 0.5 / 12, fusion.score(both), 1.0e-12); + java.util.List tied = fusion.sort(Arrays.asList(evidence("z", 1, 0), evidence("a", 1, 0))); + assertEquals(Arrays.asList("a", "z"), ids(tied)); + assertEquals(1, tied.get(0).getRank()); + assertEquals(2, tied.get(1).getRank()); + assertEquals(Collections.emptyList(), fusion.sort(Collections.emptyList())); + assertEquals(both.getStageScores(), fusion.sort(Collections.singletonList(both)).get(0).getStageScores()); + assertEquals(1.0 / ((double) Integer.MAX_VALUE + 1), + new WeightedRrfFusion(new WeightedRrfConfig(1, 0, Integer.MAX_VALUE)).score(evidence("a", 1, 0)), 1.0e-20); + } + + private static Evidence evidence(String id, int bm25Rank, int vectorRank) { + Map scores = new LinkedHashMap<>(); + if (bm25Rank > 0) { + scores.put("BM25", new ChannelScore("BM25", 10.0, null, bm25Rank)); + } + if (vectorRank > 0) { + scores.put("VECTOR", new ChannelScore("VECTOR", 0.9, null, vectorRank)); + } + double fused = (bm25Rank == 0 ? 0 : 1.0 / (60 + bm25Rank)) + + (vectorRank == 0 ? 0 : 1.0 / (60 + vectorRank)); + return new Evidence(id, EvidenceKind.CHUNK, id, Collections.emptyList(), + Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), scores, + fused, 1); + } + + private static java.util.List ids(java.util.List values) { + java.util.List result = new java.util.ArrayList<>(); + for (Evidence value : values) { + result.add(value.getEvidenceId()); + } + return result; + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/GraphRagMvpEndToEndTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/GraphRagMvpEndToEndTest.java new file mode 100644 index 000000000..91f824e4a --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/GraphRagMvpEndToEndTest.java @@ -0,0 +1,752 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 + */ +package org.apache.geaflow.ai.retrieval.service; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicLong; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalRequest; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse; +import org.apache.geaflow.ai.retrieval.codec.RetrievalApiJson; +import org.apache.geaflow.ai.retrieval.execution.RecallStageStatus; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; +import org.apache.geaflow.ai.retrieval.fusion.WeightedRrfConfig; +import org.apache.geaflow.ai.retrieval.index.bm25.Bm25IndexReader; +import org.apache.geaflow.ai.retrieval.index.bm25.LuceneBm25IndexBuilder; +import org.apache.geaflow.ai.retrieval.index.vector.OfflineVectorIndexBuilder; +import org.apache.geaflow.ai.retrieval.index.vector.VectorIndexReader; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.metadata.ChunkingConfiguration; +import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class GraphRagMvpEndToEndTest { + private static final String GRAPH = "mvp-graph"; + private static final String VERSION = "v1"; + + @TempDir + Path tempDir; + private final List readers = new ArrayList<>(); + + @AfterEach + void closeReaders() throws Exception { + for (Bm25IndexReader reader : readers) { + reader.close(); + } + } + + @Test + void executesAllMvpModesAndPreservesEvidence() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.registry); + + RetrievalResponse keyword = service.retrieve(request("KEYWORD", null)); + assertEquals("BM25_ONLY", keyword.getTrace().getSelectedMode().name()); + assertFalse(keyword.getEvidence().isEmpty()); + + RetrievalResponse bm25 = service.retrieve(request("BM25_ONLY", null)); + RetrievalResponse vector = service.retrieve(request("VECTOR_ONLY", Arrays.asList(1.0, 0.0))); + RetrievalResponse graph = service.retrieve(request("GRAPH_ONLY", null)); + RetrievalResponse hybrid = service.retrieve(request("HYBRID", Arrays.asList(1.0, 0.0))); + + assertFalse(bm25.getEvidence().isEmpty()); + assertFalse(vector.getEvidence().isEmpty()); + assertFalse(graph.getEvidence().isEmpty()); + RetrievalRequest aliasRequest = request("GRAPH_ONLY", null); + aliasRequest.setQuery("kongzi"); + assertFalse(service.retrieve(aliasRequest).getEvidence().isEmpty()); + assertFalse(hybrid.getEvidence().isEmpty()); + Evidence merged = find(hybrid.getEvidence(), "5:doc-3c3"); + assertTrue(merged.getStageScores().containsKey("BM25")); + assertTrue(merged.getStageScores().containsKey("VECTOR")); + assertTrue(merged.getStageScores().containsKey("GRAPH")); + assertEquals(1.0 / (60 + merged.getStageScores().get("BM25").getRank()) + + 1.0 / (60 + merged.getStageScores().get("VECTOR").getRank()) + + 1.0 / (60 + merged.getStageScores().get("GRAPH").getRank()), + merged.getFusedScore(), 1.0e-12); + assertFalse(merged.getPaths().isEmpty()); + assertFalse(hybrid.getPaths().isEmpty()); + assertFalse(hybrid.getSources().isEmpty()); + } + + @Test + void rejectsVersionsReadinessAndUnsupportedExecution() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.registry); + RetrievalRequest wrongVersion = request("BM25_ONLY", null); + wrongVersion.setIndexVersion("v2"); + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> service.retrieve(wrongVersion)); + + RetrievalRequest parallel = request("BM25_ONLY", null); + parallel.setExecutionMode("PARALLEL"); + assertCode(RetrievalErrorCode.UNSUPPORTED_OPTION, () -> service.retrieve(parallel)); + + RetrievalRequest badVector = request("VECTOR_ONLY", Collections.singletonList(1.0)); + assertCode(RetrievalErrorCode.INVALID_REQUEST, () -> service.retrieve(badVector)); + + RetrievalFixtureRegistry notReady = new RetrievalFixtureRegistry(); + notReady.register(GRAPH, VERSION, VERSION, fixture.bm25, fixture.vector, + fixture.chunks, fixture.entities, fixture.edges, false); + assertCode(RetrievalErrorCode.INDEX_NOT_READY, + () -> new RecallService(notReady).retrieve(request("BM25_ONLY", null))); + } + + @Test + void requiresStructuredVersions() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.registry); + + RetrievalRequest missingGraphVersion = request("BM25_ONLY", null); + missingGraphVersion.setGraphVersion(null); + assertCode(RetrievalErrorCode.INVALID_REQUEST, () -> service.retrieve(missingGraphVersion)); + + RetrievalRequest missingIndexVersion = request("BM25_ONLY", null); + missingIndexVersion.setIndexVersion(null); + assertCode(RetrievalErrorCode.INVALID_REQUEST, () -> service.retrieve(missingIndexVersion)); + + } + + @Test + void enforcesCandidateBudgetAndNoAnchorScan() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.registry); + RetrievalRequest request = request("HYBRID", Arrays.asList(1.0, 0.0)); + request.setBudget(new RetrievalBudget(1, 1000, 1, 100)); + assertCode(RetrievalErrorCode.INVALID_REQUEST, () -> service.retrieve(request)); + request.setBudget(new RetrievalBudget(1, 1000, 3, 100)); + RetrievalResponse response = service.retrieve(request); + assertEquals(1, response.getEvidence().size()); + assertEquals(1, response.getTrace().getCandidateCounts().get("BM25")); + assertEquals(1, response.getTrace().getCandidateCounts().get("VECTOR")); + assertEquals(1, response.getTrace().getCandidateCounts().get("GRAPH")); + assertEquals(3, response.getTrace().getTotalCandidatesEvaluated()); + assertEquals(Arrays.asList("BM25", "VECTOR", "GRAPH"), + response.getTrace().getSelectedChannels()); + assertEquals(response.getRequestId(), response.getTrace().getRequestId()); + + RetrievalRequest noAnchor = request("GRAPH_ONLY", null); + noAnchor.setQuery("unknown entity"); + RetrievalResponse empty = service.retrieve(noAnchor); + assertTrue(empty.getEvidence().isEmpty()); + assertEquals(0, empty.getTrace().getAnchorsConsidered()); + assertEquals(0, empty.getTrace().getEdgesExamined()); + } + + @Test + void scalesHybridRecallAcrossTenThousandChunks() throws Exception { + List chunks = new ArrayList<>(); + Map vectors = new HashMap<>(); + List graphEvidenceChunkIds = new ArrayList<>(); + for (int index = 0; index < 10_000; index++) { + String chunkId = String.format("scaled-c%05d", index); + String documentId = String.format("scaled-doc-%05d", index); + String text = index % 5 == 0 + ? "Confucius discusses ethics in document " + index + : "Background material about history in document " + index; + chunks.add(new TextChunk(chunkId, documentId, 0, 0, text.length(), 6, text)); + vectors.put(chunkId, index % 2 == 0 ? new float[] {1.0F, 0.0F} : new float[] {0.0F, 1.0F}); + if (index % 5 == 0) { + graphEvidenceChunkIds.add(chunkId); + } + } + IngestionContext context = context(); + Path bm25Path = tempDir.resolve("scaled-bm25"); + Path vectorPath = tempDir.resolve("scaled-vector"); + try (org.apache.geaflow.ai.retrieval.index.IndexArtifact ignored = + new LuceneBm25IndexBuilder(bm25Path).build(context, chunks); + org.apache.geaflow.ai.retrieval.index.IndexArtifact ignoredVector = + new OfflineVectorIndexBuilder(vectorPath, vectors, "fixture", VERSION) + .build(context, chunks)) { + // Artifacts are published before readers are opened. + } + Bm25IndexReader bm25 = new Bm25IndexReader(bm25Path.resolve("bm25-v1"), VERSION, VERSION); + readers.add(bm25); + VectorIndexReader vector = new VectorIndexReader(vectorPath.resolve("vector-v1.bin"), + VERSION, VERSION, "fixture", VERSION); + List entities = Arrays.asList( + new EntityRef("scaled-e1", "Confucius", "person"), + new EntityRef("scaled-e2", "Ethics", "topic")); + List edges = Collections.singletonList(new GraphEdgeRef( + "scaled-edge-1", "supports", "scaled-e1", "scaled-e2", graphEvidenceChunkIds)); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + registry.register(GRAPH, VERSION, VERSION, bm25, vector, chunks, + entities, edges, true); + + RetrievalRequest request = request("HYBRID", Arrays.asList(1.0, 0.0)); + request.setBudget(new RetrievalBudget(10, 5000, 400, 100)); + RecallService service = new RecallService(registry); + RetrievalResponse first = service.retrieve(request); + RetrievalResponse second = service.retrieve(request); + + assertEquals(Arrays.asList("BM25", "VECTOR", "GRAPH"), first.getTrace().getSelectedChannels()); + assertEquals(134, first.getTrace().getChannelBudgets().get("BM25")); + assertEquals(133, first.getTrace().getChannelBudgets().get("VECTOR")); + assertEquals(133, first.getTrace().getChannelBudgets().get("GRAPH")); + assertTrue(first.getTrace().getEvaluatedCounts().get("BM25") <= 134); + assertTrue(first.getTrace().getEvaluatedCounts().get("VECTOR") <= 133); + assertTrue(first.getTrace().getEvaluatedCounts().get("GRAPH") <= 133); + assertEquals(133, first.getTrace().getCandidateCounts().get("GRAPH")); + assertEquals(133, first.getTrace().getGraphCandidatesProduced()); + assertTrue(first.getTrace().getTotalCandidatesEvaluated() <= 400); + assertTrue(first.getEvidence().size() <= 10); + assertEquals(first.getEvidence(), second.getEvidence()); + assertEquals(first.getTrace().getCandidateCounts(), second.getTrace().getCandidateCounts()); + } + + @Test + void producesStableEvidenceAndTraceCounters() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.registry); + RetrievalRequest request = request("HYBRID", Arrays.asList(1.0, 0.0)); + RetrievalResponse first = service.retrieve(request); + RetrievalResponse second = service.retrieve(request); + assertEquals(first.getEvidence(), second.getEvidence()); + assertEquals(first.getTrace().getCandidateCounts(), second.getTrace().getCandidateCounts()); + assertEquals(first.getTrace().getStopReason(), second.getTrace().getStopReason()); + assertEquals(first.getTrace().getSelectedChannels(), second.getTrace().getSelectedChannels()); + assertEquals(1.0, first.getTrace().getBm25Weight()); + assertEquals(1.0, first.getTrace().getVectorWeight()); + assertEquals(60, first.getTrace().getRrfRankConstant()); + } + + @Test + void supportsExplicitPartialHybridWhenOneChannelIsUnavailable() throws Exception { + Fixture fixture = fixture(); + RetrievalFixtureRegistry partialRegistry = new RetrievalFixtureRegistry(); + partialRegistry.register(GRAPH, VERSION, VERSION, null, fixture.vector, + fixture.chunks, fixture.entities, fixture.edges, true); + RecallService service = new RecallService(partialRegistry); + RetrievalRequest request = request("HYBRID", Arrays.asList(1.0, 0.0)); + + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> service.retrieve(request)); + request.setAllowPartialResults(true); + RetrievalResponse response = service.retrieve(request); + assertFalse(response.getEvidence().isEmpty()); + assertEquals(Arrays.asList("BM25", "GRAPH"), response.getDegradedChannels()); + assertEquals("artifact unavailable", + response.getTrace().getDegradationReasons().get("BM25")); + } + + @Test + void replacementDefersReaderCloseUntilActiveRetrievalLeaseEnds() throws Exception { + Fixture fixture = fixture(); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + Bm25IndexReader oldReader = new Bm25IndexReader(tempDir.resolve("bm25/bm25-v1"), VERSION, VERSION); + Bm25IndexReader replacementReader = new Bm25IndexReader( + tempDir.resolve("bm25/bm25-v1"), VERSION, VERSION); + registry.register(GRAPH, VERSION, VERSION, oldReader, null, + fixture.chunks, fixture.entities, fixture.edges, true); + + try (RetrievalFixtureRegistry.Fixture.Lease lease = registry.acquire(GRAPH, VERSION, VERSION)) { + registry.register(GRAPH, VERSION, VERSION, replacementReader, null, + fixture.chunks, fixture.entities, fixture.edges, true); + assertFalse(lease.getFixture().getBm25().search("Confucius", 3, 3).isEmpty()); + } finally { + registry.close(); + } + + assertThrows(RetrievalException.class, () -> oldReader.search("Confucius", 3, 3)); + } + + @Test + void clearDefersReaderCloseUntilActiveRetrievalLeaseEnds() throws Exception { + Fixture fixture = fixture(); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + Bm25IndexReader reader = new Bm25IndexReader(tempDir.resolve("bm25/bm25-v1"), VERSION, VERSION); + registry.register(GRAPH, VERSION, VERSION, reader, null, + fixture.chunks, fixture.entities, fixture.edges, true); + + try (RetrievalFixtureRegistry.Fixture.Lease lease = registry.acquire(GRAPH, VERSION, VERSION)) { + registry.clear(); + assertFalse(lease.getFixture().getBm25().search("Confucius", 3, 3).isEmpty()); + } + + assertThrows(RetrievalException.class, () -> reader.search("Confucius", 3, 3)); + registry.close(); + } + + @Test + void recordsDeadlineOnTheChannelThatWasStopped() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.bm25, fixture.vector, + fixture.chunks, fixture.entities, fixture.edges); + RetrievalResponse response = service.retrieveWithTrace("Confucius", + new float[] {1.0F, 0.0F}, "HYBRID", 4, 4, System.nanoTime() - 1, true); + + assertEquals("DEADLINE", response.getTrace().getStopReason()); + assertEquals(Arrays.asList("BM25", "VECTOR", "GRAPH"), response.getTrace().getDegradedChannels()); + assertEquals(RecallStopReason.DEADLINE, response.getTrace().getChannelStopReasons().get("BM25")); + assertEquals(RecallStageStatus.NOT_RUN, response.getTrace().getChannelStatuses().get("VECTOR")); + assertEquals(0, response.getTrace().getCandidateCounts().get("VECTOR")); + } + + @Test + void customRrfConfigurationIsUsedByHybrid() throws Exception { + Fixture fixture = fixture(); + RetrievalRequest request = request("HYBRID", Arrays.asList(1.0, 0.0)); + request.setQuery("Astronomy"); + RecallService defaultService = new RecallService(fixture.registry); + RecallService vectorHeavy = new RecallService(fixture.registry, + new WeightedRrfConfig(0.0, 1.0, 60)); + + RetrievalResponse defaultResponse = defaultService.retrieve(request); + RetrievalResponse weightedResponse = vectorHeavy.retrieve(request); + assertEquals(1.0, defaultResponse.getTrace().getBm25Weight()); + assertEquals(0.0, weightedResponse.getTrace().getBm25Weight()); + assertEquals(1.0, weightedResponse.getTrace().getVectorWeight()); + assertEquals("5:doc-1c1", weightedResponse.getEvidence().get(0).getEvidenceId()); + assertEquals(1.0 / (60 + weightedResponse.getEvidence().get(0) + .getStageScores().get("VECTOR").getRank()), + weightedResponse.getEvidence().get(0).getFusedScore(), 1.0e-12); + assertTrue(defaultResponse.getEvidence().get(0).getEvidenceId() + .equals("5:doc-2c2") || defaultResponse.getEvidence().get(0).getEvidenceId() + .equals("5:doc-3c3")); + } + + @Test + void malformedVectorsAreRejectedEvenBeforeExpiredOrMissingChannel() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.registry); + for (List value : Arrays.asList(Collections.emptyList(), Arrays.asList(0.0, 0.0), + Arrays.asList(Double.NaN, 1.0), Arrays.asList(Double.MAX_VALUE, 1.0), Collections.singletonList(1.0))) { + RetrievalRequest invalid = request("HYBRID", value); + invalid.setAllowPartialResults(true); + assertCode(RetrievalErrorCode.INVALID_REQUEST, () -> service.retrieve(invalid)); + } + RecallService direct = new RecallService(fixture.bm25, fixture.vector, fixture.chunks, fixture.entities, fixture.edges); + assertCode(RetrievalErrorCode.INVALID_REQUEST, () -> direct.retrieveWithTrace("query", + new float[] {1.0F}, "HYBRID", 1, 2, System.nanoTime() - 1, true)); + assertCode(RetrievalErrorCode.UNSUPPORTED_OPTION, () -> service.retrieveWithTrace("query", + null, "ADAPTIVE", 1, 2, 0L)); + } + + @Test + void strictDeadlineFailsAndPartialDeadlineRetainsCompletedEvidence() throws Exception { + Fixture fixture = fixture(); + AtomicLong time = new AtomicLong(); + RecallService service = new RecallService(fixture.registry, WeightedRrfConfig.DEFAULT, + () -> time.getAndAdd(1_000_000L)); + RetrievalRequest timed = request("HYBRID", Arrays.asList(1.0, 0.0)); + timed.setBudget(new RetrievalBudget(2, 7, 4, 100)); + assertCode(RetrievalErrorCode.RETRIEVAL_TIMEOUT, () -> service.retrieve(timed)); + time.set(0L); + timed.setAllowPartialResults(true); + RetrievalResponse partial = service.retrieve(timed); + assertFalse(partial.getEvidence().isEmpty()); + assertEquals("DEADLINE", partial.getTrace().getStopReason()); + assertEquals(RecallStageStatus.DEGRADED, partial.getTrace().getChannelStatuses().get("BM25")); + assertEquals(RecallStageStatus.NOT_RUN, partial.getTrace().getChannelStatuses().get("VECTOR")); + assertEquals(0, partial.getTrace().getEvaluatedCounts().get("VECTOR")); + assertTrue(partial.getTrace().getTotalCandidatesEvaluated() <= 4); + } + + @Test + void corruptVectorCanDegradeWithoutLosingBm25Evidence() throws Exception { + Fixture fixture = fixture(); + Path broken = tempDir.resolve("vector/vector-v1.bin"); + Files.write(broken, new byte[] {0, 1, 2}); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + registry.load(GRAPH, VERSION, VERSION, tempDir.resolve("bm25/bm25-v1"), broken, + fixture.chunks, fixture.entities, fixture.edges, "fixture", VERSION); + readers.add(registry.get(GRAPH, VERSION, VERSION).getBm25()); + RecallService service = new RecallService(registry); + RetrievalRequest hybrid = request("HYBRID", Arrays.asList(1.0, 0.0)); + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> service.retrieve(hybrid)); + hybrid.setAllowPartialResults(true); + RetrievalResponse response = service.retrieve(hybrid); + assertFalse(response.getEvidence().isEmpty()); + assertEquals(Arrays.asList("VECTOR", "GRAPH"), response.getDegradedChannels()); + assertEquals(RecallStopReason.INDEX_NOT_READY, response.getTrace().getChannelStopReasons().get("VECTOR")); + assertEquals(RecallStageStatus.DEGRADED, response.getTrace().getChannelStatuses().get("VECTOR")); + assertTrue(response.getTrace().getDegradationReasons().get("VECTOR").contains("unable to read")); + } + + @Test + void loaderPreservesBm25FailureAndRejectsBothFailedChannels() throws Exception { + Fixture fixture = fixture(); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + registry.load(GRAPH, VERSION, VERSION, tempDir.resolve("absent"), tempDir.resolve("vector/vector-v1.bin"), + fixture.chunks, fixture.entities, fixture.edges, "fixture", VERSION); + RetrievalRequest hybrid = request("HYBRID", Arrays.asList(1.0, 0.0)); + hybrid.setAllowPartialResults(true); + RetrievalResponse response = new RecallService(registry).retrieve(hybrid); + assertFalse(response.getEvidence().isEmpty()); + assertEquals(Arrays.asList("BM25", "GRAPH"), response.getDegradedChannels()); + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> registry.load(GRAPH, VERSION, VERSION, + tempDir.resolve("absent"), tempDir.resolve("also-absent"), fixture.chunks, fixture.entities, + fixture.edges, "fixture", VERSION)); + } + + @Test + void registryRejectsMislabelledReadersAndChangedChunks() throws Exception { + Fixture fixture = fixture(); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> registry.register(GRAPH, "v2", "v2", + fixture.bm25, fixture.vector, fixture.chunks, fixture.entities, fixture.edges, true)); + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> registry.register("other-graph", VERSION, VERSION, + fixture.bm25, fixture.vector, fixture.chunks, fixture.entities, fixture.edges, true)); + List changed = new ArrayList<>(fixture.chunks); + changed.set(0, new TextChunk("c1", "doc-1", 0, 0, 7, 1, "changed")); + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> registry.register(GRAPH, VERSION, VERSION, + fixture.bm25, fixture.vector, changed, fixture.entities, fixture.edges, true)); + } + + @Test + void oddBudgetRanksAndTraceRoundTripAreConsistent() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.registry, new WeightedRrfConfig(2.0, 0.5, 10)); + RetrievalRequest hybrid = request("HYBRID", Arrays.asList(1.0, 0.0)); + hybrid.setBudget(new RetrievalBudget(2, 1000, 5, 100)); + RetrievalResponse response = service.retrieve(hybrid); + assertEquals(2, response.getTrace().getChannelBudgets().get("BM25")); + assertEquals(2, response.getTrace().getChannelBudgets().get("VECTOR")); + assertEquals(1, response.getTrace().getChannelBudgets().get("GRAPH")); + assertTrue(response.getTrace().getTotalCandidatesEvaluated() <= 5); + for (int index = 0; index < response.getEvidence().size(); index++) { + assertEquals(index + 1, response.getEvidence().get(index).getRank()); + } + assertTrue(response.getTrace().getStages().get(3).getReason().contains("weighted_rrf_k=10")); + RetrievalResponse restored = RetrievalApiJson.parseResponse(RetrievalApiJson.toJson(response)); + assertEquals(response.getTrace().getChannelStopReasons(), restored.getTrace().getChannelStopReasons()); + assertEquals(response.getTrace().getChannelStatuses(), restored.getTrace().getChannelStatuses()); + assertEquals(response.getTrace().getEvaluatedCounts(), restored.getTrace().getEvaluatedCounts()); + } + + @Test + void graphSourcesCannotExceedEvidenceBudgetAndGraphTraceRoundTrips() throws Exception { + Fixture fixture = fixture(); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + List edges = Collections.singletonList(new GraphEdgeRef("edge-1", "studies", "e1", "e2", + Arrays.asList("c1", "c2", "c3"))); + registry.register(GRAPH, VERSION, VERSION, fixture.bm25, fixture.vector, fixture.chunks, fixture.entities, edges, true); + RetrievalRequest graph = request("GRAPH_ONLY", null); + graph.setBudget(new RetrievalBudget(1, 1000, 2, 100)); + RetrievalResponse response = new RecallService(registry).retrieve(graph); + assertEquals(1, response.getEvidence().size()); + assertEquals(2, response.getTrace().getGraphCandidatesProduced()); + assertEquals(1, response.getTrace().getEdgesExamined()); + assertEquals(Collections.singletonList("GRAPH"), response.getTrace().getSelectedChannels()); + assertEquals("CANDIDATE_LIMIT", response.getTrace().getStopReason()); + RetrievalResponse restored = RetrievalApiJson.parseResponse(RetrievalApiJson.toJson(response)); + assertEquals(2, restored.getTrace().getGraphCandidatesProduced()); + assertFalse(restored.getPaths().isEmpty()); + } + + @Test + void invalidGraphProvenanceIsExcludedAndReportedInTrace() throws Exception { + Fixture fixture = fixture(); + RetrievalResponse response = new RecallService(fixture.registry).retrieve(request("GRAPH_ONLY", null)); + + assertEquals(1, response.getEvidence().size()); + assertEquals("c3", response.getEvidence().get(0).getChunks().get(0).getChunkId()); + assertTrue(response.getTrace().getValidationErrors().contains( + "edge edge-invalid references an unknown chunk")); + assertEquals(RecallStopReason.TRACE_VALIDATION_ERROR, + response.getTrace().getChannelStopReasons().get("GRAPH")); + assertEquals("TRACE_VALIDATION_ERROR", response.getTrace().getStopReason()); + assertEquals(1, response.getTrace().getGraphCandidatesProduced()); + } + + @Test + void concurrentRequestsHaveIndependentEvidenceAndTrace() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.registry); + ExecutorService executor = Executors.newFixedThreadPool(4); + try { + List> calls = new ArrayList<>(); + for (int index = 0; index < 24; index++) { + final String mode = index % 2 == 0 ? "HYBRID" : "GRAPH_ONLY"; + calls.add(executor.submit(() -> service.retrieve(request(mode, + mode.equals("HYBRID") ? Arrays.asList(1.0, 0.0) : null)))); + } + for (int index = 0; index < calls.size(); index++) { + RetrievalResponse response = calls.get(index).get(); + List channels = index % 2 == 0 + ? Arrays.asList("BM25", "VECTOR", "GRAPH") : Collections.singletonList("GRAPH"); + assertEquals(channels, response.getTrace().getSelectedChannels()); + assertEquals(Collections.singletonList("GRAPH"), response.getDegradedChannels()); + assertFalse(response.getEvidence().isEmpty()); + } + } finally { + executor.shutdownNow(); + } + } + + @Test + void bm25CountsScoredDocumentsAndStopsInsideCollector() throws Exception { + Fixture fixture = fixture(); + org.apache.geaflow.ai.retrieval.index.ChannelSearchResult limited = + fixture.bm25.searchWithStats("Confucius astronomy ethics stars", 1, 1, 0L); + assertEquals(1, limited.getCandidatesEvaluated()); + assertEquals(1, limited.getHits().size()); + assertEquals(RecallStopReason.CANDIDATE_BUDGET, limited.getStopReason()); + AtomicLong clock = new AtomicLong(); + org.apache.geaflow.ai.retrieval.index.ChannelSearchResult stopped = + fixture.bm25.searchWithStats("Confucius", 10, 10, 4L, clock::getAndIncrement); + assertEquals(1, stopped.getCandidatesEvaluated()); + assertEquals(1, stopped.getHits().size()); + assertEquals(RecallStopReason.DEADLINE, stopped.getStopReason()); + assertEquals(0, fixture.bm25.searchWithStats("not-present", 10, 10, 1L, () -> 1L).getCandidatesEvaluated()); + } + + @Test + void vectorFiniteScanIsExplicitAndDeadlinePreservesEvaluatedHits() throws Exception { + Fixture fixture = fixture(); + org.apache.geaflow.ai.retrieval.index.ChannelSearchResult prefix = + fixture.vector.searchWithStats(new float[] {0.0F, 1.0F}, 1, 1, 0L); + assertEquals(1, prefix.getCandidatesEvaluated()); + assertEquals("c1", prefix.getHits().get(0).getChunkId()); + assertEquals(RecallStopReason.CANDIDATE_BUDGET, prefix.getStopReason()); + assertEquals("c2", fixture.vector.search(new float[] {0.0F, 1.0F}, 3, 1).get(0).getChunkId()); + AtomicLong time = new AtomicLong(); + org.apache.geaflow.ai.retrieval.index.ChannelSearchResult stopped = + fixture.vector.searchWithStats(new float[] {1.0F, 0.0F}, 3, 3, 2L, time::getAndIncrement); + assertEquals(2, stopped.getCandidatesEvaluated()); + assertEquals(2, stopped.getHits().size()); + assertEquals(RecallStopReason.DEADLINE, stopped.getStopReason()); + } + + @Test + void noResultsAndCompletedTraceAreDistinctFromTopK() throws Exception { + Fixture fixture = fixture(); + RecallService service = new RecallService(fixture.registry); + RetrievalRequest query = request("BM25_ONLY", null); + query.setQuery("not-present"); + RetrievalResponse empty = service.retrieve(query); + assertEquals("NO_RESULT", empty.getTrace().getStopReason()); + assertTrue(empty.getEvidence().isEmpty()); + query.setQuery("Confucius"); + assertEquals("COMPLETED", service.retrieve(query).getTrace().getStopReason()); + query.setBudget(new RetrievalBudget(1, 1000, 10, 100)); + assertEquals("TOP_K", service.retrieve(query).getTrace().getStopReason()); + } + + @Test + void legacyArtifactsWithoutIdentityRequireRebuild() throws Exception { + Path oldBm25 = tempDir.resolve("bm25-v1"); + Files.createDirectories(oldBm25); + try (org.apache.lucene.store.Directory directory = org.apache.lucene.store.FSDirectory.open(oldBm25); + org.apache.lucene.analysis.standard.StandardAnalyzer analyzer = new org.apache.lucene.analysis.standard.StandardAnalyzer(); + org.apache.lucene.index.IndexWriter writer = new org.apache.lucene.index.IndexWriter(directory, + new org.apache.lucene.index.IndexWriterConfig(analyzer))) { + writer.commit(); + } + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> new Bm25IndexReader(oldBm25, VERSION, VERSION)); + Path oldVector = tempDir.resolve("vector-v1.bin"); + try (java.io.DataOutputStream output = new java.io.DataOutputStream(Files.newOutputStream(oldVector))) { + output.writeUTF("GEAFLOW-VECTOR-1"); + } + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> new VectorIndexReader(oldVector)); + } + + @Test + void graphDeadlineAndEdgeCapAreObservableAndRepeatedPathsAreMerged() throws Exception { + Fixture fixture = fixture(); + List repeated = Arrays.asList(fixture.edges.get(0), fixture.edges.get(0), fixture.edges.get(1)); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + registry.register(GRAPH, VERSION, VERSION, fixture.bm25, fixture.vector, fixture.chunks, fixture.entities, repeated, true); + RecallService service = new RecallService(registry); + RetrievalResponse response = service.retrieve(request("GRAPH_ONLY", null)); + assertEquals(1, response.getEvidence().size()); + assertEquals(1, response.getEvidence().get(0).getPaths().size()); + RetrievalRequest capped = request("GRAPH_ONLY", null); + capped.setBudget(new RetrievalBudget(1, 1000, 1, 100)); + assertTrue(service.retrieve(capped).getTrace().getEdgesExamined() <= 1); + AtomicLong time = new AtomicLong(); + RecallService timed = new RecallService(registry, WeightedRrfConfig.DEFAULT, () -> time.getAndAdd(1_000_000L)); + capped.setBudget(new RetrievalBudget(1, 2, 2, 100)); + assertCode(RetrievalErrorCode.RETRIEVAL_TIMEOUT, () -> timed.retrieve(capped)); + time.set(0L); + capped.setAllowPartialResults(true); + assertEquals("DEADLINE", timed.retrieve(capped).getTrace().getStopReason()); + } + + @Test + void zeroNormAndImpossibleShapeArtifactsAreTypedCorruption() throws Exception { + Fixture fixture = fixture(); + Path artifact = tempDir.resolve("invalid-vector.bin"); + for (int count : new int[] {1, Integer.MAX_VALUE}) { + try (java.io.DataOutputStream output = new java.io.DataOutputStream(Files.newOutputStream(artifact))) { + output.writeUTF("GEAFLOW-VECTOR-2"); + output.writeUTF(GRAPH); + output.writeUTF(VERSION); + output.writeUTF(VERSION); + output.writeUTF(org.apache.geaflow.ai.retrieval.index.ArtifactIdentity.fingerprint(fixture.chunks)); + output.writeUTF("fixture"); + output.writeUTF(VERSION); + output.writeInt(2); + output.writeInt(count); + output.writeUTF("c1"); + output.writeUTF("doc-1"); + output.writeFloat(0.0F); + output.writeFloat(0.0F); + } + assertCode(RetrievalErrorCode.INDEX_NOT_READY, () -> new VectorIndexReader(artifact)); + } + } + + @Test + void missingVectorAndCorruptBm25LoaderPathsRetainSuccessfulChannel() throws Exception { + Fixture fixture = fixture(); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + registry.load(GRAPH, VERSION, VERSION, tempDir.resolve("bm25/bm25-v1"), tempDir.resolve("absent"), + fixture.chunks, fixture.entities, fixture.edges, "fixture", VERSION); + readers.add(registry.get(GRAPH, VERSION, VERSION).getBm25()); + RetrievalRequest hybrid = request("HYBRID", Arrays.asList(1.0, 0.0)); + hybrid.setAllowPartialResults(true); + RetrievalResponse partial = new RecallService(registry).retrieve(hybrid); + assertEquals(Arrays.asList("VECTOR", "GRAPH"), partial.getDegradedChannels()); + assertFalse(partial.getEvidence().isEmpty()); + Path corruptBm25 = tempDir.resolve("corrupt-bm25"); + Files.createDirectories(corruptBm25); + Files.write(corruptBm25.resolve("segments_1"), new byte[] {0, 1, 2}); + registry.load(GRAPH, VERSION, VERSION, corruptBm25, tempDir.resolve("vector/vector-v1.bin"), + fixture.chunks, fixture.entities, fixture.edges, "fixture", VERSION); + partial = new RecallService(registry).retrieve(hybrid); + assertEquals(Arrays.asList("BM25", "GRAPH"), partial.getDegradedChannels()); + assertFalse(partial.getEvidence().isEmpty()); + assertEquals(RecallStopReason.INDEX_NOT_READY, partial.getTrace().getChannelStopReasons().get("BM25")); + } + + @Test + void directFacadeRequestsAlsoHaveIndependentExecutionState() throws Exception { + Fixture fixture = fixture(); + RecallService direct = new RecallService(fixture.bm25, fixture.vector, fixture.chunks, fixture.entities, fixture.edges); + ExecutorService executor = Executors.newFixedThreadPool(4); + try { + List> calls = new ArrayList<>(); + for (int index = 0; index < 24; index++) { + final String mode = index % 2 == 0 ? "HYBRID" : "GRAPH_ONLY"; + calls.add(executor.submit(() -> direct.retrieveWithTrace("Confucius", + mode.equals("HYBRID") ? new float[] {1.0F, 0.0F} : null, mode, 2, 4, 0L))); + } + for (int index = 0; index < calls.size(); index++) { + RetrievalResponse response = calls.get(index).get(); + assertEquals(index % 2 == 0 ? Arrays.asList("BM25", "VECTOR", "GRAPH") : Collections.singletonList("GRAPH"), + response.getTrace().getSelectedChannels()); + assertEquals(index % 2 == 0 ? 4 : 2, response.getTrace().getTotalCandidatesEvaluated()); + assertEquals(index % 2 == 0 ? Collections.emptyList() : Collections.singletonList("GRAPH"), + response.getDegradedChannels()); + } + } finally { + executor.shutdownNow(); + } + } + + private Fixture fixture() throws Exception { + List chunks = Arrays.asList( + new TextChunk("c1", "doc-1", 0, 0, 36, 6, "Confucius taught ethics and philosophy"), + new TextChunk("c2", "doc-2", 0, 0, 25, 4, "Astronomy studies distant stars"), + new TextChunk("c3", "doc-3", 0, 0, 35, 6, "Confucius and astronomy are studied")); + IngestionContext context = context(); + Map vectors = new HashMap<>(); + vectors.put("c1", new float[] {1.0F, 0.0F}); + vectors.put("c2", new float[] {0.0F, 1.0F}); + vectors.put("c3", new float[] {0.8F, 0.2F}); + Path bm25Path = tempDir.resolve("bm25"); + Path vectorPath = tempDir.resolve("vector"); + try (org.apache.geaflow.ai.retrieval.index.IndexArtifact ignored = + new LuceneBm25IndexBuilder(bm25Path).build(context, chunks); + org.apache.geaflow.ai.retrieval.index.IndexArtifact ignoredVector = + new OfflineVectorIndexBuilder(vectorPath, vectors, "fixture", VERSION) + .build(context, chunks)) { + // Artifacts are published before readers are opened. + } + Bm25IndexReader bm25 = new Bm25IndexReader(bm25Path.resolve("bm25-v1"), VERSION, VERSION); + readers.add(bm25); + VectorIndexReader vector = new VectorIndexReader(vectorPath.resolve("vector-v1.bin"), + VERSION, VERSION, "fixture", VERSION); + List entities = Arrays.asList( + new EntityRef("e1", "Confucius", Collections.singletonList("kongzi"), "person", + Collections.singletonList("c1")), + new EntityRef("e2", "Astronomy", "topic")); + List edges = Arrays.asList( + new GraphEdgeRef("edge-1", "studies", "e1", "e2", Collections.singletonList("c3")), + new GraphEdgeRef("edge-invalid", "bad", "e1", "e2", Collections.singletonList("missing"))); + RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); + registry.register(GRAPH, VERSION, VERSION, bm25, vector, chunks, entities, edges, true); + return new Fixture(registry, bm25, vector, chunks, entities, edges); + } + + private IngestionContext context() { + DatasetManifest manifest = new DatasetManifest(VERSION, "dataset", "revision", "dev", null, + "cache", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "policy", new ChunkingConfiguration("chunk-v1", 100, 10), "graph", "fixture", VERSION, 1L); + return new IngestionContext(manifest, new GraphVersion(GRAPH, VERSION), "test"); + } + + private RetrievalRequest request(String mode, List vector) { + RetrievalRequest request = new RetrievalRequest(); + request.setGraphName(GRAPH); + request.setGraphVersion(VERSION); + request.setIndexVersion(VERSION); + request.setQuery("Confucius"); + request.setMode(mode); + request.setQueryVector(vector); + request.setBudget(new RetrievalBudget(10, 1000, 10, 100)); + return request; + } + + private static Evidence find(List evidence, String id) { + return evidence.stream().filter(value -> id.equals(value.getEvidenceId())).findFirst() + .orElseThrow(() -> new AssertionError("missing evidence " + id)); + } + + private static void assertCode(RetrievalErrorCode code, org.junit.jupiter.api.function.Executable executable) { + RetrievalException exception = assertThrows(RetrievalException.class, executable); + assertEquals(code, exception.getCode()); + } + + private static final class Fixture { + private final RetrievalFixtureRegistry registry; + private final Bm25IndexReader bm25; + private final VectorIndexReader vector; + private final List chunks; + private final List entities; + private final List edges; + + private Fixture(RetrievalFixtureRegistry registry, Bm25IndexReader bm25, + VectorIndexReader vector, List chunks, + List entities, List edges) { + this.registry = registry; + this.bm25 = bm25; + this.vector = vector; + this.chunks = chunks; + this.entities = entities; + this.edges = edges; + } + } +} From 0e3b79da8a42a9279b9a1b303b560a7bca3e0b56 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Thu, 8 Oct 2026 21:03:34 +0800 Subject: [PATCH 23/30] feat(ai): extend retrieval contract for hybrid recall --- .../retrieval/api/model/RetrievalCommand.java | 18 ++ .../ai/retrieval/api/model/RetrievalMode.java | 12 +- .../retrieval/api/model/RetrievalRequest.java | 54 ++++ .../retrieval/api/model/RetrievalTrace.java | 238 ++++++++++++++++++ .../ai/retrieval/api/package-info.java | 14 +- .../ai/retrieval/codec/RetrievalApiJson.java | 112 +++++++++ .../retrieval/codec/RetrievalModelJson.java | 9 +- .../service/RetrievalRequestValidator.java | 62 ++++- .../service/RetrievalResponseValidator.java | 116 ++++++++- .../retrieval/service/RetrievalService.java | 7 + .../apache/geaflow/ai/MemoryServerTest.java | 8 + .../retrieval/api/RetrievalApiJsonTest.java | 91 ++++++- .../api/RetrievalRequestValidatorTest.java | 28 ++- .../service/RetrievalServiceTest.java | 15 ++ .../retrieval/api/response-empty.json | 1 + 15 files changed, 770 insertions(+), 15 deletions(-) diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalCommand.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalCommand.java index aa646f813..09dbf0dc0 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalCommand.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalCommand.java @@ -18,6 +18,10 @@ package org.apache.geaflow.ai.retrieval.api.model; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + /** Immutable request passed beyond the validation boundary. */ public final class RetrievalCommand { @@ -26,15 +30,24 @@ public final class RetrievalCommand { private final RetrievalMode mode; private final ExecutionMode executionMode; private final RetrievalBudget budget; + private final List queryVector; public RetrievalCommand(String graphName, String query, RetrievalMode mode, ExecutionMode executionMode, RetrievalBudget budget) { + this(graphName, query, mode, executionMode, budget, Collections.emptyList()); + } + + public RetrievalCommand(String graphName, String query, RetrievalMode mode, + ExecutionMode executionMode, RetrievalBudget budget, + List queryVector) { this.graphName = graphName; this.query = query; this.mode = mode; this.executionMode = executionMode; this.budget = new RetrievalBudget(budget.getTopK(), budget.getTimeoutMs(), budget.getMaxCandidates(), budget.getTokenBudget()); + this.queryVector = queryVector == null ? Collections.emptyList() + : Collections.unmodifiableList(new ArrayList<>(queryVector)); } public String getGraphName() { @@ -57,4 +70,9 @@ public RetrievalBudget getBudget() { return new RetrievalBudget(budget.getTopK(), budget.getTimeoutMs(), budget.getMaxCandidates(), budget.getTokenBudget()); } + + /** Returns the caller-supplied, precomputed query vector. */ + public List getQueryVector() { + return queryVector; + } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalMode.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalMode.java index 13688c16a..6175dda22 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalMode.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalMode.java @@ -20,5 +20,15 @@ /** Retrieval modes understood by the versioned API. */ public enum RetrievalMode { - KEYWORD + /** Legacy keyword mode, equivalent to BM25_ONLY. */ + KEYWORD, + BM25_ONLY, + VECTOR_ONLY, + GRAPH_ONLY, + HYBRID; + + /** Returns the canonical MVP mode while preserving KEYWORD compatibility. */ + public RetrievalMode canonical() { + return this == KEYWORD ? BM25_ONLY : this; + } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalRequest.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalRequest.java index 6c4194701..8b86b0652 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalRequest.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalRequest.java @@ -24,11 +24,17 @@ public class RetrievalRequest { private String graphName; + private String requestId; private String query; private List queryVector; private String mode; private String executionMode; private RetrievalBudget budget; + private String graphVersion; + private String indexVersion; + private String vectorVersion; + private String vectorSource; + private boolean allowPartialResults; public RetrievalRequest() { } @@ -37,6 +43,22 @@ public String getGraphName() { return graphName; } + public String getRequestId() { + return requestId; + } + + public void setRequestId(String requestId) { + this.requestId = requestId; + } + + public boolean isAllowPartialResults() { + return allowPartialResults; + } + + public void setAllowPartialResults(boolean allowPartialResults) { + this.allowPartialResults = allowPartialResults; + } + public void setGraphName(String graphName) { this.graphName = graphName; } @@ -80,4 +102,36 @@ public RetrievalBudget getBudget() { public void setBudget(RetrievalBudget budget) { this.budget = budget; } + + public String getGraphVersion() { + return graphVersion; + } + + public void setGraphVersion(String graphVersion) { + this.graphVersion = graphVersion; + } + + public String getIndexVersion() { + return indexVersion; + } + + public void setIndexVersion(String indexVersion) { + this.indexVersion = indexVersion; + } + + public String getVectorVersion() { + return vectorVersion; + } + + public void setVectorVersion(String vectorVersion) { + this.vectorVersion = vectorVersion; + } + + public String getVectorSource() { + return vectorSource; + } + + public void setVectorSource(String vectorSource) { + this.vectorSource = vectorSource; + } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalTrace.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalTrace.java index b5babc669..ea66cf1df 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalTrace.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalTrace.java @@ -19,17 +19,47 @@ package org.apache.geaflow.ai.retrieval.api.model; import java.util.ArrayList; +import java.util.LinkedHashMap; import java.util.List; +import java.util.Map; +import org.apache.geaflow.ai.retrieval.execution.RecallStageStatus; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; /** Stable minimum trace shape; future fields must be additive. */ public class RetrievalTrace { private String traceVersion; + private String requestId; private String originalQuery; private RetrievalMode selectedMode; private ExecutionMode executionMode; private List stages = new ArrayList<>(); private String stopReason; + private String graphVersion; + private String indexVersion; + private String vectorVersion; + private int effectiveCandidateBudget; + private long elapsedNanos; + private Map candidateCounts = new LinkedHashMap<>(); + private int anchorsConsidered; + private int edgesExamined; + private int neighborsSampled; + private int verticesReached; + private int graphCandidatesProduced; + private List degradedChannels = new ArrayList<>(); + private List validationErrors = new ArrayList<>(); + private List selectedChannels = new ArrayList<>(); + private int totalCandidatesEvaluated; + private Map channelBudgets = new LinkedHashMap<>(); + private Map degradationReasons = new LinkedHashMap<>(); + private double bm25Weight = 1.0; + private double vectorWeight = 1.0; + private double graphWeight = 1.0; + private int rrfRankConstant = 60; + private int effectiveTopK; + private Map evaluatedCounts = new LinkedHashMap<>(); + private Map channelStatuses = new LinkedHashMap<>(); + private Map channelStopReasons = new LinkedHashMap<>(); public RetrievalTrace() { } @@ -42,6 +72,14 @@ public void setTraceVersion(String traceVersion) { this.traceVersion = traceVersion; } + public String getRequestId() { + return requestId; + } + + public void setRequestId(String requestId) { + this.requestId = requestId; + } + public String getOriginalQuery() { return originalQuery; } @@ -81,4 +119,204 @@ public String getStopReason() { public void setStopReason(String stopReason) { this.stopReason = stopReason; } + + public String getGraphVersion() { + return graphVersion; + } + + public void setGraphVersion(String graphVersion) { + this.graphVersion = graphVersion; + } + + public String getIndexVersion() { + return indexVersion; + } + + public void setIndexVersion(String indexVersion) { + this.indexVersion = indexVersion; + } + + public String getVectorVersion() { + return vectorVersion; + } + + public void setVectorVersion(String vectorVersion) { + this.vectorVersion = vectorVersion; + } + + public int getEffectiveCandidateBudget() { + return effectiveCandidateBudget; + } + + public void setEffectiveCandidateBudget(int value) { + this.effectiveCandidateBudget = value; + } + + public long getElapsedNanos() { + return elapsedNanos; + } + + public void setElapsedNanos(long value) { + this.elapsedNanos = value; + } + + public Map getCandidateCounts() { + return candidateCounts; + } + + public void setCandidateCounts(Map value) { + candidateCounts = value == null ? new LinkedHashMap<>() : new LinkedHashMap<>(value); + } + + public int getAnchorsConsidered() { + return anchorsConsidered; + } + + public void setAnchorsConsidered(int value) { + anchorsConsidered = value; + } + + public int getEdgesExamined() { + return edgesExamined; + } + + public void setEdgesExamined(int value) { + edgesExamined = value; + } + + public int getNeighborsSampled() { + return neighborsSampled; + } + + public void setNeighborsSampled(int value) { + neighborsSampled = value; + } + + public int getVerticesReached() { + return verticesReached; + } + + public void setVerticesReached(int value) { + verticesReached = value; + } + + public int getGraphCandidatesProduced() { + return graphCandidatesProduced; + } + + public void setGraphCandidatesProduced(int value) { + graphCandidatesProduced = value; + } + + public List getDegradedChannels() { + return degradedChannels; + } + + public void setDegradedChannels(List value) { + degradedChannels = value == null ? new ArrayList() : new ArrayList<>(value); + } + + public List getValidationErrors() { + return validationErrors; + } + + public void setValidationErrors(List value) { + validationErrors = value == null ? new ArrayList() : new ArrayList<>(value); + } + + public List getSelectedChannels() { + return selectedChannels; + } + + public void setSelectedChannels(List value) { + selectedChannels = value == null ? new ArrayList() : new ArrayList<>(value); + } + + public int getTotalCandidatesEvaluated() { + return totalCandidatesEvaluated; + } + + public void setTotalCandidatesEvaluated(int value) { + totalCandidatesEvaluated = value; + } + + public Map getChannelBudgets() { + return channelBudgets; + } + + public void setChannelBudgets(Map value) { + channelBudgets = value == null ? new LinkedHashMap<>() : new LinkedHashMap<>(value); + } + + public Map getDegradationReasons() { + return degradationReasons; + } + + public void setDegradationReasons(Map value) { + degradationReasons = value == null ? new LinkedHashMap<>() : new LinkedHashMap<>(value); + } + + public double getBm25Weight() { + return bm25Weight; + } + + public void setBm25Weight(double value) { + bm25Weight = value; + } + + public double getVectorWeight() { + return vectorWeight; + } + + public void setVectorWeight(double value) { + vectorWeight = value; + } + + public double getGraphWeight() { + return graphWeight; + } + + public void setGraphWeight(double value) { + graphWeight = value; + } + + public int getRrfRankConstant() { + return rrfRankConstant; + } + + public void setRrfRankConstant(int value) { + rrfRankConstant = value; + } + + public int getEffectiveTopK() { + return effectiveTopK; + } + + public void setEffectiveTopK(int value) { + effectiveTopK = value; + } + + public Map getEvaluatedCounts() { + return evaluatedCounts; + } + + public void setEvaluatedCounts(Map value) { + evaluatedCounts = value == null ? new LinkedHashMap<>() : new LinkedHashMap<>(value); + } + + public Map getChannelStatuses() { + return channelStatuses; + } + + public void setChannelStatuses(Map value) { + channelStatuses = value == null ? new LinkedHashMap<>() : new LinkedHashMap<>(value); + } + + public Map getChannelStopReasons() { + return channelStopReasons; + } + + public void setChannelStopReasons(Map value) { + channelStopReasons = value == null ? new LinkedHashMap<>() : new LinkedHashMap<>(value); + } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/package-info.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/package-info.java index 09cdd2847..d6f781d5d 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/package-info.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/package-info.java @@ -20,10 +20,16 @@ /** * Versioned retrieval API contract for Java and HTTP clients. * - *

Version {@code v1} supports {@code KEYWORD} retrieval with {@code SEQUENTIAL} execution. - * {@code PARALLEL}, {@code CASCADED}, and non-empty query vectors are reserved and must be - * reported as {@code UNSUPPORTED_OPTION}. Unknown JSON fields are accepted for additive wire - * compatibility; duplicate fields are rejected. Response collections are always JSON arrays.

+ *

The core version {@code v1} facade supports {@code BM25_ONLY}, {@code VECTOR_ONLY}, + * {@code GRAPH_ONLY}, and fixed {@code HYBRID} retrieval with {@code SEQUENTIAL} execution. + * {@code KEYWORD} is retained as an alias for {@code BM25_ONLY}. A non-empty query vector is + * required by {@code VECTOR_ONLY} and {@code HYBRID}; it is unsupported for the other modes. + * {@code PARALLEL} and {@code CASCADED} execution are reserved and reported as + * {@code UNSUPPORTED_OPTION}. The HTTP endpoint currently exposes {@code KEYWORD} only. + * Unknown JSON fields are accepted for additive wire compatibility; duplicate fields are + * rejected. Response collections are always JSON arrays. Hybrid responses include selected + * channels, channel budgets and statuses, evaluated counts, versions, and stop/degradation + * reasons in the trace.

* *

Error codes map to HTTP status as follows: {@code INVALID_REQUEST}=400, * {@code UNSUPPORTED_OPTION}=400, {@code GRAPH_NOT_FOUND}=404, diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java index 1a51985a4..03b710760 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalApiJson.java @@ -160,6 +160,11 @@ private static JsonObject object(String json) { private static void validateRequestShape(JsonObject object) { requireString(object, "graphName"); requireString(object, "query"); + optionalString(object, "requestId"); + optionalString(object, "graphVersion"); + optionalString(object, "indexVersion"); + optionalString(object, "vectorVersion"); + optionalString(object, "vectorSource"); optionalString(object, "mode"); optionalString(object, "executionMode"); if (object.has("queryVector") && !object.get("queryVector").isJsonNull()) { @@ -201,6 +206,7 @@ private static RetrievalTrace parseTrace(JsonObject object) { } RetrievalTrace trace = new RetrievalTrace(); trace.setTraceVersion(string(object, "traceVersion")); + trace.setRequestId(optionalString(object, "requestId")); trace.setOriginalQuery(string(object, "originalQuery")); trace.setSelectedMode(enumValue(object, "selectedMode", org.apache.geaflow.ai.retrieval.api.model.RetrievalMode.class)); @@ -208,9 +214,115 @@ private static RetrievalTrace parseTrace(JsonObject object) { requireArray(object, "stages"); trace.setStages(traceStages(object, "stages")); trace.setStopReason(optionalString(object, "stopReason")); + trace.setGraphVersion(optionalString(object, "graphVersion")); + trace.setIndexVersion(optionalString(object, "indexVersion")); + trace.setVectorVersion(optionalString(object, "vectorVersion")); + trace.setEffectiveCandidateBudget(optionalInt(object, "effectiveCandidateBudget") == null + ? 0 : optionalInt(object, "effectiveCandidateBudget")); + trace.setElapsedNanos(optionalLong(object, "elapsedNanos")); + trace.setCandidateCounts(integerMap(object, "candidateCounts")); + trace.setEvaluatedCounts(integerMap(object, "evaluatedCounts")); + trace.setEffectiveTopK(optionalInt(object, "effectiveTopK") == null ? 0 : optionalInt(object, "effectiveTopK")); + trace.setChannelStatuses(enumMap(object, "channelStatuses", + org.apache.geaflow.ai.retrieval.execution.RecallStageStatus.class)); + trace.setChannelStopReasons(enumMap(object, "channelStopReasons", + org.apache.geaflow.ai.retrieval.execution.RecallStopReason.class)); + trace.setAnchorsConsidered(optionalInt(object, "anchorsConsidered") == null ? 0 : optionalInt(object, "anchorsConsidered")); + trace.setEdgesExamined(optionalInt(object, "edgesExamined") == null ? 0 : optionalInt(object, "edgesExamined")); + trace.setNeighborsSampled(optionalInt(object, "neighborsSampled") == null ? 0 : optionalInt(object, "neighborsSampled")); + trace.setVerticesReached(optionalInt(object, "verticesReached") == null ? 0 : optionalInt(object, "verticesReached")); + trace.setGraphCandidatesProduced(optionalInt(object, "graphCandidatesProduced") == null ? 0 : optionalInt(object, "graphCandidatesProduced")); + trace.setTotalCandidatesEvaluated(optionalInt(object, "totalCandidatesEvaluated") == null + ? 0 : optionalInt(object, "totalCandidatesEvaluated")); + trace.setSelectedChannels(strings(object, "selectedChannels")); + trace.setChannelBudgets(integerMap(object, "channelBudgets")); + trace.setDegradedChannels(strings(object, "degradedChannels")); + trace.setDegradationReasons(stringMap(object, "degradationReasons")); + trace.setBm25Weight(optionalDouble(object, "bm25Weight", 1.0)); + trace.setVectorWeight(optionalDouble(object, "vectorWeight", 1.0)); + trace.setGraphWeight(optionalDouble(object, "graphWeight", 1.0)); + trace.setRrfRankConstant(optionalInt(object, "rrfRankConstant") == null + ? 60 : optionalInt(object, "rrfRankConstant")); + trace.setValidationErrors(strings(object, "validationErrors")); return trace; } + private static > java.util.Map enumMap(JsonObject object, String name, Class type) { + java.util.Map result = new java.util.LinkedHashMap<>(); + for (java.util.Map.Entry entry : stringMap(object, name).entrySet()) { + try { + result.put(entry.getKey(), Enum.valueOf(type, entry.getValue())); + } catch (IllegalArgumentException error) { + throw new JsonParseException("unknown " + name + ": " + entry.getValue(), error); + } + } + return result; + } + + private static long optionalLong(JsonObject object, String name) { + JsonElement value = object.get(name); + if (value == null || value.isJsonNull()) { + return 0L; + } + if (!value.isJsonPrimitive() || !value.getAsJsonPrimitive().isNumber() + || !Double.isFinite(value.getAsDouble()) || value.getAsDouble() != Math.rint(value.getAsDouble())) { + throw new JsonParseException(name + " must be an integer"); + } + return value.getAsLong(); + } + + private static double optionalDouble(JsonObject object, String name, double fallback) { + JsonElement value = object.get(name); + if (value == null || value.isJsonNull()) { + return fallback; + } + if (!value.isJsonPrimitive() || !value.getAsJsonPrimitive().isNumber() + || !Double.isFinite(value.getAsDouble())) { + throw new JsonParseException(name + " must be a finite number"); + } + return value.getAsDouble(); + } + + private static java.util.Map integerMap(JsonObject object, String name) { + java.util.Map result = new java.util.LinkedHashMap<>(); + JsonElement value = object.get(name); + if (value == null || value.isJsonNull()) { + return result; + } + if (!value.isJsonObject()) { + throw new JsonParseException(name + " must be an object"); + } + for (java.util.Map.Entry entry : value.getAsJsonObject().entrySet()) { + JsonElement item = entry.getValue(); + if (!item.isJsonPrimitive() || !item.getAsJsonPrimitive().isNumber() + || !Double.isFinite(item.getAsDouble()) || item.getAsDouble() != Math.rint(item.getAsDouble()) + || item.getAsDouble() < Integer.MIN_VALUE || item.getAsDouble() > Integer.MAX_VALUE) { + throw new JsonParseException(name + " values must be integers"); + } + result.put(entry.getKey(), item.getAsInt()); + } + return result; + } + + private static java.util.Map stringMap(JsonObject object, String name) { + java.util.Map result = new java.util.LinkedHashMap<>(); + JsonElement value = object.get(name); + if (value == null || value.isJsonNull()) { + return result; + } + if (!value.isJsonObject()) { + throw new JsonParseException(name + " must be an object"); + } + for (java.util.Map.Entry entry : value.getAsJsonObject().entrySet()) { + JsonElement item = entry.getValue(); + if (!item.isJsonPrimitive() || !item.getAsJsonPrimitive().isString()) { + throw new JsonParseException(name + " values must be strings"); + } + result.put(entry.getKey(), item.getAsString()); + } + return result; + } + private static List traceStages(JsonObject object, String name) { JsonElement element = object.get(name); if (element == null || element.isJsonNull()) { diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java index 2c73715ba..69e0ccc90 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java @@ -36,6 +36,7 @@ import org.apache.geaflow.ai.retrieval.model.evidence.ChannelScore; import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; import org.apache.geaflow.ai.retrieval.model.evidence.EvidenceKind; +import org.apache.geaflow.ai.retrieval.model.evidence.GraphEvidenceTrace; import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; import org.apache.geaflow.ai.retrieval.model.graph.GraphEdgeRef; import org.apache.geaflow.ai.retrieval.model.graph.GraphPathRef; @@ -131,7 +132,13 @@ private static Object parse(JsonObject object, Class type) { models(object, "entities", EntityRef.class), models(object, "paths", GraphPathRef.class), models(object, "sources", SourceRef.class), scores(object), optionalDouble(object, "fusedScore"), optionalDouble(object, "finalScore"), - optionalInt(object, "rank")); + optionalInt(object, "rank"), models(object, "graphTraces", GraphEvidenceTrace.class)); + } else if (type == GraphEvidenceTrace.class) { + JsonObject path = object.getAsJsonObject("path"); + return new GraphEvidenceTrace(requiredString(object, "queryId"), + requiredString(object, "graphVersion"), requiredString(object, "anchorEntityId"), + requiredString(object, "anchorMatchType"), requiredDouble(object, "anchorConfidence"), + (GraphPathRef) parse(path, GraphPathRef.class), strings(object, "supportingChunkIds")); } throw new JsonParseException("unsupported retrieval model: " + type.getName()); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java index a139fe736..0202dd689 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalRequestValidator.java @@ -19,6 +19,7 @@ package org.apache.geaflow.ai.retrieval.service; import java.util.List; +import java.util.Locale; import java.util.Objects; import org.apache.geaflow.ai.retrieval.api.model.ExecutionMode; import org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget; @@ -34,6 +35,7 @@ public class RetrievalRequestValidator { private static final int MAX_GRAPH_NAME_LENGTH = 128; private static final int MAX_QUERY_LENGTH = 4096; + private static final int MAX_IDENTIFIER_LENGTH = 256; private final RetrievalProperties properties; public RetrievalRequestValidator(RetrievalProperties properties) { @@ -47,20 +49,38 @@ public RetrievalCommand validate(RetrievalRequest request) { } final String graphName = trimmed(request.getGraphName(), "graphName", MAX_GRAPH_NAME_LENGTH); final String query = trimmed(request.getQuery(), "query", MAX_QUERY_LENGTH); + optionalText(request.getRequestId(), "requestId", MAX_IDENTIFIER_LENGTH); + optionalText(request.getGraphVersion(), "graphVersion", MAX_IDENTIFIER_LENGTH); + optionalText(request.getIndexVersion(), "indexVersion", MAX_IDENTIFIER_LENGTH); + optionalText(request.getVectorVersion(), "vectorVersion", MAX_IDENTIFIER_LENGTH); + optionalText(request.getVectorSource(), "vectorSource", MAX_IDENTIFIER_LENGTH); final RetrievalMode mode = parseMode(request.getMode() == null ? properties.getDefaultMode() : request.getMode()); final ExecutionMode executionMode = parseExecutionMode(request.getExecutionMode() == null ? properties.getDefaultExecutionMode() : request.getExecutionMode()); List vector = request.getQueryVector(); - if (vector != null && !vector.isEmpty()) { - throw unsupported("non-empty queryVector is not supported in v1"); + if (vector != null) { + validateVector(vector); + } + RetrievalMode canonicalMode = mode.canonical(); + boolean requiresVector = canonicalMode == RetrievalMode.VECTOR_ONLY + || canonicalMode == RetrievalMode.HYBRID; + if (requiresVector && (vector == null || vector.isEmpty())) { + throw invalid("queryVector is required for mode " + mode); + } + if (vector != null && !vector.isEmpty() + && canonicalMode != RetrievalMode.VECTOR_ONLY && canonicalMode != RetrievalMode.HYBRID) { + throw unsupported("queryVector is not supported for mode " + mode); } RetrievalBudget input = request.getBudget(); final RetrievalBudget effectiveBudget = RetrievalBudgetValidator.validate(input, properties, false); + if (canonicalMode == RetrievalMode.HYBRID && effectiveBudget.getMaxCandidates() < 3) { + throw invalid("HYBRID requires maxCandidates >= 3"); + } return new RetrievalCommand(graphName, query, mode, executionMode, - effectiveBudget); + effectiveBudget, vector); } private static String trimmed(String value, String name, int maxLength) { @@ -74,11 +94,41 @@ private static String trimmed(String value, String name, int maxLength) { return result; } + private static void optionalText(String value, String name, int maxLength) { + if (value == null) { + return; + } + if (value.trim().isEmpty()) { + throw invalid(name + " must not be blank"); + } + if (value.trim().length() > maxLength) { + throw invalid(name + " exceeds maximum length " + maxLength); + } + } + private static RetrievalMode parseMode(String value) { - if (value == null || !RetrievalMode.KEYWORD.name().equals(value)) { + if (value == null) { + throw unsupported("unsupported retrieval mode: null"); + } + try { + return RetrievalMode.valueOf(value.trim().toUpperCase(java.util.Locale.ROOT)); + } catch (IllegalArgumentException e) { throw unsupported("unsupported retrieval mode: " + value); } - return RetrievalMode.KEYWORD; + } + + private static void validateVector(List vector) { + double norm = 0.0; + for (Double value : vector) { + if (value == null || !Double.isFinite(value) || !Float.isFinite(value.floatValue())) { + throw invalid("queryVector must contain finite float values"); + } + double converted = value.floatValue(); + norm += converted * converted; + } + if (!vector.isEmpty() && (!Double.isFinite(norm) || norm == 0.0)) { + throw invalid("queryVector must have a finite, non-zero norm"); + } } private static ExecutionMode parseExecutionMode(String value) { @@ -86,7 +136,7 @@ private static ExecutionMode parseExecutionMode(String value) { throw unsupported("unsupported execution mode: null"); } try { - ExecutionMode mode = ExecutionMode.valueOf(value); + ExecutionMode mode = ExecutionMode.valueOf(value.trim().toUpperCase(Locale.ROOT)); if (mode != ExecutionMode.SEQUENTIAL) { throw unsupported("unsupported execution mode: " + value); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalResponseValidator.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalResponseValidator.java index 02a846496..3aea8dcb9 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalResponseValidator.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalResponseValidator.java @@ -20,7 +20,10 @@ package org.apache.geaflow.ai.retrieval.service; import java.util.Collections; +import java.util.HashSet; import java.util.List; +import java.util.Map; +import java.util.Set; import org.apache.geaflow.ai.retrieval.api.model.ExecutionMode; import org.apache.geaflow.ai.retrieval.api.model.RetrievalBudget; import org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode; @@ -30,6 +33,8 @@ import org.apache.geaflow.ai.retrieval.api.model.RetrievalTrace; import org.apache.geaflow.ai.retrieval.api.model.TraceStage; import org.apache.geaflow.ai.retrieval.config.RetrievalProperties; +import org.apache.geaflow.ai.retrieval.execution.RecallStageStatus; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; /** Validates the complete success response before it crosses the API boundary. */ public final class RetrievalResponseValidator { @@ -50,8 +55,12 @@ public static RetrievalResponse validate(RetrievalResponse response, throw invalid("trace is required"); } required(trace.getTraceVersion(), "traceVersion"); + required(trace.getRequestId(), "trace.requestId"); + if (!response.getRequestId().equals(trace.getRequestId())) { + throw invalid("trace.requestId must match requestId"); + } required(trace.getOriginalQuery(), "originalQuery"); - if (trace.getSelectedMode() != RetrievalMode.KEYWORD) { + if (trace.getSelectedMode() == null) { throw unsupported("unsupported retrieval mode in trace"); } if (trace.getExecutionMode() != ExecutionMode.SEQUENTIAL) { @@ -72,11 +81,116 @@ public static RetrievalResponse validate(RetrievalResponse response, || response.getSources() == null || response.getDegradedChannels() == null) { throw invalid("response collections are required"); } + if (trace.getGraphVersion() != null + && !response.getGraphVersion().equals(trace.getGraphVersion())) { + throw invalid("trace.graphVersion must match graphVersion"); + } RetrievalBudget effectiveBudget = response.getEffectiveBudget(); RetrievalBudgetValidator.validate(effectiveBudget, properties, true); + if (trace.getSelectedMode().canonical() == RetrievalMode.HYBRID) { + validateHybridTrace(trace, effectiveBudget); + } return response; } + private static void validateHybridTrace(RetrievalTrace trace, RetrievalBudget effectiveBudget) { + required(trace.getGraphVersion(), "trace.graphVersion"); + required(trace.getIndexVersion(), "trace.indexVersion"); + required(trace.getStopReason(), "trace.stopReason"); + try { + RecallStopReason.valueOf(trace.getStopReason()); + } catch (IllegalArgumentException error) { + throw invalid("unsupported trace.stopReason: " + trace.getStopReason()); + } + if (!Double.isFinite(trace.getBm25Weight()) || trace.getBm25Weight() < 0.0 + || !Double.isFinite(trace.getVectorWeight()) || trace.getVectorWeight() < 0.0 + || !Double.isFinite(trace.getGraphWeight()) || trace.getGraphWeight() < 0.0 + || trace.getBm25Weight() == 0.0 && trace.getVectorWeight() == 0.0 + && trace.getGraphWeight() == 0.0) { + throw invalid("trace fusion weights must be finite, non-negative, and non-zero"); + } + if (trace.getRrfRankConstant() < 1) { + throw invalid("trace.rrfRankConstant must be positive"); + } + Set expected = new HashSet<>(); + expected.add("BM25"); + expected.add("VECTOR"); + if (trace.getSelectedChannels() == null) { + throw invalid("HYBRID trace selectedChannels is required"); + } + if (trace.getSelectedChannels().contains("GRAPH")) { + expected.add("GRAPH"); + } + if (trace.getSelectedChannels().size() != expected.size() + || !expected.equals(new HashSet<>(trace.getSelectedChannels()))) { + throw invalid("HYBRID trace must select BM25 and VECTOR, with optional GRAPH channel"); + } + validateChannelMapKeys(trace.getChannelBudgets(), expected, "channelBudgets"); + validateChannelMapKeys(trace.getEvaluatedCounts(), expected, "evaluatedCounts"); + validateChannelMapKeys(trace.getCandidateCounts(), expected, "candidateCounts"); + validateChannelMapKeys(trace.getChannelStatuses(), expected, "channelStatuses"); + validateChannelMapKeys(trace.getChannelStopReasons(), expected, "channelStopReasons"); + int budgetSum = 0; + int evaluatedSum = 0; + Set degradedChannels = new HashSet<>(); + for (String channel : expected) { + Integer channelBudget = trace.getChannelBudgets().get(channel); + Integer evaluated = trace.getEvaluatedCounts().get(channel); + Integer candidates = trace.getCandidateCounts().get(channel); + if (channelBudget == null || channelBudget < 1) { + throw invalid("channel budget must be positive: " + channel); + } + if (evaluated == null || evaluated < 0 || evaluated > channelBudget) { + throw invalid("invalid evaluated count: " + channel); + } + if (candidates == null || candidates < 0 + || candidates > ("GRAPH".equals(channel) ? channelBudget : evaluated)) { + throw invalid("invalid candidate count: " + channel); + } + if (trace.getChannelStatuses().get(channel) == null + || trace.getChannelStopReasons().get(channel) == null) { + throw invalid("channel status and stop reason are required: " + channel); + } + budgetSum += channelBudget; + evaluatedSum += evaluated; + boolean degraded = trace.getChannelStatuses().get(channel) != RecallStageStatus.SUCCESS; + if (degraded) { + degradedChannels.add(channel); + } + if (degraded != trace.getDegradedChannels().contains(channel)) { + throw invalid("degradedChannels does not match channel status: " + channel); + } + if (degraded && !hasText(trace.getDegradationReasons().get(channel))) { + throw invalid("degradation reason is required: " + channel); + } + } + if (trace.getTotalCandidatesEvaluated() != evaluatedSum) { + throw invalid("totalCandidatesEvaluated does not match channel counts"); + } + if (degradedChannels.size() != trace.getDegradedChannels().size() + || !degradedChannels.equals(new HashSet<>(trace.getDegradedChannels()))) { + throw invalid("degradedChannels must contain exactly degraded channels"); + } + if (!degradedChannels.equals(trace.getDegradationReasons().keySet())) { + throw invalid("degradationReasons must contain exactly degraded channels"); + } + if (trace.getEffectiveCandidateBudget() != budgetSum + || trace.getEffectiveCandidateBudget() != effectiveBudget.getMaxCandidates() + || trace.getEffectiveTopK() != effectiveBudget.getTopK()) { + throw invalid("invalid effective trace budget"); + } + } + + private static void validateChannelMapKeys(Map values, Set expected, String name) { + if (values == null || !expected.equals(values.keySet())) { + throw invalid(name + " must contain the selected Hybrid channels"); + } + } + + private static boolean hasText(String value) { + return value != null && !value.trim().isEmpty(); + } + public static RetrievalResponse normalizeCollections(RetrievalResponse response) { response.setEvidence(response.getEvidence() == null ? Collections.emptyList() : response.getEvidence()); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalService.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalService.java index 5bfd0e624..5f7fc3684 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalService.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalService.java @@ -85,6 +85,11 @@ public RetrievalResponse retrieve(RetrievalRequest request) { public RetrievalResponse retrieve(RetrievalRequest request, String requestId) { long startedAt = System.nanoTime(); try { + if (request != null && request.getMode() != null + && !"KEYWORD".equalsIgnoreCase(request.getMode().trim())) { + throw new RetrievalException(RetrievalErrorCode.UNSUPPORTED_OPTION, + "HTTP retrieval currently supports KEYWORD only"); + } RetrievalCommand command = validator.validate(request); ServerMemoryCache.ReadyStatus readiness = cache.keywordReadiness(command.getGraphName()); if ("GRAPH_NOT_LOADED".equals(readiness.getReason())) { @@ -119,7 +124,9 @@ public RetrievalResponse retrieve(RetrievalRequest request, String requestId) { response.setEffectiveBudget(budget); RetrievalTrace trace = new RetrievalTrace(); trace.setTraceVersion("v1"); + trace.setRequestId(requestId); trace.setOriginalQuery(command.getQuery()); + trace.setGraphVersion(response.getGraphVersion()); trace.setSelectedMode(command.getMode()); trace.setExecutionMode(command.getExecutionMode()); trace.setStages(Collections.singletonList(new TraceStage("keyword", "SUCCESS", null))); diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java index 60757e216..4bdf14821 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java @@ -103,6 +103,14 @@ void legacyWorkflow() { assertTrue(response.getEvidence().get(0).getFinalScore() > 0); assertTrue(response.getEvidence().get(0).getStageScores().containsKey("keyword")); + for (String mode : new String[] {"BM25_ONLY", "VECTOR_ONLY", "GRAPH_ONLY", "HYBRID"}) { + String unsupportedRequest = request.replace("\"query\":", "\"mode\":\"" + mode + + "\",\"queryVector\":[1,0],\"query\":"); + HttpTestResponse unsupported = client.post("/api/v1/retrievals", unsupportedRequest); + assertEquals(400, unsupported.getStatus(), unsupported.getBody()); + assertTrue(unsupported.getBody().contains("UNSUPPORTED_OPTION"), unsupported.getBody()); + } + HttpTestResponse noMatch = client.post("/api/v1/retrievals", request.replace("Confucius", "not-present")); assertEquals(200, noMatch.getStatus(), noMatch.getBody()); diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java index 5a9e68f07..5aa15291c 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java @@ -38,6 +38,8 @@ import org.apache.geaflow.ai.retrieval.api.model.RetrievalMode; import org.apache.geaflow.ai.retrieval.api.model.TraceStage; import org.apache.geaflow.ai.retrieval.codec.RetrievalApiJson; +import org.apache.geaflow.ai.retrieval.execution.RecallStageStatus; +import org.apache.geaflow.ai.retrieval.execution.RecallStopReason; import org.apache.geaflow.ai.retrieval.model.evidence.Evidence; import org.apache.geaflow.ai.retrieval.model.evidence.EvidenceKind; import org.junit.jupiter.api.Assertions; @@ -46,6 +48,25 @@ /** JSON contract tests independent of Solon and HTTP. */ public class RetrievalApiJsonTest { + @Test + void hybridContractAcceptsValidVectorAndRejectsMalformedInput() { + String valid = "{\"graphName\":\"graph\",\"query\":\"query\",\"mode\":\"HYBRID\",\"queryVector\":[1,0]}"; + org.apache.geaflow.ai.retrieval.service.RetrievalRequestValidator validator = + new org.apache.geaflow.ai.retrieval.service.RetrievalRequestValidator(new org.apache.geaflow.ai.retrieval.config.RetrievalProperties()); + Assertions.assertEquals(org.apache.geaflow.ai.retrieval.api.model.RetrievalMode.HYBRID, + validator.validate(RetrievalApiJson.parseRequest(valid)).getMode()); + for (String vector : new String[] {"[null]", "[\"1\"]", "{}", "[1e400]"}) { + Assertions.assertThrows(com.google.gson.JsonParseException.class, + () -> RetrievalApiJson.parseRequest(valid.replace("[1,0]", vector))); + } + for (String vector : new String[] {"[]", "[0,0]", "[1e100,0]"}) { + org.apache.geaflow.ai.retrieval.api.model.RetrievalException error = Assertions.assertThrows( + org.apache.geaflow.ai.retrieval.api.model.RetrievalException.class, + () -> validator.validate(RetrievalApiJson.parseRequest(valid.replace("[1,0]", vector)))); + Assertions.assertEquals(org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.INVALID_REQUEST, error.getCode()); + } + } + @Test public void parsesMinimalAndFullRequests() { RetrievalRequest minimal = RetrievalApiJson.parseRequest( @@ -122,10 +143,40 @@ public void responseKeepsRequiredCollectionsAndRoundTrips() { RetrievalTrace trace = new RetrievalTrace(); trace.setTraceVersion("v1"); trace.setOriginalQuery("confucius"); - trace.setSelectedMode(RetrievalMode.KEYWORD); + trace.setSelectedMode(RetrievalMode.HYBRID); trace.setExecutionMode(ExecutionMode.SEQUENTIAL); trace.setStages(Collections.singletonList(new TraceStage("keyword", "COMPLETED", null))); trace.setStopReason("COMPLETED"); + trace.setRequestId("req-1"); + trace.setGraphVersion("v1"); + trace.setIndexVersion("v1"); + trace.setSelectedChannels(java.util.Arrays.asList("BM25", "VECTOR")); + Map channelBudgets = new HashMap<>(); + channelBudgets.put("BM25", 50); + channelBudgets.put("VECTOR", 50); + trace.setChannelBudgets(channelBudgets); + Map candidateCounts = new HashMap<>(); + candidateCounts.put("BM25", 2); + candidateCounts.put("VECTOR", 2); + trace.setCandidateCounts(candidateCounts); + Map evaluatedCounts = new HashMap<>(); + evaluatedCounts.put("BM25", 2); + evaluatedCounts.put("VECTOR", 2); + trace.setEvaluatedCounts(evaluatedCounts); + Map channelStatuses = new HashMap<>(); + channelStatuses.put("BM25", RecallStageStatus.SUCCESS); + channelStatuses.put("VECTOR", RecallStageStatus.SUCCESS); + trace.setChannelStatuses(channelStatuses); + Map channelStopReasons = new HashMap<>(); + channelStopReasons.put("BM25", RecallStopReason.COMPLETED); + channelStopReasons.put("VECTOR", RecallStopReason.COMPLETED); + trace.setChannelStopReasons(channelStopReasons); + trace.setEffectiveCandidateBudget(100); + trace.setEffectiveTopK(10); + trace.setTotalCandidatesEvaluated(4); + trace.setBm25Weight(1.0); + trace.setVectorWeight(1.0); + trace.setRrfRankConstant(60); response.setTrace(trace); response.setEffectiveBudget(new RetrievalBudget(10, 3000, 100, 4096)); @@ -138,8 +189,45 @@ public void responseKeepsRequiredCollectionsAndRoundTrips() { Assertions.assertNotNull(restored.getSources()); Assertions.assertNotNull(restored.getDegradedChannels()); Assertions.assertEquals("v1", restored.getTrace().getTraceVersion()); + Assertions.assertEquals("req-1", restored.getTrace().getRequestId()); + Assertions.assertEquals(4, restored.getTrace().getTotalCandidatesEvaluated()); + Assertions.assertEquals(java.util.Arrays.asList("BM25", "VECTOR"), + restored.getTrace().getSelectedChannels()); Assertions.assertEquals(Integer.valueOf(4096), restored.getEffectiveBudget().getTokenBudget()); + + response.getTrace().getCandidateCounts().put("BM25", 3); + Assertions.assertThrows(RuntimeException.class, () -> RetrievalApiJson.toJson(response)); + } + + @Test + public void rejectsIncompleteHybridTraceContract() { + RetrievalResponse response = new RetrievalResponse(); + response.setRequestId("req-hybrid"); + response.setGraphName("graph"); + response.setGraphVersion("v1"); + response.setEffectiveBudget(new RetrievalBudget(1, 100, 2, 10)); + RetrievalTrace trace = new RetrievalTrace(); + trace.setTraceVersion("v1"); + trace.setRequestId("req-hybrid"); + trace.setOriginalQuery("q"); + trace.setSelectedMode(RetrievalMode.HYBRID); + trace.setExecutionMode(ExecutionMode.SEQUENTIAL); + trace.setGraphVersion("v1"); + trace.setIndexVersion("v1"); + trace.setSelectedChannels(java.util.Arrays.asList("BM25", "VECTOR")); + trace.setChannelBudgets(Collections.singletonMap("BM25", 2)); + response.setTrace(trace); + + Assertions.assertThrows(RuntimeException.class, + () -> RetrievalApiJson.toJson(response)); + } + + @Test + public void rejectsMismatchedTraceRequestId() { + String json = responseJson("KEYWORD", "SEQUENTIAL", 10, 3000, 100, 4096) + .replace("\"stages\":[]", "\"requestId\":\"other\",\"stages\":[]"); + Assertions.assertThrows(RuntimeException.class, () -> RetrievalApiJson.parseResponse(json)); } @Test @@ -148,6 +236,7 @@ public void emptyResponseUsesArraysInsteadOfNullCollections() { + "\"graphVersion\":\"v1\",\"evidence\":[],\"paths\":[]," + "\"sources\":[],\"degradedChannels\":[]," + "\"trace\":{\"traceVersion\":\"v1\"," + + "\"requestId\":\"req-2\"," + "\"originalQuery\":\"none\",\"selectedMode\":\"KEYWORD\"," + "\"executionMode\":\"SEQUENTIAL\",\"stages\":[]}," + "\"effectiveBudget\":{\"topK\":10,\"timeoutMs\":3000," diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalRequestValidatorTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalRequestValidatorTest.java index 1880d1fb8..15801cadc 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalRequestValidatorTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalRequestValidatorTest.java @@ -88,7 +88,7 @@ public void rejectsMissingAndOverlongText() { public void rejectsUnsupportedModeExecutionAndVector() { RetrievalRequest mode = request(); mode.setMode("HYBRID"); - assertCode(mode, RetrievalErrorCode.UNSUPPORTED_OPTION); + assertCode(mode, RetrievalErrorCode.INVALID_REQUEST); RetrievalRequest execution = request(); execution.setExecutionMode("PARALLEL"); @@ -107,6 +107,32 @@ public void rejectsUnsupportedModeExecutionAndVector() { Assertions.assertDoesNotThrow(() -> validator().validate(emptyVector)); } + @Test + public void acceptsHybridVectorAndCaseInsensitiveExecutionMode() { + RetrievalRequest request = request(); + request.setMode("hybrid"); + request.setExecutionMode(" sequential "); + request.setQueryVector(Collections.singletonList(0.1)); + + RetrievalCommand command = validator().validate(request); + + Assertions.assertEquals(RetrievalMode.HYBRID, command.getMode()); + Assertions.assertEquals(ExecutionMode.SEQUENTIAL, command.getExecutionMode()); + } + + @Test + public void rejectsVectorsForNonVectorModesAndBlankMetadata() { + for (String mode : new String[] {"BM25_ONLY", "KEYWORD", "GRAPH_ONLY"}) { + RetrievalRequest request = request(); + request.setMode(mode); + request.setQueryVector(Collections.singletonList(0.1)); + assertCode(request, RetrievalErrorCode.UNSUPPORTED_OPTION); + } + RetrievalRequest invalid = request(); + invalid.setRequestId(" "); + assertCode(invalid, RetrievalErrorCode.INVALID_REQUEST); + } + @Test public void rejectsBudgetRangesAndRelationships() { assertBudget(new RetrievalBudget(0, 3000, 100, 4096)); diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/RetrievalServiceTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/RetrievalServiceTest.java index c1de6d2a2..bed731116 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/RetrievalServiceTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/RetrievalServiceTest.java @@ -119,6 +119,21 @@ void distinguishesLoadedGraphWithUnavailableIndex() { assertEquals(RetrievalErrorCode.INDEX_NOT_READY, exception.getCode()); } + @Test + void httpServiceRejectsCoreModesBeforeReadinessOrKeywordExecution() { + RetrievalService service = new RetrievalService(new ServerMemoryCache(), new RetrievalProperties()); + for (String mode : new String[] {"BM25_ONLY", "VECTOR_ONLY", "GRAPH_ONLY", "HYBRID", "unknown"}) { + RetrievalRequest request = new RetrievalRequest(); + request.setGraphName("week1-service-graph"); + request.setQuery("Confucius"); + request.setMode(mode); + request.setQueryVector(java.util.Arrays.asList(1.0, 0.0)); + RetrievalException error = org.junit.jupiter.api.Assertions.assertThrows(RetrievalException.class, + () -> service.retrieve(request)); + assertEquals(RetrievalErrorCode.UNSUPPORTED_OPTION, error.getCode()); + } + } + private static MemoryGraph createGraph() { VertexSchema schema = new VertexSchema("chunk", "id", Collections.singletonList("text")); GraphSchema graphSchema = new GraphSchema(); diff --git a/geaflow-ai/src/test/resources/retrieval/api/response-empty.json b/geaflow-ai/src/test/resources/retrieval/api/response-empty.json index b5649e936..2ca7db8f9 100644 --- a/geaflow-ai/src/test/resources/retrieval/api/response-empty.json +++ b/geaflow-ai/src/test/resources/retrieval/api/response-empty.json @@ -7,6 +7,7 @@ "sources": [], "trace": { "traceVersion": "v1", + "requestId": "request-1", "originalQuery": "unknown", "selectedMode": "KEYWORD", "executionMode": "SEQUENTIAL", From 98cd90f13aa541d613f21ce4a4a4ab323b1c5147 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Thu, 8 Oct 2026 21:03:47 +0800 Subject: [PATCH 24/30] docs(ai): document fixed hybrid recall MVP --- geaflow-ai/README-hybrid-recall.md | 127 +++++++++++++++++++++++++++ geaflow-ai/README-week1-retrieval.md | 3 + 2 files changed, 130 insertions(+) create mode 100644 geaflow-ai/README-hybrid-recall.md diff --git a/geaflow-ai/README-hybrid-recall.md b/geaflow-ai/README-hybrid-recall.md new file mode 100644 index 000000000..ccb8d275b --- /dev/null +++ b/geaflow-ai/README-hybrid-recall.md @@ -0,0 +1,127 @@ + + +# Fixed Hybrid Recall MVP + +`RecallService` is a session-free Java facade over offline, versioned fixtures. Its structured request +supports `BM25_ONLY`, `VECTOR_ONLY`, `GRAPH_ONLY`, and `HYBRID`; `KEYWORD` aliases `BM25_ONLY` in this core +facade. The existing HTTP `/api/v1/retrievals` endpoint supports only `KEYWORD`. Other modes return +`UNSUPPORTED_OPTION` there; HTTP Hybrid serving is not part of this implementation. + +## Load and retrieve + +Build both artifacts with `LuceneBm25IndexBuilder` and `OfflineVectorIndexBuilder`, using the same +`IngestionContext` and canonical chunks. Supply precomputed vectors; no embedding service is called. +Load the published artifacts and graph records once: + +```java +RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); +registry.load(graphName, graphVersion, indexVersion, bm25Path, vectorPath, + chunks, entities, edges, vectorSource, vectorVersion); +RecallService service = new RecallService(registry, new WeightedRrfConfig(1.0, 1.0, 1.0, 60)); + +RetrievalRequest request = new RetrievalRequest(); +request.setGraphName(graphName); +request.setGraphVersion(graphVersion); +request.setIndexVersion(indexVersion); +request.setMode("HYBRID"); +request.setQuery("Confucius"); +request.setQueryVector(Arrays.asList(1.0, 0.0)); // Must match the artifact's dimensions. +request.setBudget(new RetrievalBudget(10, 3000, 100, 4096)); +request.setAllowPartialResults(true); +RetrievalResponse result = service.retrieve(request); +``` + +`graphVersion` and `indexVersion` are required and must match. The registry validates the persisted +artifact graph name, graph/index versions, and a fingerprint of the complete canonical chunk payload. +Optional request vector source/version constraints must match the loaded vector artifact. Keep the +readers alive while their fixtures are in use; callers own the lifecycle of registered readers. +The legacy overload accepting text, a float vector, and an absolute deadline requires a service +constructed directly with readers, rather than a registry-backed service. + +## Budgets, ranking, and limitations + +Hybrid runs BM25, Vector, and one-hop Graph under one monotonic deadline starting at request entry. It +allocates a deterministic three-way budget: BM25 receives the base plus the first remainder, Vector +receives the base plus the second remainder, and Graph receives the base, where `base = floor(maxCandidates / 3)`. +Hybrid requires `maxCandidates >= 3`; all modes require `1 <= topK <= maxCandidates`. +`topK` is applied only after merging and ranking. Individual retriever convenience methods still +return at most their requested `topK`; the facade uses the full bounded candidate lists for fusion. + +`maxCandidates` caps actual scoring work: matched Lucene documents scored for BM25, vectors evaluated +for Vector, and edges examined for GRAPH_ONLY. Lucene uses a collector that checks the deadline before +scoring each matched document. Parsing, query setup, and sorting are checked at stage boundaries; +deadlines are cooperative, rather than a thread-interruption guarantee. + +Both text and vector channels are bounded approximate recall when their evaluation budget is smaller +than the searchable data. BM25 collects Lucene traversal order; Vector scans stable chunk-ID order and +sorts only the evaluated prefix by cosine similarity. The vector result is **not** a global nearest +neighbor guarantee. Increasing the budget widens the evaluated set. ANN search and adaptive candidate +selection remain future work. + +Weighted RRF uses `sum(weight[channel] / (rankConstant + rank[channel]))`, with defaults `1.0`, `1.0`, +`1.0`, and `60`. Weights must be finite and non-negative, at least one positive; the rank constant is positive. +Zero-weight channels still run and retain their scores/provenance. Evidence merges by canonical +document/chunk identity, uses stable evidence ID for ties, and receives consecutive final ranks. + +GRAPH_ONLY retains deterministic anchor resolution and one-hop expansion. Both edge scans and unique +chunk evidence are capped by `maxCandidates`, including edges that cite several chunks. Paths and +anchor provenance survive merging. Hybrid always selects BM25, Vector, and Graph; graph provenance +contains the stable query, graph version, anchor match, path, and supporting chunk IDs. + +## Errors and trace + +Missing, empty, non-finite, float-overflowing, zero-norm, or dimension-mismatched query vectors return +`INVALID_REQUEST` for vector modes. Unknown modes and non-sequential execution return +`UNSUPPORTED_OPTION`. Malformed JSON is rejected at the codec boundary. + +The loader records channel-specific missing/corrupt artifact failures as `INDEX_NOT_READY` and only +publishes a READY fixture when at least one channel is usable. With partial results disabled, a failed +selected channel fails the request. With partial results enabled, its evidence is omitted and the +successful channel's evidence remains. If the request deadline expires, strict requests return +`RETRIEVAL_TIMEOUT`; partial requests retain fully evaluated evidence and mark unfinished channels. + +Trace fields include versions, selected channels, effective topK, per-channel budgets, actual evaluated +counts, candidate counts before fusion/topK, graph anchors/edges/neighbors/vertices, validation errors, +elapsed nanoseconds, and the effective three-channel RRF configuration. +`channelStatuses` distinguishes `SUCCESS`, `DEGRADED`, and `NOT_RUN`; `channelStopReasons` carries typed +`COMPLETED`, `CANDIDATE_BUDGET`, `CANDIDATE_LIMIT`, `EDGE_SCAN_LIMIT`, `NO_RELIABLE_ANCHOR`, `DEADLINE`, +or `INDEX_NOT_READY`. `degradationReasons` carries explanatory text separately. Overall `stopReason` +prioritizes deadline, graph trace validation, graph edge/candidate limits, candidate budget, then final topK +truncation, then no results or normal completion. Candidate-budget exhaustion is a valid bounded stage, +not an index failure. A fusion-stage deadline is also reflected in its stage status. + +## Artifact compatibility and verification + +BM25 stores identity in Lucene commit metadata. Vector artifacts use `GEAFLOW-VECTOR-2`, adding identity +before vector source/version and the vector payload. Legacy artifacts without this metadata are +rejected with `INDEX_NOT_READY`; rebuild and publish a fresh graph/index version before switching +requests. Existing immutable artifacts are validated rather than silently overwritten. + +Run the retrieval and HTTP regression suite with JDK 11, which is compatible with the repository's +JaCoCo 0.8.8: + +```bash +mvn -pl geaflow-ai -Dtest='org.apache.geaflow.ai.retrieval.**.*Test,MemoryServerTest' test +``` + +The Maven lifecycle also runs Checkstyle and Apache RAT. New retrieval production sources have no +blanket Checkstyle suppressions. The suite covers bounded scoring, partial-channel failures, +deadlines using a controlled clock, graph limits, version mismatches, JSON trace round trips, +concurrent request isolation, and rejection of core modes at the HTTP boundary. diff --git a/geaflow-ai/README-week1-retrieval.md b/geaflow-ai/README-week1-retrieval.md index 590868824..586efddb0 100644 --- a/geaflow-ai/README-week1-retrieval.md +++ b/geaflow-ai/README-week1-retrieval.md @@ -39,3 +39,6 @@ curl -i http://localhost:8080/health curl -i http://localhost:8080/ready curl -sS http://localhost:8080/metrics/retrieval ``` + +The HTTP endpoint supports `KEYWORD` only. Offline BM25, vector, graph-only, and Hybrid retrieval +are available through the [session-free core facade](README-hybrid-recall.md). From 3676d433b12ad326d98271c70e3c38590121aa2f Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 10 Oct 2026 20:30:47 +0800 Subject: [PATCH 25/30] fix(ai): make ingestion and artifact publication exclusive --- .../ai/retrieval/index/ArtifactIdentity.java | 1 + .../ai/retrieval/index/ArtifactPublisher.java | 89 ++++++++ .../index/bm25/LuceneBm25IndexBuilder.java | 16 +- .../graph/OfflineGraphArtifactBuilder.java | 13 +- .../vector/OfflineVectorIndexBuilder.java | 18 +- .../ai/retrieval/ingest/ImportAttempt.java | 44 ++++ .../retrieval/ingest/IngestionPipeline.java | 43 +++- .../retrieval/ingest/MetadataPublisher.java | 6 +- .../ingest/MetadataStorePublisher.java | 40 ++-- .../index/ArtifactPublicationTest.java | 216 ++++++++++++++++++ .../ingest/IngestionPipelineTest.java | 169 +++++++++++++- 11 files changed, 586 insertions(+), 69 deletions(-) create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactPublisher.java create mode 100644 geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ImportAttempt.java create mode 100644 geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/ArtifactPublicationTest.java diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentity.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentity.java index 5c9015148..441decdec 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentity.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactIdentity.java @@ -105,6 +105,7 @@ public static String fingerprint(List chunks) { append(digest, Integer.toString(chunk.getTokenEstimate())); append(digest, chunk.getPolicyVersion()); append(digest, chunk.getTextHash()); + append(digest, chunk.getSourceUri()); append(digest, Integer.toString(chunk.getStartOffset())); append(digest, Integer.toString(chunk.getEndOffset())); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactPublisher.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactPublisher.java new file mode 100644 index 000000000..235526a2c --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/ArtifactPublisher.java @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index; + +import java.io.IOException; +import java.nio.channels.FileChannel; +import java.nio.channels.FileLock; +import java.nio.file.FileAlreadyExistsException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.StandardCopyOption; +import java.nio.file.StandardOpenOption; +import java.util.concurrent.locks.ReentrantLock; + +/** Serializes artifact publication per target and never replaces an existing artifact. */ +public final class ArtifactPublisher { + + private static final ReentrantLock[] LOCKS = createLocks(); + + private ArtifactPublisher() { + } + + /** + * Publishes a completed staging path, or validates and reuses the winner if it already exists. + * + * @return true when this call moved the staging artifact into place + */ + public static boolean publish(Path staging, Path target, ExistingArtifactValidator validator) + throws IOException { + // Canonicalize the parent so aliases of one output directory share the same JVM lock. + Path normalizedTarget = target.toAbsolutePath().getParent().toRealPath() + .resolve(target.getFileName()); + Path lockPath = normalizedTarget.resolveSibling(normalizedTarget.getFileName() + ".publish.lock"); + ReentrantLock processLock = LOCKS[(normalizedTarget.hashCode() & Integer.MAX_VALUE) % LOCKS.length]; + processLock.lock(); + try { + try (FileChannel lockChannel = FileChannel.open(lockPath, StandardOpenOption.CREATE, + StandardOpenOption.WRITE); FileLock ignored = lockChannel.lock()) { + if (Files.exists(normalizedTarget)) { + validator.validate(normalizedTarget); + return false; + } + try { + try { + Files.move(staging, normalizedTarget, StandardCopyOption.ATOMIC_MOVE); + } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { + Files.move(staging, normalizedTarget); + } + return true; + } catch (FileAlreadyExistsException winnerPublished) { + validator.validate(normalizedTarget); + return false; + } + } + } finally { + processLock.unlock(); + } + } + + private static ReentrantLock[] createLocks() { + ReentrantLock[] locks = new ReentrantLock[64]; + for (int index = 0; index < locks.length; index++) { + locks[index] = new ReentrantLock(); + } + return locks; + } + + /** Checks whether an existing target is equivalent to the artifact being built. */ + public interface ExistingArtifactValidator { + void validate(Path existing) throws IOException; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java index 290fffb54..c69373434 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/bm25/LuceneBm25IndexBuilder.java @@ -22,7 +22,6 @@ import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; -import java.nio.file.StandardCopyOption; import java.util.ArrayList; import java.util.Comparator; import java.util.HashSet; @@ -30,6 +29,7 @@ import java.util.Objects; import java.util.Set; import org.apache.geaflow.ai.retrieval.index.ArtifactIdentity; +import org.apache.geaflow.ai.retrieval.index.ArtifactPublisher; import org.apache.geaflow.ai.retrieval.index.Bm25IndexBuilder; import org.apache.geaflow.ai.retrieval.index.IndexArtifact; import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; @@ -103,18 +103,8 @@ public IndexArtifact build(IngestionContext context, List chunks) thr } } } - if (Files.exists(published)) { - validateExisting(published, ordered, context); - return new LuceneIndexArtifact(new IndexBuildMetadata(context.getGraphVersion(), - new IndexVersion(INDEX_NAME, version, version), INDEX_NAME, BUILDER_VERSION, - published.toString(), true), published); - } - try { - Files.move(staging, published, StandardCopyOption.ATOMIC_MOVE); - } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { - Files.move(staging, published); - } - publishedStaging = true; + publishedStaging = ArtifactPublisher.publish(staging, published, + existing -> validateExisting(existing, ordered, context)); IndexVersion indexVersion = new IndexVersion(INDEX_NAME, version, version); IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, INDEX_NAME, BUILDER_VERSION, published.toString(), true); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/OfflineGraphArtifactBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/OfflineGraphArtifactBuilder.java index c796cb27e..4cf6a26b9 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/OfflineGraphArtifactBuilder.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/graph/OfflineGraphArtifactBuilder.java @@ -29,6 +29,7 @@ import java.util.List; import java.util.Set; import org.apache.geaflow.ai.retrieval.index.ArtifactIdentity; +import org.apache.geaflow.ai.retrieval.index.ArtifactPublisher; import org.apache.geaflow.ai.retrieval.ingest.GraphArtifact; import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; import org.apache.geaflow.ai.retrieval.metadata.GraphBuildMetadata; @@ -80,16 +81,8 @@ public GraphArtifact build(IngestionContext context, List chunks, writeStrings(output, edge.getSourceChunkIds()); } } - if (Files.exists(published)) { - validateExisting(published, context, orderedEntities, orderedEdges, chunks); - return new Artifact(new GraphBuildMetadata(context.getGraphVersion(), published.toString(), - "offline-graph", "offline-graph-v1", true), published); - } - try { - Files.move(staging, published, java.nio.file.StandardCopyOption.ATOMIC_MOVE); - } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { - Files.move(staging, published); - } + ArtifactPublisher.publish(staging, published, + existing -> validateExisting(existing, context, orderedEntities, orderedEdges, chunks)); return new Artifact(new GraphBuildMetadata(context.getGraphVersion(), published.toString(), "offline-graph", "offline-graph-v1", true), published); } finally { diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java index c9596a3d5..f239541a5 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/index/vector/OfflineVectorIndexBuilder.java @@ -32,6 +32,7 @@ import java.util.Objects; import java.util.Set; import org.apache.geaflow.ai.retrieval.index.ArtifactIdentity; +import org.apache.geaflow.ai.retrieval.index.ArtifactPublisher; import org.apache.geaflow.ai.retrieval.index.IndexArtifact; import org.apache.geaflow.ai.retrieval.index.VectorIndexBuilder; import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; @@ -89,20 +90,9 @@ public IndexArtifact build(IngestionContext context, List chunks) thr } } } - if (Files.exists(published)) { - VectorArtifact.validateArtifact(published, ordered, dimensions, vectorSource, vectorVersion, context, vectors); - return new VectorArtifact(new IndexBuildMetadata(context.getGraphVersion(), - new IndexVersion(INDEX_NAME, context.getGraphVersion().getVersion(), - context.getGraphVersion().getVersion()), INDEX_NAME, - BUILDER_VERSION + ":" + vectorSource + ":" + vectorVersion, - published.toString(), true), published); - } - try { - Files.move(staging, published, java.nio.file.StandardCopyOption.ATOMIC_MOVE); - } catch (java.nio.file.AtomicMoveNotSupportedException unsupported) { - Files.move(staging, published); - } - publishedStaging = true; + publishedStaging = ArtifactPublisher.publish(staging, published, + existing -> VectorArtifact.validateArtifact(existing, ordered, dimensions, + vectorSource, vectorVersion, context, vectors)); IndexVersion indexVersion = new IndexVersion(INDEX_NAME, context.getGraphVersion().getVersion(), context.getGraphVersion().getVersion()); IndexBuildMetadata metadata = new IndexBuildMetadata(context.getGraphVersion(), indexVersion, diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ImportAttempt.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ImportAttempt.java new file mode 100644 index 000000000..cdc346092 --- /dev/null +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/ImportAttempt.java @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.ingest; + +import java.util.Objects; +import java.util.UUID; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; + +/** Ownership token for one ingestion lifecycle attempt. */ +public final class ImportAttempt { + + private final GraphVersion graphVersion; + private final String attemptId; + + public ImportAttempt(GraphVersion graphVersion) { + this.graphVersion = Objects.requireNonNull(graphVersion, "graphVersion"); + this.attemptId = UUID.randomUUID().toString(); + } + + public GraphVersion getGraphVersion() { + return graphVersion; + } + + String getAttemptId() { + return attemptId; + } +} diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java index d3ca791ec..38cf24e54 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipeline.java @@ -22,7 +22,9 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Objects; import org.apache.geaflow.ai.retrieval.index.Bm25IndexBuilder; import org.apache.geaflow.ai.retrieval.index.IndexArtifact; @@ -61,10 +63,13 @@ public ImportMetadata run(IngestionContext context) throws IOException { Objects.requireNonNull(context, "context"); List graphArtifacts = new ArrayList<>(); List indexArtifacts = new ArrayList<>(); + ImportAttempt attempt = null; try { - metadataPublisher.begin(context); + attempt = metadataPublisher.begin(context); + Objects.requireNonNull(attempt, "metadata publisher returned a null attempt"); List documents = requireList(sourceLoader.load(context), "documents"); - List chunks = requireList(normalizer.normalize(documents, context), "chunks"); + List chunks = withSourceUris(documents, + requireList(normalizer.normalize(documents, context), "chunks")); ExtractionResult extraction = Objects.requireNonNull( entityExtractor.extract(chunks, context), "extraction"); GraphArtifact graph = Objects.requireNonNull( @@ -76,17 +81,19 @@ public ImportMetadata run(IngestionContext context) throws IOException { indexArtifacts.add(vector); context.setQualityCounters(new QualityCounters(documents.size(), chunks.size(), extraction.getEntities().size(), extraction.getEdges().size(), 0, 0)); - return metadataPublisher.publish(context, graph, + return metadataPublisher.publish(context, attempt, graph, Collections.unmodifiableList(indexArtifacts)); } catch (Exception failure) { closeAll(indexArtifacts, failure); indexArtifacts.clear(); closeAll(graphArtifacts, failure); graphArtifacts.clear(); - try { - metadataPublisher.fail(context, failure); - } catch (IOException publicationFailure) { - failure.addSuppressed(publicationFailure); + if (attempt != null) { + try { + metadataPublisher.fail(context, attempt, failure); + } catch (Exception publicationFailure) { + failure.addSuppressed(publicationFailure); + } } if (failure instanceof IOException) { throw (IOException) failure; @@ -109,6 +116,28 @@ private static List requireList(List values, String name) { return values; } + private static List withSourceUris(List documents, List chunks) { + Map byId = new HashMap<>(); + for (SourceDocument document : documents) { + if (byId.put(document.getDocumentId(), document) != null) { + throw new IllegalArgumentException("duplicate source document ID"); + } + } + List result = new ArrayList<>(chunks.size()); + for (TextChunk chunk : chunks) { + SourceDocument document = byId.get(chunk.getDocumentId()); + if (document == null || document.getSourceUri() == null + || document.getSourceUri().trim().isEmpty()) { + throw new IllegalArgumentException("chunk requires a source document with a source URI"); + } + if (chunk.getSourceUri() != null && !document.getSourceUri().equals(chunk.getSourceUri())) { + throw new IllegalArgumentException("chunk source URI does not match its document"); + } + result.add(chunk.withSourceUri(document.getSourceUri())); + } + return result; + } + private static void closeAll(List resources, Exception failure) throws IOException { IOException closeFailure = null; diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataPublisher.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataPublisher.java index 0cd58af14..f0ca2eec9 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataPublisher.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataPublisher.java @@ -27,10 +27,10 @@ /** Owns import state changes and atomic publication of completed artifacts. */ public interface MetadataPublisher { - void begin(IngestionContext context) throws IOException; + ImportAttempt begin(IngestionContext context) throws IOException; - ImportMetadata publish(IngestionContext context, GraphArtifact graph, + ImportMetadata publish(IngestionContext context, ImportAttempt attempt, GraphArtifact graph, List indexes) throws IOException; - void fail(IngestionContext context, Exception failure) throws IOException; + void fail(IngestionContext context, ImportAttempt attempt, Exception failure) throws IOException; } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataStorePublisher.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataStorePublisher.java index f4a03e60c..f7c574297 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataStorePublisher.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/ingest/MetadataStorePublisher.java @@ -49,7 +49,7 @@ public final class MetadataStorePublisher implements MetadataPublisher { private final MetadataStore store; private final SourceStreamProvider sourceProvider; - private final Map attempts = new HashMap<>(); + private final Map attempts = new HashMap<>(); public MetadataStorePublisher(MetadataStore store, SourceStreamProvider sourceProvider) { this.store = Objects.requireNonNull(store, "store"); @@ -57,28 +57,31 @@ public MetadataStorePublisher(MetadataStore store, SourceStreamProvider sourcePr } @Override - public synchronized void begin(IngestionContext context) { + public synchronized ImportAttempt begin(IngestionContext context) { Objects.requireNonNull(context, "context"); GraphVersion version = context.getGraphVersion(); if (attempts.containsKey(version)) { throw new MetadataException(MetadataException.Code.VERSION_CONFLICT, "ingestion attempt already started"); } + ImportAttempt token = new ImportAttempt(version); GraphBuildMetadata graph = new GraphBuildMetadata(version, null, "graph", "pending", false, REQUIRED_INDEXES); GraphVersion expected = store.getPublishedVersion(version.getGraphName()).orElse(null); store.begin(context.getManifest(), context.getImporterVersion(), graph); - attempts.put(version, new Attempt(expected)); + attempts.put(version, new AttemptState(token, expected)); + return token; } @Override - public synchronized ImportMetadata publish(IngestionContext context, GraphArtifact graph, + public synchronized ImportMetadata publish(IngestionContext context, ImportAttempt token, GraphArtifact graph, List indexes) throws IOException { Objects.requireNonNull(context, "context"); + Objects.requireNonNull(token, "attempt"); Objects.requireNonNull(graph, "graph"); Objects.requireNonNull(indexes, "indexes"); GraphVersion version = context.getGraphVersion(); - final Attempt attempt = requireAttempt(version); + final AttemptState attempt = requireAttempt(version, token); GraphBuildMetadata graphMetadata = graph.getMetadata(); if (!version.equals(graphMetadata.getGraphVersion()) || !REQUIRED_INDEXES.equals(graphMetadata.getRequiredIndexes())) { @@ -99,14 +102,17 @@ public synchronized ImportMetadata publish(IngestionContext context, GraphArtifa } @Override - public synchronized void fail(IngestionContext context, Exception failure) throws IOException { + public synchronized void fail(IngestionContext context, ImportAttempt token, Exception failure) + throws IOException { Objects.requireNonNull(context, "context"); + Objects.requireNonNull(token, "attempt"); Objects.requireNonNull(failure, "failure"); GraphVersion version = context.getGraphVersion(); - Attempt attempt = attempts.remove(version); - if (attempt == null) { + AttemptState attempt = attempts.get(version); + if (attempt == null || !sameToken(attempt.token, token)) { return; } + attempts.remove(version); ImportMetadata current = store.find(version).orElse(null); if (current == null || current.getState() == ImportState.FAILED || current.getState() == ImportState.READY) { @@ -126,19 +132,25 @@ private InputStream openSource(IngestionContext context) throws IOException { return source; } - private Attempt requireAttempt(GraphVersion version) { - Attempt attempt = attempts.get(version); - if (attempt == null) { + private AttemptState requireAttempt(GraphVersion version, ImportAttempt token) { + AttemptState attempt = attempts.get(version); + if (attempt == null || !sameToken(attempt.token, token)) { throw new MetadataException(MetadataException.Code.INVALID_TRANSITION, - "ingestion attempt has not started"); + "ingestion attempt does not own this version"); } return attempt; } - private static final class Attempt { + private static boolean sameToken(ImportAttempt expected, ImportAttempt actual) { + return expected == actual && expected.getAttemptId().equals(actual.getAttemptId()); + } + + private static final class AttemptState { + private final ImportAttempt token; private final GraphVersion expectedPublishedVersion; - private Attempt(GraphVersion expectedPublishedVersion) { + private AttemptState(ImportAttempt token, GraphVersion expectedPublishedVersion) { + this.token = token; this.expectedPublishedVersion = expectedPublishedVersion; } } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/ArtifactPublicationTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/ArtifactPublicationTest.java new file mode 100644 index 000000000..b889ce21a --- /dev/null +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/ArtifactPublicationTest.java @@ -0,0 +1,216 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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. + */ + +package org.apache.geaflow.ai.retrieval.index; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; +import org.apache.geaflow.ai.retrieval.index.bm25.Bm25IndexReader; +import org.apache.geaflow.ai.retrieval.index.bm25.LuceneBm25IndexBuilder; +import org.apache.geaflow.ai.retrieval.index.graph.GraphArtifactReader; +import org.apache.geaflow.ai.retrieval.index.graph.OfflineGraphArtifactBuilder; +import org.apache.geaflow.ai.retrieval.index.vector.OfflineVectorIndexBuilder; +import org.apache.geaflow.ai.retrieval.index.vector.VectorIndexReader; +import org.apache.geaflow.ai.retrieval.ingest.IngestionContext; +import org.apache.geaflow.ai.retrieval.metadata.ChunkingConfiguration; +import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; +import org.apache.geaflow.ai.retrieval.model.document.TextChunk; +import org.apache.geaflow.ai.retrieval.model.graph.EntityRef; +import org.apache.geaflow.ai.retrieval.model.version.GraphVersion; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class ArtifactPublicationTest { + @TempDir + Path directory; + + @Test + void concurrentEquivalentBuildersReuseWinnerForEveryArtifactType() throws Exception { + for (String kind : new String[] {"vector", "graph", "bm25"}) { + Path output = directory.resolve(kind); + runRace(kind, output, false); + build(kind, output, 0); + validateReadable(kind, output); + assertNoStaging(output); + } + } + + @Test + void concurrentConflictingBuildersCannotOverwriteWinnerForEveryArtifactType() throws Exception { + for (String kind : new String[] {"vector", "graph", "bm25"}) { + Path output = directory.resolve(kind); + int winner = runRace(kind, output, true); + build(kind, output, winner); + assertThrows(Exception.class, () -> build(kind, output, 1 - winner)); + validateReadable(kind, output); + assertNoStaging(output); + } + } + + @Test + void separateJvmsSerializePublicationAndValidateTheWinner() throws Exception { + Path output = Files.createDirectory(directory.resolve("processes")); + String java = Paths.get(System.getProperty("java.home"), "bin", "java").toString(); + String classpath = System.getProperty("surefire.test.class.path", System.getProperty("java.class.path")); + Process first = new ProcessBuilder(java, "-cp", classpath, PublisherProcess.class.getName(), + output.toString(), "0").redirectErrorStream(true).redirectOutput(output.resolve("first.log").toFile()).start(); + Process second = new ProcessBuilder(java, "-cp", classpath, PublisherProcess.class.getName(), + output.toString(), "1").redirectErrorStream(true).redirectOutput(output.resolve("second.log").toFile()).start(); + try { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (!Files.exists(output.resolve("ready-0")) || !Files.exists(output.resolve("ready-1"))) { + if (System.nanoTime() > deadline) { + throw new TimeoutException("publisher processes did not start"); + } + Thread.sleep(10); + } + Files.createFile(output.resolve("start")); + org.junit.jupiter.api.Assertions.assertTrue(first.waitFor(10, TimeUnit.SECONDS)); + org.junit.jupiter.api.Assertions.assertTrue(second.waitFor(10, TimeUnit.SECONDS)); + assertEquals(1, (first.exitValue() == 0 ? 1 : 0) + (second.exitValue() == 0 ? 1 : 0)); + assertEquals(first.exitValue() == 0 ? "0" : "1", + new String(Files.readAllBytes(output.resolve("artifact")), StandardCharsets.UTF_8)); + } finally { + first.destroyForcibly(); + second.destroyForcibly(); + } + } + + public static final class PublisherProcess { + public static void main(String[] args) throws Exception { + Path output = Paths.get(args[0]); + byte[] payload = args[1].getBytes(StandardCharsets.UTF_8); + Path staging = Files.createTempFile(output, "artifact", ".tmp"); + Files.write(staging, payload); + Files.createFile(output.resolve("ready-" + args[1])); + try { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (!Files.exists(output.resolve("start"))) { + if (System.nanoTime() > deadline) { + throw new TimeoutException("publisher process start timed out"); + } + Thread.sleep(10); + } + ArtifactPublisher.publish(staging, output.resolve("artifact"), existing -> { + if (!Arrays.equals(payload, Files.readAllBytes(existing))) { + throw new IOException("artifact conflict"); + } + }); + } finally { + Files.deleteIfExists(staging); + } + } + } + + private int runRace(String kind, Path output, boolean conflicting) throws Exception { + ExecutorService executor = Executors.newFixedThreadPool(2); + CyclicBarrier start = new CyclicBarrier(2); + try { + Future first = executor.submit(() -> publish(kind, output, 0, start)); + Future second = executor.submit(() -> publish(kind, output, conflicting ? 1 : 0, start)); + boolean firstSuccess = first.get(10, TimeUnit.SECONDS); + boolean secondSuccess = second.get(10, TimeUnit.SECONDS); + assertEquals(conflicting ? 1 : 2, (firstSuccess ? 1 : 0) + (secondSuccess ? 1 : 0)); + return firstSuccess ? 0 : 1; + } finally { + executor.shutdownNow(); + } + } + + private boolean publish(String kind, Path output, int variant, CyclicBarrier start) throws Exception { + start.await(5, TimeUnit.SECONDS); + try { + build(kind, output, variant); + return true; + } catch (IOException | RetrievalException conflict) { + return false; + } + } + + private void build(String kind, Path output, int variant) throws Exception { + List chunks = Collections.singletonList(new TextChunk("chunk", "doc", 0, 0, 10, 3, + "graph evidence " + ("bm25".equals(kind) ? variant : 0)).withSourceUri("fixture://doc")); + if ("vector".equals(kind)) { + try (IndexArtifact ignored = new OfflineVectorIndexBuilder(output, + Collections.singletonMap("chunk", variant == 0 ? new float[] {1, 0} : new float[] {0, 1}), + "fixture", "v1").build(context(), chunks)) { + assertEquals("vector", ignored.getMetadata().getIndexType()); + } + } else if ("graph".equals(kind)) { + try (org.apache.geaflow.ai.retrieval.ingest.GraphArtifact ignored = + new OfflineGraphArtifactBuilder(output).build(context(), chunks, + Collections.singletonList(new EntityRef("entity", "name " + variant, + Collections.emptyList(), "topic", Collections.singletonList("chunk"))), + Collections.emptyList())) { + assertEquals("v1", ignored.getMetadata().getGraphVersion().getVersion()); + } + } else { + try (IndexArtifact ignored = new LuceneBm25IndexBuilder(output).build(context(), chunks)) { + assertEquals("bm25", ignored.getMetadata().getIndexType()); + } + } + } + + private void validateReadable(String kind, Path output) throws Exception { + if ("vector".equals(kind)) { + try (VectorIndexReader ignored = new VectorIndexReader(output.resolve("vector-v1.bin"))) { + assertEquals("v1", ignored.getIdentity().toMap().get("graphVersion")); + } + } else if ("graph".equals(kind)) { + try (GraphArtifactReader reader = new GraphArtifactReader(output.resolve("graph-v1.bin"), "v1", "v1")) { + assertEquals(1, reader.getEntities().size()); + } + } else { + try (Bm25IndexReader ignored = new Bm25IndexReader(output.resolve("bm25-v1"))) { + assertEquals("v1", ignored.getIdentity().toMap().get("graphVersion")); + } + } + } + + private void assertNoStaging(Path output) throws IOException { + try (java.util.stream.Stream paths = Files.list(output)) { + assertFalse(paths.anyMatch(path -> path.getFileName().toString().endsWith(".tmp") + || path.getFileName().toString().startsWith("bm25-v1-"))); + } + } + + private IngestionContext context() { + DatasetManifest manifest = new DatasetManifest("v1", "dataset", "release", "dev", null, + "cache", "0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef", + "parser", new ChunkingConfiguration("chunk-v1", 100, 10), "graph-v1", "fixture", "v1", 1L); + return new IngestionContext(manifest, new GraphVersion("graph", "v1"), "test"); + } +} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java index bdd4dfcdd..c603f51bf 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/ingest/IngestionPipelineTest.java @@ -22,16 +22,20 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import java.io.ByteArrayInputStream; import java.util.Arrays; import java.util.Collections; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; -import org.apache.geaflow.ai.retrieval.index.Bm25IndexBuilder; import org.apache.geaflow.ai.retrieval.index.IndexArtifact; -import org.apache.geaflow.ai.retrieval.index.VectorIndexBuilder; import org.apache.geaflow.ai.retrieval.metadata.ChunkingConfiguration; import org.apache.geaflow.ai.retrieval.metadata.DatasetManifest; import org.apache.geaflow.ai.retrieval.metadata.GraphBuildMetadata; @@ -91,12 +95,14 @@ public void runsStagesInOrderAndReportsCounters() throws Exception { }, new MetadataPublisher() { @Override - public void begin(IngestionContext value) { + public ImportAttempt begin(IngestionContext value) { calls.add("begin"); + return new ImportAttempt(value.getGraphVersion()); } @Override - public ImportMetadata publish(IngestionContext value, GraphArtifact valueGraph, + public ImportMetadata publish(IngestionContext value, ImportAttempt attempt, + GraphArtifact valueGraph, List indexes) { calls.add("publish"); assertEquals(2, indexes.size()); @@ -105,7 +111,7 @@ public ImportMetadata publish(IngestionContext value, GraphArtifact valueGraph, } @Override - public void fail(IngestionContext value, Exception failure) { + public void fail(IngestionContext value, ImportAttempt attempt, Exception failure) { throw new AssertionError("unexpected failure", failure); } }); @@ -135,18 +141,20 @@ public void failureClosesArtifactsAndNeverPublishes() throws Exception { }, new MetadataPublisher() { @Override - public void begin(IngestionContext value) { + public ImportAttempt begin(IngestionContext value) { + return new ImportAttempt(value.getGraphVersion()); } @Override - public ImportMetadata publish(IngestionContext value, GraphArtifact valueGraph, + public ImportMetadata publish(IngestionContext value, ImportAttempt attempt, + GraphArtifact valueGraph, List indexes) { published.set(true); return null; } @Override - public void fail(IngestionContext value, Exception failure) { + public void fail(IngestionContext value, ImportAttempt attempt, Exception failure) { failed.set(true); assertTrue(failure.getMessage().contains("vector build failed")); } @@ -164,6 +172,40 @@ public void fail(IngestionContext value, Exception failure) { assertTrue(bm25.closed); } + @Test + public void beginFailureDoesNotFailAnAttemptItDoesNotOwn() { + AtomicBoolean failed = new AtomicBoolean(); + IngestionPipeline pipeline = new IngestionPipeline( + value -> Collections.singletonList(document()), + (documents, value) -> Collections.singletonList(chunk()), + (chunks, value) -> new ExtractionResult(Collections.emptyList(), Collections.emptyList()), + (value, documents, chunks, extraction) -> new TrackingGraphArtifact( + new java.util.ArrayList<>(), "graph"), + (value, chunks) -> new TrackingIndexArtifact(new java.util.ArrayList<>(), "bm25"), + (value, chunks) -> new TrackingIndexArtifact(new java.util.ArrayList<>(), "vector"), + new MetadataPublisher() { + @Override + public ImportAttempt begin(IngestionContext value) { + throw new IllegalStateException("begin failed"); + } + + @Override + public ImportMetadata publish(IngestionContext value, ImportAttempt attempt, + GraphArtifact graph, List indexes) { + throw new AssertionError("unexpected publish"); + } + + @Override + public void fail(IngestionContext value, ImportAttempt attempt, Exception failure) { + failed.set(true); + } + }); + + org.junit.jupiter.api.Assertions.assertThrows(IllegalStateException.class, + () -> pipeline.run(context())); + assertFalse(failed.get()); + } + @Test public void metadataStorePublisherCompletesReadyLifecycle() throws Exception { String source = "source"; @@ -198,6 +240,117 @@ public void metadataStorePublisherCompletesReadyLifecycle() throws Exception { assertEquals(2, result.getIndexes().size()); } + @Test + public void rejectedConcurrentDuplicateLeavesOwnerAbleToPublishReady() throws Exception { + String source = "source"; + String checksum = "41cf6794ba4200b839c53531555f0f3998df4cbb01a4d5cb0b94e3ca5e23947d"; + GraphVersion version = new GraphVersion("graph", "g1"); + DatasetManifest manifest = new DatasetManifest("v1", "dataset", "release", "dev", "uri", null, + checksum, "preprocess-v1", new ChunkingConfiguration("chunk-v1", 100, 10), + "schema-v1", "model", "v1", 7); + IngestionContext ownerContext = new IngestionContext(manifest, version, "importer-v1"); + IngestionContext duplicateContext = new IngestionContext(manifest, version, "importer-v1"); + InMemoryMetadataStore store = new InMemoryMetadataStore(); + MetadataStorePublisher publisher = new MetadataStorePublisher(store, + value -> new ByteArrayInputStream(source.getBytes(java.nio.charset.StandardCharsets.UTF_8))); + CountDownLatch ownerStarted = new CountDownLatch(1); + CountDownLatch ownerReleased = new CountDownLatch(1); + TrackingGraphArtifact graph = new TrackingGraphArtifact(new java.util.ArrayList<>(), "graph", + version, Arrays.asList("bm25", "vector")); + TrackingIndexArtifact bm25 = new TrackingIndexArtifact(new java.util.ArrayList<>(), "bm25", version); + TrackingIndexArtifact vector = new TrackingIndexArtifact(new java.util.ArrayList<>(), "vector", version); + IngestionPipeline pipeline = new IngestionPipeline(value -> { + ownerStarted.countDown(); + try { + if (!ownerReleased.await(5, TimeUnit.SECONDS)) { + throw new java.io.IOException("owner release timed out"); + } + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new java.io.IOException(error); + } + return Collections.singletonList(document()); + }, (documents, value) -> Collections.singletonList(chunk()), + (chunks, value) -> new ExtractionResult(Collections.singletonList(entity()), + Collections.singletonList(edge())), + (value, documents, chunks, extraction) -> graph, + (value, chunks) -> bm25, (value, chunks) -> vector, publisher); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future owner = executor.submit(() -> pipeline.run(ownerContext)); + assertTrue(ownerStarted.await(5, TimeUnit.SECONDS)); + MetadataException rejected = assertThrows(MetadataException.class, + () -> pipeline.run(duplicateContext)); + assertEquals(MetadataException.Code.VERSION_CONFLICT, rejected.getCode()); + assertEquals(ImportState.IMPORTING, store.find(version).get().getState()); + assertFalse(store.getPublishedVersion(version.getGraphName()).isPresent()); + + ownerReleased.countDown(); + ImportMetadata result = owner.get(5, TimeUnit.SECONDS); + assertNotNull(result); + assertEquals(ImportState.READY, result.getState()); + assertEquals(ImportState.READY, store.find(version).get().getState()); + assertEquals(version, store.getPublishedVersion(version.getGraphName()).get()); + assertEquals(Arrays.asList("bm25", "vector"), result.getGraph().getRequiredIndexes()); + assertEquals(2, result.getIndexes().size()); + assertTrue(graph.closed); + assertTrue(bm25.closed); + assertTrue(vector.closed); + } finally { + ownerReleased.countDown(); + executor.shutdownNow(); + } + } + + @Test + public void concurrentDuplicateDoesNotFailOwnerAndForeignTokenCannotPublish() throws Exception { + IngestionContext context = context(); + InMemoryMetadataStore store = new InMemoryMetadataStore(); + MetadataStorePublisher publisher = new MetadataStorePublisher(store, + ignored -> new ByteArrayInputStream(new byte[0])); + CountDownLatch ownerStarted = new CountDownLatch(1); + CountDownLatch ownerReleased = new CountDownLatch(1); + AtomicBoolean ownerFailed = new AtomicBoolean(); + IngestionPipeline pipeline = new IngestionPipeline(value -> { + ownerStarted.countDown(); + try { + if (!ownerReleased.await(5, TimeUnit.SECONDS)) { + throw new java.io.IOException("owner release timed out"); + } + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw new java.io.IOException(error); + } + throw new java.io.IOException("owner build failure"); + }, (documents, value) -> Collections.singletonList(chunk()), + (chunks, value) -> new ExtractionResult(Collections.emptyList(), Collections.emptyList()), + (value, documents, chunks, extraction) -> new TrackingGraphArtifact(new java.util.ArrayList<>(), "graph"), + (value, chunks) -> new TrackingIndexArtifact(new java.util.ArrayList<>(), "bm25"), + (value, chunks) -> new TrackingIndexArtifact(new java.util.ArrayList<>(), "vector"), publisher); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try { + Future owner = executor.submit(() -> { + assertThrows(java.io.IOException.class, () -> pipeline.run(context)); + ownerFailed.set(true); + }); + assertTrue(ownerStarted.await(5, TimeUnit.SECONDS)); + assertThrows(MetadataException.class, () -> pipeline.run(context)); + assertEquals(ImportState.IMPORTING, store.find(context.getGraphVersion()).get().getState()); + ImportAttempt foreign = new ImportAttempt(context.getGraphVersion()); + publisher.fail(context, foreign, new java.io.IOException("foreign failure")); + assertThrows(MetadataException.class, () -> publisher.publish(context, foreign, + new TrackingGraphArtifact(new java.util.ArrayList<>(), "graph"), Collections.emptyList())); + assertEquals(ImportState.IMPORTING, store.find(context.getGraphVersion()).get().getState()); + ownerReleased.countDown(); + owner.get(5, TimeUnit.SECONDS); + assertTrue(ownerFailed.get()); + assertEquals(ImportState.FAILED, store.find(context.getGraphVersion()).get().getState()); + } finally { + ownerReleased.countDown(); + executor.shutdownNow(); + } + } + private static IngestionContext context() { DatasetManifest manifest = new DatasetManifest("v1", "hotpotqa", "release-1", "train", null, "cache/hotpotqa", SHA256, "normalizer-1", From 42189141ea8d820709d7558128c2efe69b2572cc Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 10 Oct 2026 20:31:53 +0800 Subject: [PATCH 26/30] fix(ai): secure and cache memory retrieval service --- .../geaflow/ai/GeaFlowMemoryServer.java | 61 ++++- .../apache/geaflow/ai/GraphMemoryServer.java | 208 ++++++++++++++---- .../geaflow/ai/operator/GraphSearchStore.java | 17 +- .../geaflow/ai/operator/SearchStore.java | 83 +++++-- .../api/model/RetrievalErrorCode.java | 2 + .../retrieval/config/RetrievalProperties.java | 50 +++++ .../geaflow/ai/service/ServerMemoryCache.java | 18 +- geaflow-ai/src/main/resources/application.yml | 4 +- .../ai/GraphMemoryKeywordSearchTest.java | 47 +++- .../apache/geaflow/ai/MemoryServerTest.java | 35 +++ .../geaflow/ai/operator/SearchStoreTest.java | 21 ++ 11 files changed, 486 insertions(+), 60 deletions(-) diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java index d9967240b..963d91136 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GeaFlowMemoryServer.java @@ -19,6 +19,8 @@ package org.apache.geaflow.ai; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; import java.util.ArrayList; import java.util.HashMap; import java.util.List; @@ -79,6 +81,19 @@ public static void main(String[] args) { app.cfg().loadAdd("application.yml"); app.cfg().put("server.port", port); LOGGER.info("Starting {} on port {}", SERVER_NAME, port); + app.filter(-100, (ctx, chain) -> { + String path = ctx.path(); + if (path.startsWith("/graph/") || path.startsWith("/query/")) { + String error = authenticationError(ctx, app.context().getBean(RetrievalProperties.class), + resolveRequestId(ctx.header("X-Request-Id"))); + if (error != null) { + ctx.output(error); + ctx.setHandled(true); + return; + } + } + chain.doFilter(ctx); + }); app.get("/", ctx -> { ctx.output("GeaFlow AI Server is running..."); }); @@ -111,9 +126,14 @@ public String retrieve(org.noear.solon.core.handle.Context ctx, @Body String inp long startedAt = System.nanoTime(); ctx.headerSet("X-Request-Id", requestId); ctx.contentType("application/json; charset=utf-8"); + RetrievalProperties properties = properties(); + String authenticationFailure = authenticationError(ctx, properties, requestId); + if (authenticationFailure != null) { + return authenticationFailure; + } try { - String body = RetrievalApiJson.toJson( - runtimeRetrievalService().retrieve(RetrievalApiJson.parseRequest(input), requestId)); + String body = RetrievalApiJson.toJson(runtimeRetrievalService().retrieve( + RetrievalApiJson.parseRequest(input), requestId), properties); ctx.status(200); return body; } catch (RetrievalException exception) { @@ -142,6 +162,43 @@ public String retrieve(org.noear.solon.core.handle.Context ctx, @Body String inp } } + private static String authenticationError(org.noear.solon.core.handle.Context ctx, + RetrievalProperties properties, String requestId) { + if (!properties.isRemoteAccessEnabled()) { + return null; + } + String authorization = ctx.header("Authorization"); + org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode code; + String message; + if (authorization == null || authorization.trim().isEmpty()) { + code = org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.AUTHENTICATION_REQUIRED; + message = "Bearer token required"; + } else { + String token = bearerToken(authorization); + if (token != null && properties.getApiToken() != null + && MessageDigest.isEqual(token.getBytes(StandardCharsets.UTF_8), + properties.getApiToken().getBytes(StandardCharsets.UTF_8))) { + return null; + } + code = org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.FORBIDDEN; + message = "Invalid bearer token"; + } + ctx.status(code.getHttpStatus()); + ctx.headerSet("X-Request-Id", requestId); + ctx.contentType("application/json; charset=utf-8"); + return RetrievalApiJson.toJson(new RetrievalError(requestId, code, message)); + } + + private static String bearerToken(String authorization) { + String value = authorization.trim(); + int separator = value.indexOf(' '); + if (separator < 1 || !"Bearer".equalsIgnoreCase(value.substring(0, separator))) { + return null; + } + String token = value.substring(separator + 1).trim(); + return token.isEmpty() ? null : token; + } + static String resolveRequestId(String requestedId) { if (requestedId == null) { return java.util.UUID.randomUUID().toString(); diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java index 02d29087f..36cbb607f 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/GraphMemoryServer.java @@ -23,6 +23,10 @@ import java.util.HashSet; import java.util.List; import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReadWriteLock; +import java.util.concurrent.locks.ReentrantReadWriteLock; import java.util.stream.Collectors; import org.apache.geaflow.ai.graph.GraphAccessor; import org.apache.geaflow.ai.graph.GraphEntity; @@ -41,15 +45,27 @@ import org.apache.geaflow.ai.verbalization.Context; import org.apache.geaflow.ai.verbalization.VerbalizationFunction; -public class GraphMemoryServer { +public class GraphMemoryServer implements AutoCloseable { private final SessionManagement sessionManagement = new SessionManagement(); private final List graphAccessors = new ArrayList<>(); private final List indexStores = new ArrayList<>(); + private final ReadWriteLock keywordStoreLock = new ReentrantReadWriteLock(); + private GraphSearchStore keywordSearchStore; + private GraphAccessor keywordSearchAccessor; + private EntityAttributeIndexStore keywordSearchIndex; + private int keywordSearchStoreBuildCount; + private boolean closed; public void addGraphAccessor(GraphAccessor graph) { if (graph != null) { - graphAccessors.add(graph); + keywordStoreLock.writeLock().lock(); + try { + closeKeywordSearchStore(); + graphAccessors.add(graph); + } finally { + keywordStoreLock.writeLock().unlock(); + } } } @@ -59,7 +75,13 @@ public List getGraphAccessors() { public void addIndexStore(IndexStore indexStore) { if (indexStore != null) { - indexStores.add(indexStore); + keywordStoreLock.writeLock().lock(); + try { + closeKeywordSearchStore(); + indexStores.add(indexStore); + } finally { + keywordStoreLock.writeLock().unlock(); + } } } @@ -73,47 +95,153 @@ public List searchKeyword(String query, int topK, int maxCand if (graphAccessors.isEmpty()) { throw new IllegalStateException("No graph accessor available"); } - for (IndexStore indexStore : indexStores) { - if (indexStore instanceof EntityAttributeIndexStore) { - EntityAttributeIndexStore keywordIndex = (EntityAttributeIndexStore) indexStore; - if (!keywordIndex.isInitialized()) { - throw new IllegalStateException("Keyword index is not initialized"); + EntityAttributeIndexStore keywordIndex = findKeywordIndex(); + if (keywordIndex == null) { + throw new IllegalStateException("Keyword index is not available"); + } + GraphAccessor accessor = graphAccessors.get(0); + GraphSearchStore store = acquireKeywordSearchStore(accessor, keywordIndex, deadlineNanos); + try { + if (System.nanoTime() > deadlineNanos) { + throw retrievalTimeout(); + } + return store.searchScored(query, accessor, topK, maxCandidates, deadlineNanos); + } finally { + keywordStoreLock.readLock().unlock(); + } + } + + /** Invalidates the cached keyword index after graph data or schemas change. */ + public void invalidateKeywordSearchIndex() { + keywordStoreLock.writeLock().lock(); + try { + closeKeywordSearchStore(); + } finally { + keywordStoreLock.writeLock().unlock(); + } + } + + @Override + public void close() { + keywordStoreLock.writeLock().lock(); + try { + closed = true; + closeKeywordSearchStore(); + } finally { + keywordStoreLock.writeLock().unlock(); + } + } + + public int getKeywordSearchStoreBuildCount() { + keywordStoreLock.readLock().lock(); + try { + return keywordSearchStoreBuildCount; + } finally { + keywordStoreLock.readLock().unlock(); + } + } + + private GraphSearchStore acquireKeywordSearchStore(GraphAccessor accessor, + EntityAttributeIndexStore keywordIndex, + long deadlineNanos) { + lockUntilDeadline(keywordStoreLock.readLock(), deadlineNanos); + if (closed) { + keywordStoreLock.readLock().unlock(); + throw new IllegalStateException("graph memory server is closed"); + } + if (keywordSearchStore != null && keywordSearchAccessor == accessor + && keywordSearchIndex == keywordIndex) { + return keywordSearchStore; + } + keywordStoreLock.readLock().unlock(); + lockUntilDeadline(keywordStoreLock.writeLock(), deadlineNanos); + try { + if (closed) { + throw new IllegalStateException("graph memory server is closed"); + } + if (keywordSearchStore == null || keywordSearchAccessor != accessor + || keywordSearchIndex != keywordIndex) { + closeKeywordSearchStore(); + keywordSearchStore = buildKeywordSearchStore(accessor, keywordIndex, deadlineNanos); + keywordSearchAccessor = accessor; + keywordSearchIndex = keywordIndex; + keywordSearchStoreBuildCount++; + } + keywordStoreLock.readLock().lock(); + return keywordSearchStore; + } finally { + keywordStoreLock.writeLock().unlock(); + } + } + + private static void lockUntilDeadline(Lock lock, long deadlineNanos) { + long remaining = deadlineNanos - System.nanoTime(); + try { + if (remaining <= 0 || !lock.tryLock(remaining, TimeUnit.NANOSECONDS)) { + throw retrievalTimeout(); + } + } catch (InterruptedException error) { + Thread.currentThread().interrupt(); + throw retrievalTimeout(); + } + } + + private static GraphSearchStore buildKeywordSearchStore(GraphAccessor accessor, + EntityAttributeIndexStore keywordIndex, + long deadlineNanos) { + if (!keywordIndex.isInitialized()) { + throw new IllegalStateException("Keyword index is not initialized"); + } + GraphSearchStore store = new GraphSearchStore(); + boolean complete = false; + try { + java.util.Iterator vertices = accessor.scanVertex(); + while (vertices.hasNext()) { + if (System.nanoTime() > deadlineNanos) { + throw retrievalTimeout(); } - GraphAccessor accessor = graphAccessors.get(0); - GraphSearchStore store = new GraphSearchStore(); - try { - java.util.Iterator vertices = accessor.scanVertex(); - int candidates = 0; - while (vertices.hasNext()) { - if (candidates >= maxCandidates || System.nanoTime() > deadlineNanos) { - if (System.nanoTime() > deadlineNanos) { - throw new org.apache.geaflow.ai.retrieval.api.model.RetrievalException( - org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.RETRIEVAL_TIMEOUT, - "retrieval deadline exceeded"); - } - break; - } - candidates++; - GraphVertex vertex = vertices.next(); - List vectors = - keywordIndex.getEntityIndex(vertex); - if (vectors != null && !vectors.isEmpty()) { - store.indexVertex(vertex, vectors); - } - } - if (System.nanoTime() > deadlineNanos) { - throw new org.apache.geaflow.ai.retrieval.api.model.RetrievalException( - org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.RETRIEVAL_TIMEOUT, - "retrieval deadline exceeded"); - } - store.finishWriting(); - return store.searchScored(query, accessor, topK, maxCandidates, deadlineNanos); - } finally { - store.close(); + GraphVertex vertex = vertices.next(); + List vectors = keywordIndex.getEntityIndex(vertex); + if (vectors != null && !vectors.isEmpty()) { + store.indexVertex(vertex, vectors); } } + if (System.nanoTime() > deadlineNanos) { + throw retrievalTimeout(); + } + store.finishWriting(); + complete = true; + return store; + } finally { + if (!complete) { + store.close(); + } } - throw new IllegalStateException("Keyword index is not available"); + } + + private EntityAttributeIndexStore findKeywordIndex() { + for (IndexStore indexStore : indexStores) { + if (indexStore instanceof EntityAttributeIndexStore) { + return (EntityAttributeIndexStore) indexStore; + } + } + return null; + } + + private void closeKeywordSearchStore() { + if (keywordSearchStore != null) { + final GraphSearchStore previous = keywordSearchStore; + keywordSearchStore = null; + keywordSearchAccessor = null; + keywordSearchIndex = null; + previous.close(); + } + } + + private static org.apache.geaflow.ai.retrieval.api.model.RetrievalException retrievalTimeout() { + return new org.apache.geaflow.ai.retrieval.api.model.RetrievalException( + org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.RETRIEVAL_TIMEOUT, + "retrieval deadline exceeded"); } public String createSession() { diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java index 43a4da513..3b20e9f02 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/GraphSearchStore.java @@ -72,12 +72,14 @@ public int getRank() { private SearchStore store; private long entityNum = 0L; + private boolean closed; public GraphSearchStore() { this.store = new SearchStore(); } public boolean indexVertex(GraphVertex graphVertex, List indexVectors) { + ensureOpen(); Map kv = new HashMap<>(); Vertex vertex = graphVertex.getVertex(); kv.put(SearchConstants.ID, vertex.getId()); @@ -99,6 +101,7 @@ public boolean indexVertex(GraphVertex graphVertex, List indexVectors) } public boolean indexEdge(GraphEdge graphEdge, List indexVectors) { + ensureOpen(); Map kv = new HashMap<>(); Edge edge = graphEdge.getEdge(); kv.put(SearchConstants.SRC, edge.getSrcId()); @@ -133,6 +136,7 @@ public List search(String key1, GraphAccessor graphAccessor) { public List searchScored(String key1, GraphAccessor graphAccessor, int topK, int maxCandidates, long deadlineNanos) { + ensureOpen(); try { String query = SearchUtils.formatQuery(key1); TopDocs docs = store.searchDoc(SearchConstants.CONTENT, query, @@ -208,7 +212,11 @@ private void addItem() { entityNum++; } - public void close() { + public synchronized void close() { + if (closed) { + return; + } + closed = true; try { store.close(); } catch (Throwable e) { @@ -217,6 +225,7 @@ public void close() { } public void finishWriting() { + ensureOpen(); try { store.finishWriting(); } catch (Throwable e) { @@ -236,5 +245,11 @@ public IndexWriterConfig getConfig() { return store.getConfig(); } + private void ensureOpen() { + if (closed) { + throw new IllegalStateException("graph search store is closed"); + } + } + } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java index 6749246d1..589d7edfb 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/operator/SearchStore.java @@ -48,11 +48,13 @@ public class SearchStore { private IndexReader reader; private IndexSearcher searcher; private boolean readStats = false; + private boolean closed; public SearchStore() { } - public void addDoc(Map kv) throws IOException { + public synchronized void addDoc(Map kv) throws IOException { + ensureOpen(); initWriter(); Document doc = new Document(); for (Map.Entry entry : kv.entrySet()) { @@ -65,7 +67,9 @@ public TopDocs searchDoc(String field, String content) throws ParseException, IO return searchDoc(field, content, Constants.GRAPH_SEARCH_STORE_DEFAULT_TOPN); } - public TopDocs searchDoc(String field, String content, int topN) throws ParseException, IOException { + public synchronized TopDocs searchDoc(String field, String content, int topN) + throws ParseException, IOException { + ensureOpen(); if (topN < 1) { throw new IllegalArgumentException("topN must be positive"); } @@ -78,7 +82,8 @@ public TopDocs searchDoc(String field, String content, int topN) throws ParseExc return searcher.search(parser.parse(content), topN); } - public Document getDoc(int docId) { + public synchronized Document getDoc(int docId) { + ensureOpen(); try { if (!readStats) { reader = DirectoryReader.open(directory); @@ -92,28 +97,80 @@ public Document getDoc(int docId) { } - public void initWriter() throws IOException { + public synchronized void initWriter() throws IOException { + ensureOpen(); if (!writeStats) { writer = new IndexWriter(directory, config); writeStats = true; } } - public void close() throws IOException { - finishWriting(); + public synchronized void close() throws IOException { + if (closed) { + return; + } + closed = true; + IOException failure = null; + if (writeStats) { + try { + writer.close(); + } catch (IOException error) { + failure = error; + } finally { + writer = null; + writeStats = false; + } + } if (readStats) { - reader.close(); - readStats = false; + try { + reader.close(); + } catch (IOException error) { + failure = append(failure, error); + } finally { + reader = null; + searcher = null; + readStats = false; + } + } + try { + directory.close(); + } catch (IOException error) { + failure = append(failure, error); + } + try { + analyzer.close(); + } catch (RuntimeException error) { + IOException wrapped = new IOException("failed to close analyzer", error); + failure = append(failure, wrapped); + } + if (failure != null) { + throw failure; } - directory.close(); - analyzer.close(); } - public void finishWriting() throws IOException { + public synchronized void finishWriting() throws IOException { if (writeStats) { - writer.close(); - writeStats = false; + try { + writer.close(); + } finally { + writer = null; + writeStats = false; + } + } + } + + private void ensureOpen() { + if (closed) { + throw new IllegalStateException("search store is closed"); + } + } + + private static IOException append(IOException current, IOException next) { + if (current == null) { + return next; } + current.addSuppressed(next); + return current; } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalErrorCode.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalErrorCode.java index 7a52f43f7..71e09ceeb 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalErrorCode.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/api/model/RetrievalErrorCode.java @@ -22,6 +22,8 @@ public enum RetrievalErrorCode { INVALID_REQUEST(400, false), UNSUPPORTED_OPTION(400, false), + AUTHENTICATION_REQUIRED(401, false), + FORBIDDEN(403, false), GRAPH_NOT_FOUND(404, false), INDEX_NOT_READY(503, true), RETRIEVAL_TIMEOUT(504, true), diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java index f00b38079..b43455f98 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/config/RetrievalProperties.java @@ -18,6 +18,8 @@ package org.apache.geaflow.ai.retrieval.config; +import java.net.InetAddress; +import java.net.UnknownHostException; import org.apache.geaflow.ai.retrieval.api.model.ExecutionMode; import org.apache.geaflow.ai.retrieval.api.model.RetrievalException; import org.apache.geaflow.ai.retrieval.api.model.RetrievalMode; @@ -53,6 +55,12 @@ public class RetrievalProperties { private int defaultTokenBudget = 4096; @Inject("${retrieval.max-token-budget:16384}") private int maxTokenBudget = 16384; + @Inject("${server.host:127.0.0.1}") + private String serverHost = "127.0.0.1"; + @Inject("${retrieval.remote-access-enabled:false}") + private boolean remoteAccessEnabled; + @Inject("${retrieval.api-token:}") + private String apiToken; @Init public void validateConfiguration() { @@ -78,6 +86,13 @@ public void validateConfiguration() { if (defaultTopK > defaultMaxCandidates) { throw invalid("default-top-k must not exceed default-max-candidates"); } + boolean loopback = isLoopbackHost(serverHost); + if (!loopback && !remoteAccessEnabled) { + throw invalid("non-loopback server.host requires retrieval.remote-access-enabled=true"); + } + if (remoteAccessEnabled && blank(apiToken)) { + throw invalid("retrieval.api-token is required when remote access is enabled"); + } } private static void positiveLimit(int defaultValue, int maxValue, String name) { @@ -90,6 +105,17 @@ private static boolean blank(String value) { return value == null || value.trim().isEmpty(); } + private static boolean isLoopbackHost(String host) { + if (blank(host)) { + return false; + } + try { + return InetAddress.getByName(host.trim()).isLoopbackAddress(); + } catch (UnknownHostException error) { + throw invalid("server.host is not a resolvable address"); + } + } + private static RetrievalException invalid(String message) { return new RetrievalException( org.apache.geaflow.ai.retrieval.api.model.RetrievalErrorCode.INVALID_REQUEST, message); @@ -190,4 +216,28 @@ public int getMaxTokenBudget() { public void setMaxTokenBudget(int maxTokenBudget) { this.maxTokenBudget = maxTokenBudget; } + + public String getServerHost() { + return serverHost; + } + + public void setServerHost(String serverHost) { + this.serverHost = serverHost; + } + + public boolean isRemoteAccessEnabled() { + return remoteAccessEnabled; + } + + public void setRemoteAccessEnabled(boolean remoteAccessEnabled) { + this.remoteAccessEnabled = remoteAccessEnabled; + } + + public String getApiToken() { + return apiToken; + } + + public void setApiToken(String apiToken) { + this.apiToken = apiToken; + } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/service/ServerMemoryCache.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/service/ServerMemoryCache.java index bedd7a676..7e47152c0 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/service/ServerMemoryCache.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/service/ServerMemoryCache.java @@ -26,6 +26,7 @@ import org.apache.geaflow.ai.consolidate.ConsolidateServer; import org.apache.geaflow.ai.graph.Graph; import org.noear.solon.annotation.Component; +import org.noear.solon.annotation.Destroy; @Component public class ServerMemoryCache { @@ -47,7 +48,10 @@ public void putServer(GraphMemoryServer server) { throw new RuntimeException("Cannot register server without graph accessor"); } String name = server.getGraphAccessors().get(0).getGraphSchema().getName(); - name2Server.put(name, server); + GraphMemoryServer previous = name2Server.put(name, server); + if (previous != null && previous != server) { + previous.close(); + } graphVersions.putIfAbsent(name, new AtomicLong(1L)); } @@ -75,12 +79,24 @@ public ConsolidateServer getConsolidateServer() { return consolidateServer; } + @Destroy + public void close() { + for (GraphMemoryServer server : name2Server.values()) { + server.close(); + } + name2Server.clear(); + } + public String getGraphVersion(String graphName) { AtomicLong version = graphVersions.get(graphName); return version == null ? null : "v" + version.get(); } public String markGraphUpdated(String graphName) { + GraphMemoryServer server = name2Server.get(graphName); + if (server != null) { + server.invalidateKeywordSearchIndex(); + } AtomicLong version = graphVersions.get(graphName); if (version == null) { throw new IllegalArgumentException("Unknown graph: " + graphName); diff --git a/geaflow-ai/src/main/resources/application.yml b/geaflow-ai/src/main/resources/application.yml index b808dcc56..01f7aeb49 100644 --- a/geaflow-ai/src/main/resources/application.yml +++ b/geaflow-ai/src/main/resources/application.yml @@ -22,7 +22,7 @@ solon: server: port: ${GEAFLOW_SERVER_PORT:8080} - host: 0.0.0.0 + host: ${GEAFLOW_SERVER_HOST:127.0.0.1} logging: level: @@ -42,3 +42,5 @@ retrieval: max-candidates: 1000 default-token-budget: 4096 max-token-budget: 16384 + remote-access-enabled: ${GEAFLOW_REMOTE_ACCESS_ENABLED:false} + api-token: ${GEAFLOW_RETRIEVAL_API_TOKEN:} diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/GraphMemoryKeywordSearchTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/GraphMemoryKeywordSearchTest.java index 401840566..06006eded 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/GraphMemoryKeywordSearchTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/GraphMemoryKeywordSearchTest.java @@ -28,6 +28,10 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import org.apache.geaflow.ai.graph.LocalMemoryGraphAccessor; import org.apache.geaflow.ai.graph.io.EntityGroup; import org.apache.geaflow.ai.graph.io.GraphSchema; @@ -45,12 +49,13 @@ class GraphMemoryKeywordSearchTest { @Test - void maxCandidatesStopsScanningBeforeTailMatch() { + void candidateBudgetLimitsSearchAcrossCompleteCachedIndex() { GraphMemoryServer server = createServer(); List hits = server.searchKeyword("stars", 10, 1, Long.MAX_VALUE); - assertEquals(0, hits.size()); + assertEquals(1, hits.size()); + server.close(); } @Test @@ -63,6 +68,44 @@ void expiredDeadlineIsReportedAsTimeout() { assertEquals(RetrievalErrorCode.RETRIEVAL_TIMEOUT, exception.getCode()); } + @Test + void reusesCachedIndexUntilGraphUpdateInvalidatesIt() { + GraphMemoryServer server = createServer(); + + server.searchKeyword("confucius", 1, 10, Long.MAX_VALUE); + server.searchKeyword("stars", 1, 10, Long.MAX_VALUE); + assertEquals(1, server.getKeywordSearchStoreBuildCount()); + + server.invalidateKeywordSearchIndex(); + server.searchKeyword("stars", 1, 10, Long.MAX_VALUE); + + assertEquals(2, server.getKeywordSearchStoreBuildCount()); + server.invalidateKeywordSearchIndex(); + server.close(); + server.close(); + assertThrows(IllegalStateException.class, + () -> server.searchKeyword("stars", 1, 10, Long.MAX_VALUE)); + } + + @Test + void concurrentRequestsBuildOneCachedStore() throws Exception { + GraphMemoryServer server = createServer(); + ExecutorService executor = Executors.newFixedThreadPool(4); + try { + List>> calls = new ArrayList<>(); + for (int index = 0; index < 16; index++) { + calls.add(executor.submit(() -> server.searchKeyword("stars", 1, 10, Long.MAX_VALUE))); + } + for (Future> call : calls) { + assertEquals(1, call.get(5, TimeUnit.SECONDS).size()); + } + assertEquals(1, server.getKeywordSearchStoreBuildCount()); + } finally { + executor.shutdownNow(); + server.close(); + } + } + private static GraphMemoryServer createServer() { VertexSchema vertexSchema = new VertexSchema("chunk", "id", Collections.singletonList("text")); diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java index 4bdf14821..ef715adc7 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/MemoryServerTest.java @@ -34,15 +34,20 @@ import org.apache.geaflow.ai.graph.io.Vertex; import org.apache.geaflow.ai.retrieval.api.model.RetrievalResponse; import org.apache.geaflow.ai.retrieval.codec.RetrievalApiJson; +import org.apache.geaflow.ai.retrieval.config.RetrievalProperties; import org.apache.geaflow.ai.retrieval.support.HttpTestClient; import org.apache.geaflow.ai.retrieval.support.HttpTestResponse; import org.apache.geaflow.ai.retrieval.support.RetrievalTestFixture; import org.junit.jupiter.api.Test; +import org.noear.solon.annotation.Inject; import org.noear.solon.test.SolonTest; @SolonTest(GeaFlowMemoryServer.class) public class MemoryServerTest { + @Inject + private RetrievalProperties properties; + @Test void legacyWorkflow() { RetrievalTestFixture fixture = new RetrievalTestFixture(); @@ -103,6 +108,36 @@ void legacyWorkflow() { assertTrue(response.getEvidence().get(0).getFinalScore() > 0); assertTrue(response.getEvidence().get(0).getStageScores().containsKey("keyword")); + int previousMaxTopK = properties.getMaxTopK(); + try { + properties.setMaxTopK(200); + String largerRequest = request.replace("\"topK\":1", "\"topK\":120") + .replace("\"maxCandidates\":10", "\"maxCandidates\":120"); + HttpTestResponse larger = client.post("/api/v1/retrievals", largerRequest); + assertEquals(200, larger.getStatus(), larger.getBody()); + assertEquals(120, RetrievalApiJson.parseResponse(larger.getBody(), properties) + .getEffectiveBudget().getTopK()); + properties.setApiToken("http-test-token"); + properties.setRemoteAccessEnabled(true); + assertEquals(401, client.post("/api/v1/retrievals", request).getStatus()); + assertEquals(403, client.postWithHeaders("/api/v1/retrievals", request, + Collections.singletonMap("Authorization", "Bearer wrong")).getStatus()); + assertEquals(403, client.postWithHeaders("/api/v1/retrievals", request, + Collections.singletonMap("Authorization", "Basic http-test-token")).getStatus()); + HttpTestResponse authorized = client.postWithHeaders("/api/v1/retrievals", request, + Collections.singletonMap("Authorization", "Bearer http-test-token")); + assertEquals(200, authorized.getStatus(), authorized.getBody()); + assertEquals(401, client.post("/api/v1/retrievals", "{").getStatus()); + assertEquals(401, client.post("/graph/getGraphSchema", "", graph).getStatus()); + assertEquals(403, client.postWithHeaders("/query/result", "", + Collections.singletonMap("Authorization", "Bearer wrong")).getStatus()); + assertEquals(200, client.get("/health").getStatus()); + } finally { + properties.setMaxTopK(previousMaxTopK); + properties.setRemoteAccessEnabled(false); + properties.setApiToken(null); + } + for (String mode : new String[] {"BM25_ONLY", "VECTOR_ONLY", "GRAPH_ONLY", "HYBRID"}) { String unsupportedRequest = request.replace("\"query\":", "\"mode\":\"" + mode + "\",\"queryVector\":[1,0],\"query\":"); diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/SearchStoreTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/SearchStoreTest.java index 211d7aac6..57bc6a36b 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/SearchStoreTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/operator/SearchStoreTest.java @@ -20,6 +20,7 @@ package org.apache.geaflow.ai.operator; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import java.util.Collections; import org.apache.lucene.search.TopDocs; @@ -41,6 +42,26 @@ void searchDocHonorsRequestedTopN() throws Exception { assertEquals(1, docs.scoreDocs.length); } finally { store.close(); + store.close(); } + assertThrows(IllegalStateException.class, + () -> store.searchDoc(SearchConstants.CONTENT, "confucius", 1)); + } + + @Test + void emptyAndWriteOnlyStoresAndGraphWrapperCanBeClosedRepeatedly() throws Exception { + SearchStore empty = new SearchStore(); + empty.close(); + empty.close(); + SearchStore writeOnly = new SearchStore(); + writeOnly.addDoc(Collections.singletonMap(SearchConstants.CONTENT, "text")); + writeOnly.close(); + writeOnly.close(); + assertThrows(IllegalStateException.class, + () -> writeOnly.addDoc(Collections.singletonMap(SearchConstants.CONTENT, "text"))); + GraphSearchStore graph = new GraphSearchStore(); + graph.close(); + graph.close(); + assertThrows(IllegalStateException.class, graph::finishWriting); } } From 1dc145ad616065f0fd6e112dbd4a095721a60bc1 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 10 Oct 2026 20:33:04 +0800 Subject: [PATCH 27/30] fix(ai): preserve retrieval provenance and trace identity --- .../retrieval/codec/RetrievalModelJson.java | 3 +- .../retrieval/model/document/TextChunk.java | 30 ++++++++++++++++--- .../ai/retrieval/service/RecallService.java | 27 +++++++++-------- .../service/RetrievalFixtureRegistry.java | 4 +-- .../retrieval/api/RetrievalApiJsonTest.java | 2 ++ .../api/RetrievalPropertiesTest.java | 16 ++++++++++ .../index/graph/GraphArtifactTest.java | 4 +-- .../service/GraphRagMvpEndToEndTest.java | 27 +++++++++++++---- 8 files changed, 87 insertions(+), 26 deletions(-) diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java index 69e0ccc90..d9deefb1c 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/codec/RetrievalModelJson.java @@ -95,7 +95,8 @@ private static Object parse(JsonObject object, Class type) { requiredString(object, "documentId"), requiredInt(object, "chunkIndex"), requiredInt(object, "startOffset"), requiredInt(object, "endOffset"), requiredInt(object, "tokenEstimate"), requiredString(object, "text"), - optionalString(object, "policyVersion"), optionalString(object, "textHash")); + optionalString(object, "policyVersion"), optionalString(object, "textHash"), + optionalString(object, "sourceUri")); } else if (type == EntityRef.class) { return new EntityRef(requiredString(object, "entityId"), requiredString(object, "canonicalName"), strings(object, "aliases"), diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/document/TextChunk.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/document/TextChunk.java index b52a8d3b8..097456cb0 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/document/TextChunk.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/model/document/TextChunk.java @@ -45,15 +45,25 @@ public final class TextChunk { private final String policyVersion; @SerializedName("textHash") private final String textHash; + @SerializedName("sourceUri") + private final String sourceUri; public TextChunk(String chunkId, String documentId, int chunkIndex, int startOffset, int endOffset, int tokenEstimate, String text) { - this(chunkId, documentId, chunkIndex, startOffset, endOffset, tokenEstimate, text, null, null); + this(chunkId, documentId, chunkIndex, startOffset, endOffset, tokenEstimate, text, + null, null, null); } public TextChunk(String chunkId, String documentId, int chunkIndex, int startOffset, int endOffset, int tokenEstimate, String text, String policyVersion, String textHash) { + this(chunkId, documentId, chunkIndex, startOffset, endOffset, tokenEstimate, text, + policyVersion, textHash, null); + } + + public TextChunk(String chunkId, String documentId, int chunkIndex, int startOffset, + int endOffset, int tokenEstimate, String text, + String policyVersion, String textHash, String sourceUri) { this.chunkId = ModelValidation.required(chunkId, "chunkId"); this.documentId = ModelValidation.required(documentId, "documentId"); this.chunkIndex = ModelValidation.nonNegative(chunkIndex, "chunkIndex"); @@ -66,6 +76,7 @@ public TextChunk(String chunkId, String documentId, int chunkIndex, int startOff this.text = ModelValidation.required(text, "text"); this.policyVersion = ModelValidation.optionalNonBlank(policyVersion, "policyVersion"); this.textHash = ModelValidation.optionalNonBlank(textHash, "textHash"); + this.sourceUri = ModelValidation.optionalNonBlank(sourceUri, "sourceUri"); } public String getChunkId() { @@ -104,6 +115,15 @@ public String getTextHash() { return textHash; } + public String getSourceUri() { + return sourceUri; + } + + public TextChunk withSourceUri(String uri) { + return new TextChunk(chunkId, documentId, chunkIndex, startOffset, endOffset, + tokenEstimate, text, policyVersion, textHash, ModelValidation.required(uri, "sourceUri")); + } + public boolean sameIdentityAs(TextChunk other) { return other != null && Objects.equals(chunkId, other.chunkId) && Objects.equals(documentId, other.documentId) @@ -112,7 +132,8 @@ public boolean sameIdentityAs(TextChunk other) { && endOffset == other.endOffset && Objects.equals(text, other.text) && Objects.equals(policyVersion, other.policyVersion) - && Objects.equals(textHash, other.textHash); + && Objects.equals(textHash, other.textHash) + && Objects.equals(sourceUri, other.sourceUri); } @Override @@ -130,12 +151,13 @@ public boolean equals(Object object) { && Objects.equals(documentId, that.documentId) && Objects.equals(text, that.text) && Objects.equals(policyVersion, that.policyVersion) - && Objects.equals(textHash, that.textHash); + && Objects.equals(textHash, that.textHash) + && Objects.equals(sourceUri, that.sourceUri); } @Override public int hashCode() { return Objects.hash(chunkId, documentId, chunkIndex, startOffset, endOffset, - tokenEstimate, text, policyVersion, textHash); + tokenEstimate, text, policyVersion, textHash, sourceUri); } } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RecallService.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RecallService.java index 8f9b81e21..cd3d965d2 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RecallService.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RecallService.java @@ -161,8 +161,10 @@ public RetrievalResponse retrieve(RetrievalRequest request) { throw notReady("vector source/version mismatch"); } RecallService delegate = new RecallService(fixture, fusion.getConfig(), clock); + String requestId = request.getRequestId() == null || request.getRequestId().trim().isEmpty() + ? UUID.randomUUID().toString() : request.getRequestId(); RetrievalResponse response = delegate.execute(command, toFloat(command.getQueryVector()), deadline, - request.isAllowPartialResults(), started, fixture); + request.isAllowPartialResults(), started, fixture, requestId); response.setGraphName(fixture.getGraphName()); response.setGraphVersion(fixture.getGraphVersion()); response.getTrace().setGraphVersion(fixture.getGraphVersion()); @@ -170,10 +172,6 @@ public RetrievalResponse retrieve(RetrievalRequest request) { if (fixture.getVector() != null) { response.getTrace().setVectorVersion(fixture.getVector().getVectorVersion()); } - if (request.getRequestId() != null) { - response.setRequestId(request.getRequestId()); - response.getTrace().setRequestId(request.getRequestId()); - } response.getTrace().setElapsedNanos(clock.getAsLong() - started); return response; } @@ -217,11 +215,13 @@ public RetrievalResponse retrieveWithTrace(String query, float[] queryVector, St if (registry != null) { throw notReady("registry retrieval requires a structured request with graph/index versions"); } - return execute(command, queryVector, deadlineNanos, allowPartialResults, started, null); + return execute(command, queryVector, deadlineNanos, allowPartialResults, started, null, + UUID.randomUUID().toString()); } private RetrievalResponse execute(RetrievalCommand command, float[] queryVector, long deadline, - boolean partial, long started, RetrievalFixtureRegistry.Fixture fixture) { + boolean partial, long started, RetrievalFixtureRegistry.Fixture fixture, + String requestId) { RetrievalBudget budget = command.getBudget(); RecallPlan plan = new RecallPlan(command.getMode(), budget.getTopK(), budget.getMaxCandidates(), deadline, fixture == null ? null : fixture.getGraphVersion(), fixture == null ? null : fixture.getIndexVersion(), @@ -321,7 +321,7 @@ private RetrievalResponse execute(RetrievalCommand command, float[] queryVector, if (!builder.paths.contains(hit.getPath())) { builder.paths.add(hit.getPath()); } - builder.graphTraces.add(new GraphEvidenceTrace(command.getQuery(), + builder.graphTraces.add(new GraphEvidenceTrace(requestId, plan.getGraphVersion() == null ? "fixture" : plan.getGraphVersion(), hit.getAnchor().getEntity().getEntityId(), hit.getAnchor().getMatchType(), hit.getAnchor().getConfidence(), hit.getPath(), @@ -416,9 +416,8 @@ private RetrievalResponse execute(RetrievalCommand command, float[] queryVector, } response.setSources(sources); response.setPaths(paths); - String responseId = UUID.randomUUID().toString(); - response.setRequestId(responseId); - trace.setRequestId(responseId); + response.setRequestId(requestId); + trace.setRequestId(requestId); trace.setElapsedNanos(clock.getAsLong() - started); response.setTrace(trace); response.setEffectiveBudget(budget); @@ -488,7 +487,11 @@ private Evidence build() { for (ChannelScore score : scores.values()) { fused += 1.0 / (60.0 + score.getRank()); } - SourceRef source = new SourceRef(chunk.getDocumentId(), chunk.getDocumentId(), chunk.getStartOffset(), chunk.getEndOffset()); + if (chunk.getSourceUri() == null) { + throw notReady("chunk source URI is missing"); + } + SourceRef source = new SourceRef(chunk.getDocumentId(), chunk.getSourceUri(), + chunk.getStartOffset(), chunk.getEndOffset()); return new Evidence(id, kind, chunk.getText(), Collections.singletonList(chunk), entities, paths, Collections.singletonList(source), scores, fused, null, 1, graphTraces); } diff --git a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalFixtureRegistry.java b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalFixtureRegistry.java index 95811706d..4fef21111 100644 --- a/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalFixtureRegistry.java +++ b/geaflow-ai/src/main/java/org/apache/geaflow/ai/retrieval/service/RetrievalFixtureRegistry.java @@ -64,8 +64,8 @@ private synchronized void register(String graphName, String graphVersion, String } java.util.Set chunkIds = new java.util.HashSet<>(); for (TextChunk chunk : chunks) { - if (chunk == null || !chunkIds.add(chunk.getChunkId())) { - throw notReady("invalid or duplicate fixture chunk"); + if (chunk == null || !chunkIds.add(chunk.getChunkId()) || chunk.getSourceUri() == null) { + throw notReady("invalid, duplicate, or untraceable fixture chunk"); } } if (bm25 != null) { diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java index 5aa15291c..fafea04d8 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalApiJsonTest.java @@ -316,6 +316,8 @@ public void documentsEveryErrorCodeMapping() { Map statuses = new HashMap<>(); statuses.put(RetrievalErrorCode.INVALID_REQUEST, 400); statuses.put(RetrievalErrorCode.UNSUPPORTED_OPTION, 400); + statuses.put(RetrievalErrorCode.AUTHENTICATION_REQUIRED, 401); + statuses.put(RetrievalErrorCode.FORBIDDEN, 403); statuses.put(RetrievalErrorCode.GRAPH_NOT_FOUND, 404); statuses.put(RetrievalErrorCode.INDEX_NOT_READY, 503); statuses.put(RetrievalErrorCode.RETRIEVAL_TIMEOUT, 504); diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java index 0a8a639dc..623b703ba 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/api/RetrievalPropertiesTest.java @@ -94,4 +94,20 @@ private static void assertInvalid(RetrievalProperties properties) { properties::validateConfiguration); Assertions.assertEquals(RetrievalErrorCode.INVALID_REQUEST, exception.getCode()); } + + @Test + public void remoteBindingRequiresOptInAndToken() { + RetrievalProperties properties = new RetrievalProperties(); + Assertions.assertEquals("127.0.0.1", properties.getServerHost()); + properties.setServerHost("0.0.0.0"); + assertInvalid(properties); + properties.setRemoteAccessEnabled(true); + assertInvalid(properties); + properties.setApiToken("test-token"); + properties.validateConfiguration(); + properties.setServerHost("::"); + properties.validateConfiguration(); + properties.setRemoteAccessEnabled(false); + assertInvalid(properties); + } } diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactTest.java index c04b176cd..2ad00ad2f 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/index/graph/GraphArtifactTest.java @@ -225,8 +225,8 @@ private static int[] graphCountOffsets(byte[] artifact) throws IOException { private static List chunks() { return Arrays.asList( - new TextChunk("c1", "doc-1", 0, 0, 10, 2, "Confucius"), - new TextChunk("c2", "doc-2", 0, 0, 9, 2, "philosophy")); + new TextChunk("c1", "doc-1", 0, 0, 10, 2, "Confucius").withSourceUri("fixture://doc-1"), + new TextChunk("c2", "doc-2", 0, 0, 9, 2, "philosophy").withSourceUri("fixture://doc-2")); } private static List entities() { diff --git a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/GraphRagMvpEndToEndTest.java b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/GraphRagMvpEndToEndTest.java index 91f824e4a..775133440 100644 --- a/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/GraphRagMvpEndToEndTest.java +++ b/geaflow-ai/src/test/java/org/apache/geaflow/ai/retrieval/service/GraphRagMvpEndToEndTest.java @@ -79,12 +79,23 @@ void executesAllMvpModesAndPreservesEvidence() throws Exception { RetrievalResponse bm25 = service.retrieve(request("BM25_ONLY", null)); RetrievalResponse vector = service.retrieve(request("VECTOR_ONLY", Arrays.asList(1.0, 0.0))); - RetrievalResponse graph = service.retrieve(request("GRAPH_ONLY", null)); + RetrievalRequest graphRequest = request("GRAPH_ONLY", null); + graphRequest.setRequestId("request-graph-trace"); + RetrievalResponse graph = service.retrieve(graphRequest); RetrievalResponse hybrid = service.retrieve(request("HYBRID", Arrays.asList(1.0, 0.0))); assertFalse(bm25.getEvidence().isEmpty()); assertFalse(vector.getEvidence().isEmpty()); assertFalse(graph.getEvidence().isEmpty()); + assertEquals("request-graph-trace", graph.getRequestId()); + for (Evidence evidence : graph.getEvidence()) { + for (org.apache.geaflow.ai.retrieval.model.evidence.GraphEvidenceTrace trace : evidence.getGraphTraces()) { + assertEquals(graph.getRequestId(), trace.getQueryId()); + assertFalse(graphRequest.getQuery().equals(trace.getQueryId())); + } + assertEquals("https://example.test/" + evidence.getChunks().get(0).getDocumentId(), + evidence.getSources().get(0).getSourceUri()); + } RetrievalRequest aliasRequest = request("GRAPH_ONLY", null); aliasRequest.setQuery("kongzi"); assertFalse(service.retrieve(aliasRequest).getEvidence().isEmpty()); @@ -176,7 +187,8 @@ void scalesHybridRecallAcrossTenThousandChunks() throws Exception { String text = index % 5 == 0 ? "Confucius discusses ethics in document " + index : "Background material about history in document " + index; - chunks.add(new TextChunk(chunkId, documentId, 0, 0, text.length(), 6, text)); + chunks.add(new TextChunk(chunkId, documentId, 0, 0, text.length(), 6, text, + null, null, "fixture://scaled/" + documentId)); vectors.put(chunkId, index % 2 == 0 ? new float[] {1.0F, 0.0F} : new float[] {0.0F, 1.0F}); if (index % 5 == 0) { graphEvidenceChunkIds.add(chunkId); @@ -206,6 +218,7 @@ void scalesHybridRecallAcrossTenThousandChunks() throws Exception { entities, edges, true); RetrievalRequest request = request("HYBRID", Arrays.asList(1.0, 0.0)); + request.setRequestId("stable-request"); request.setBudget(new RetrievalBudget(10, 5000, 400, 100)); RecallService service = new RecallService(registry); RetrievalResponse first = service.retrieve(request); @@ -231,6 +244,7 @@ void producesStableEvidenceAndTraceCounters() throws Exception { Fixture fixture = fixture(); RecallService service = new RecallService(fixture.registry); RetrievalRequest request = request("HYBRID", Arrays.asList(1.0, 0.0)); + request.setRequestId("stable-request"); RetrievalResponse first = service.retrieve(request); RetrievalResponse second = service.retrieve(request); assertEquals(first.getEvidence(), second.getEvidence()); @@ -668,9 +682,12 @@ void directFacadeRequestsAlsoHaveIndependentExecutionState() throws Exception { private Fixture fixture() throws Exception { List chunks = Arrays.asList( - new TextChunk("c1", "doc-1", 0, 0, 36, 6, "Confucius taught ethics and philosophy"), - new TextChunk("c2", "doc-2", 0, 0, 25, 4, "Astronomy studies distant stars"), - new TextChunk("c3", "doc-3", 0, 0, 35, 6, "Confucius and astronomy are studied")); + new TextChunk("c1", "doc-1", 0, 0, 36, 6, + "Confucius taught ethics and philosophy", null, null, "https://example.test/doc-1"), + new TextChunk("c2", "doc-2", 0, 0, 25, 4, + "Astronomy studies distant stars", null, null, "https://example.test/doc-2"), + new TextChunk("c3", "doc-3", 0, 0, 35, 6, + "Confucius and astronomy are studied", null, null, "https://example.test/doc-3")); IngestionContext context = context(); Map vectors = new HashMap<>(); vectors.put("c1", new float[] {1.0F, 0.0F}); From a93c5978f71c819f100e22b83811ac500b51c364 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 10 Oct 2026 20:34:13 +0800 Subject: [PATCH 28/30] fix(graphrag): make dataset ingestion identities and spills deterministic --- tools/graphrag/ingest/normalize.py | 14 +- tools/graphrag/ingest/pipeline.py | 162 ++++++++++++++---- .../graphrag/ingest/sources/source_loader.py | 32 ++-- tools/graphrag/tests/test_normalize.py | 20 +++ tools/graphrag/tests/test_pipeline.py | 59 ++++++- tools/graphrag/tests/test_source_loader.py | 28 +++ 6 files changed, 268 insertions(+), 47 deletions(-) diff --git a/tools/graphrag/ingest/normalize.py b/tools/graphrag/ingest/normalize.py index 51bf42796..af67dc6c0 100644 --- a/tools/graphrag/ingest/normalize.py +++ b/tools/graphrag/ingest/normalize.py @@ -3,6 +3,7 @@ from __future__ import annotations import hashlib +import json import re import unicodedata from dataclasses import dataclass @@ -25,6 +26,12 @@ def __post_init__(self) -> None: if self.token_characters < 1: raise ValueError("token_characters must be positive") + def policy_version(self) -> str: + effective = json.dumps({"size": self.size, "overlap": self.overlap, + "token_characters": self.token_characters}, sort_keys=True, + separators=(",", ":")) + return CHUNKING_VERSION + ":" + hashlib.sha256(effective.encode("utf-8")).hexdigest() + def normalize_text(value: str) -> str: text = unicodedata.normalize("NFC", value) @@ -61,6 +68,7 @@ def documents(records: Iterable[Dict[str, Any]]) -> Iterator[Dict[str, Any]]: "answers": record.get("answers", []), "supporting_facts": record.get("supporting_facts", []), "normalization_version": NORMALIZATION_VERSION, + "source_uri": record.get("source_uri"), } @@ -82,7 +90,8 @@ def chunk_document(document: Dict[str, Any], config: ChunkingConfig) -> List[Dic actual_end = actual_start + len(chunk_text) if chunk_text: digest = hashlib.sha256(chunk_text.encode("utf-8")).hexdigest() - chunk_id = hashlib.sha256((document["document_id"] + "\0" + str(index) + "\0" + digest) + chunk_id = hashlib.sha256((document["document_id"] + "\0" + str(index) + "\0" + digest + + "\0" + config.policy_version()) .encode("utf-8")).hexdigest() result.append({ "chunk_id": chunk_id, @@ -93,7 +102,8 @@ def chunk_document(document: Dict[str, Any], config: ChunkingConfig) -> List[Dic "token_estimate": (len(chunk_text) + config.token_characters - 1) // config.token_characters, "text": chunk_text, - "policy_version": CHUNKING_VERSION, + "policy_version": config.policy_version(), + "source_uri": document.get("source_uri"), "text_hash": digest, "source_hash": document["source_hash"], }) diff --git a/tools/graphrag/ingest/pipeline.py b/tools/graphrag/ingest/pipeline.py index 667378195..e9229430b 100644 --- a/tools/graphrag/ingest/pipeline.py +++ b/tools/graphrag/ingest/pipeline.py @@ -1,16 +1,18 @@ -"""Restartable offline ingestion orchestration for the Week 2 artifacts.""" +"""Restartable offline ingestion with bounded batches and disk-backed graph merging.""" from __future__ import annotations import hashlib import json import os +import shutil +import sqlite3 import tempfile from pathlib import Path -from typing import Any, Callable, Dict, Iterable, List, Optional +from typing import Any, Callable, Dict, Iterable, Iterator -from .entity_graph import graph_records, write_graph_artifact -from .normalize import ChunkingConfig, chunks, documents +from .entity_graph import graph_record_stream, write_graph_artifact +from .normalize import ChunkingConfig, chunk_document, documents from .sources.source_loader import SourceLoader, SourceManifest @@ -18,53 +20,145 @@ class ImportFailure(RuntimeError): """Raised after an import is recorded as FAILED.""" +class JsonlRecords: + """Repeatable, sized iterable that opens a fresh stream for each builder pass.""" + + def __init__(self, path: Path, count: int): + self.path = path + self.count = count + + def __len__(self) -> int: + return self.count + + def __iter__(self) -> Iterator[Dict[str, Any]]: + with self.path.open(encoding="utf-8") as stream: + for line in stream: + yield json.loads(line) + + +def _write_record(output: Any, record: Dict[str, Any]) -> None: + output.write(json.dumps(record, sort_keys=True, ensure_ascii=False, separators=(",", ":")) + "\n") + + +def _merge_graph(connection: sqlite3.Connection, kind: str, record: Dict[str, Any]) -> None: + key = record["vertex_id" if kind == "vertex" else "edge_id"] + record = dict(record) + for field in ("source_document_ids", "source_chunk_ids"): + for reference in record.pop(field, []): + connection.execute("INSERT OR IGNORE INTO provenance(kind,id,field,reference) VALUES (?,?,?,?)", + (kind, key, field, reference)) + row = connection.execute("SELECT payload FROM graph WHERE kind=? AND id=?", (kind, key)).fetchone() + if row is not None: + previous = json.loads(row[0]) + if kind == "vertex": + previous["canonical_name"] = min(previous.get("canonical_name", ""), + record.get("canonical_name", "")) + record = previous + connection.execute("INSERT OR REPLACE INTO graph(kind,id,payload) VALUES (?,?,?)", + (kind, key, json.dumps(record, sort_keys=True, ensure_ascii=False))) + + +def _graph_values(connection: sqlite3.Connection, kind: str) -> Iterator[Dict[str, Any]]: + for key, payload in connection.execute("SELECT id,payload FROM graph WHERE kind=? ORDER BY id", (kind,)): + value = json.loads(payload) + fields = ("source_document_ids", "source_chunk_ids") if kind == "edge" else ("source_document_ids",) + for field in fields: + value[field] = [row[0] for row in connection.execute( + "SELECT reference FROM provenance WHERE kind=? AND id=? AND field=? ORDER BY reference", + (kind, key, field))] + yield value + + def run_import(manifest: SourceManifest, output_directory: Path, loader: SourceLoader, - chunking: ChunkingConfig, bm25_builder: Callable[[Path, List[Dict[str, Any]]], Path], - vector_builder: Callable[[Path, List[Dict[str, Any]]], Path]) -> Dict[str, Any]: - """Build one immutable version; publish the metadata file only after every artifact succeeds.""" + chunking: ChunkingConfig, bm25_builder: Callable[[Path, Iterable[Dict[str, Any]]], Path], + vector_builder: Callable[[Path, Iterable[Dict[str, Any]]], Path]) -> Dict[str, Any]: + """Publish one immutable version only after all builders finish reading canonical spills.""" output_directory.mkdir(parents=True, exist_ok=True) - version = hashlib.sha256((manifest.dataset + "\0" + manifest.split + "\0" + manifest.sha256 - + "\0" + manifest.parser_version + "\0" + repr(chunking)) + identity = {"dataset": manifest.dataset, "dataset_release": manifest.dataset_release, + "split": manifest.split, "sha256": manifest.sha256, + "parser_version": manifest.parser_version, "policy_version": chunking.policy_version(), + "pipeline_version": "disk-spill-v1", "source_uri": manifest.source_uri, + "cache_uri": Path(manifest.cache_path).expanduser().resolve().as_uri() + if not manifest.source_uri and manifest.cache_path else None} + version = hashlib.sha256(json.dumps(identity, sort_keys=True, separators=(",", ":")) .encode("utf-8")).hexdigest()[:16] final = output_directory / version if final.exists() and (final / "READY.json").is_file(): return json.loads((final / "READY.json").read_text(encoding="utf-8")) staging = Path(tempfile.mkdtemp(prefix=version + ".", dir=output_directory)) - metadata: Dict[str, Any] = { - "dataset": manifest.dataset, "dataset_release": manifest.dataset_release, - "split": manifest.split, "sha256": manifest.sha256, "version": version, - "state": "IMPORTING", "parser_version": manifest.parser_version, - "errors": [], - } + metadata: Dict[str, Any] = dict(identity, version=version, state="IMPORTING", errors=[]) try: - source_records = list(loader.load(manifest)) - normalized = list(documents(source_records)) - normalized_chunks = list(chunks(normalized, chunking)) - (staging / "chunks.jsonl").write_text( - "".join(json.dumps(item, sort_keys=True, ensure_ascii=False) + "\n" - for item in normalized_chunks), encoding="utf-8") - vertices, edges = graph_records(source_records, normalized_chunks) - graph_manifest = write_graph_artifact(staging, "graph", vertices, edges) + document_count = 0 + chunk_count = 0 + database_path = staging / "graph-staging.sqlite" + connection = sqlite3.connect(database_path) + try: + connection.execute("PRAGMA cache_size=-2048") + connection.execute("CREATE TABLE graph(kind TEXT, id TEXT, payload TEXT, PRIMARY KEY(kind,id))") + connection.execute("CREATE TABLE provenance(kind TEXT, id TEXT, field TEXT, reference TEXT, " + "PRIMARY KEY(kind,id,field,reference))") + connection.execute("CREATE TABLE documents(id TEXT PRIMARY KEY)") + with (staging / "documents.jsonl").open("w", encoding="utf-8") as document_output, \ + (staging / "chunks.jsonl").open("w", encoding="utf-8") as chunk_output: + for batch in loader.load_batches(manifest): + for record in batch: + connection.execute("INSERT INTO documents(id) VALUES (?)", (record["document_id"],)) + document = next(documents([record])) + if not document.get("source_uri"): + raise ImportFailure("document source URI is missing") + _write_record(document_output, document) + document_count += 1 + document_chunks = chunk_document(document, chunking) + for chunk in document_chunks: + _write_record(chunk_output, chunk) + chunk_count += 1 + # Only one document's chunks are resident; merged graph records live on disk. + for vertex, edge in graph_record_stream([record], document_chunks): + if vertex is not None: + _merge_graph(connection, "vertex", vertex) + if edge is not None: + _merge_graph(connection, "edge", edge) + connection.commit() + write_graph_artifact(staging, "graph", _graph_values(connection, "vertex"), + _graph_values(connection, "edge")) + finally: + connection.close() + database_path.unlink(missing_ok=True) + graph_manifest = json.loads((staging / "graph/manifest.json").read_text(encoding="utf-8")) metadata["state"] = "INDEXING" - bm25_path = bm25_builder(staging, normalized_chunks) - vector_path = vector_builder(staging, normalized_chunks) + canonical_chunks = JsonlRecords(staging / "chunks.jsonl", chunk_count) + bm25_path = bm25_builder(staging, canonical_chunks) + vector_path = vector_builder(staging, canonical_chunks) metadata.update({ - "state": "READY", "documents": len(normalized), "chunks": len(normalized_chunks), - "vertices": len(vertices), "edges": len(edges), "graph_manifest": "graph/manifest.json", + "state": "READY", "documents": document_count, "chunks": chunk_count, + "vertices": graph_manifest["vertex_count"], "edges": graph_manifest["edge_count"], + "graph_manifest": "graph/manifest.json", "bm25_artifact": str(Path(bm25_path).relative_to(staging)), "vector_artifact": str(Path(vector_path).relative_to(staging)), }) for relative in (metadata["graph_manifest"], metadata["bm25_artifact"], metadata["vector_artifact"]): - if not (staging / relative).exists(): - raise ImportFailure("artifact does not exist before publication: %s" % relative) - ready_path = staging / "READY.json" - ready_path.write_text(json.dumps(metadata, sort_keys=True, indent=2) + "\n", encoding="utf-8") - os.replace(staging, final) + candidate = staging / relative + if not candidate.exists() or staging.resolve() not in candidate.resolve().parents: + raise ImportFailure("artifact is missing or outside staging: %s" % relative) + (staging / "READY.json").write_text(json.dumps(metadata, sort_keys=True, indent=2) + "\n", + encoding="utf-8") + # The target is a nonempty immutable directory; rename cannot replace a READY winner. + try: + os.rename(staging, final) + except OSError: + if not (final / "READY.json").is_file(): + raise + winner = json.loads((final / "READY.json").read_text(encoding="utf-8")) + if winner != metadata: + raise ImportFailure("existing import metadata conflicts with this attempt") + return winner return metadata except Exception as failure: metadata["state"] = "FAILED" metadata["errors"] = [str(failure)] - (staging / "FAILED.json").write_text(json.dumps(metadata, sort_keys=True, indent=2) + "\n", - encoding="utf-8") + (output_directory / (staging.name + ".FAILED.json")).write_text( + json.dumps(metadata, sort_keys=True, indent=2) + "\n", encoding="utf-8") raise ImportFailure("import failed for %s/%s: %s" % (manifest.dataset, manifest.split, failure)) from failure + finally: + shutil.rmtree(staging, ignore_errors=True) diff --git a/tools/graphrag/ingest/sources/source_loader.py b/tools/graphrag/ingest/sources/source_loader.py index 4f70ed278..6e09cb079 100644 --- a/tools/graphrag/ingest/sources/source_loader.py +++ b/tools/graphrag/ingest/sources/source_loader.py @@ -51,8 +51,8 @@ def from_file(cls, path: Path) -> "SourceManifest": checksum = str(raw["sha256"]).lower() if len(checksum) != 64 or any(char not in "0123456789abcdef" for char in checksum): raise SourceError("manifest %s has invalid sha256" % path) - if raw.get("dataset") == "2wikimultihopqa" and raw.get("split") not in ("dev", "test"): - raise SourceError("2wikimultihopqa manifest split must be dev or test") + if raw.get("dataset") == "2wikimultihopqa" and raw.get("split") not in ("train", "dev", "test"): + raise SourceError("2wikimultihopqa manifest split must be train, dev or test") return cls( dataset=str(raw["dataset"]), dataset_release=str(raw["dataset_release"]), @@ -89,8 +89,10 @@ def load_batches(self, manifest: SourceManifest) -> Iterator[List[Dict[str, Any] if path.stat().st_size > self.max_download_bytes: raise SourceError("source exceeds configured size limit: %s" % path) for number, raw in enumerate(self._records(path, self.max_uncompressed_bytes, - self.max_zip_entries), start=1): - batch.append(self._canonical_record(manifest, number, raw)) + self.max_zip_entries, manifest.split), start=1): + record = self._canonical_record(manifest, number, raw) + record["source_uri"] = manifest.source_uri or path.resolve().as_uri() + batch.append(record) if len(batch) == self.batch_size: yield batch batch = [] @@ -161,22 +163,32 @@ def _verify_checksum(path: Path, manifest: SourceManifest) -> None: @staticmethod def _records(path: Path, max_uncompressed_bytes: int = 1024 * 1024 * 1024, - max_zip_entries: int = 16) -> Iterable[Any]: + max_zip_entries: int = 16, split: Optional[str] = None) -> Iterable[Any]: if path.suffix == ".zip": try: with zipfile.ZipFile(path) as archive: names = sorted(name for name in archive.namelist() - if name.endswith((".json", ".jsonl", ".json.gz", ".jsonl.gz"))) + if name.lower().endswith((".json", ".jsonl", ".json.gz", ".jsonl.gz"))) if len(archive.namelist()) > max_zip_entries: raise SourceError("ZIP source %s has too many entries" % path) - if len(names) != 1: - raise SourceError("ZIP source %s must contain exactly one JSON data file" % path) - descriptor, temporary_name = tempfile.mkstemp(suffix=Path(names[0]).suffix) + matches = [name for name in names if Path(name).name.lower() + in {str(split).lower() + suffix for suffix in + (".json", ".jsonl", ".json.gz", ".jsonl.gz")}] + if not matches and len(names) == 1: + # Preserve single-file archive compatibility, but reject a wrong named split. + basename = Path(names[0]).name.lower().split(".")[0] + if basename not in ("train", "dev", "test"): + matches = names + if len(matches) != 1: + raise SourceError("ZIP source %s split %s has %d matching data files" % + (path, split, len(matches))) + selected = matches[0] + descriptor, temporary_name = tempfile.mkstemp(suffix=Path(selected).suffix.lower()) os.close(descriptor) extracted = Path(temporary_name) try: total = 0 - with archive.open(names[0]) as source, extracted.open("wb") as target: + with archive.open(selected) as source, extracted.open("wb") as target: while True: block = source.read(1024 * 1024) if not block: diff --git a/tools/graphrag/tests/test_normalize.py b/tools/graphrag/tests/test_normalize.py index 8a3a616fb..5082ef562 100644 --- a/tools/graphrag/tests/test_normalize.py +++ b/tools/graphrag/tests/test_normalize.py @@ -38,6 +38,26 @@ def test_rejects_invalid_overlap(self): with self.assertRaises(ValueError): ChunkingConfig(size=10, overlap=10) + def test_chunk_policy_is_part_of_identity(self): + record = { + "dataset": "d", "dataset_release": "r", "split": "dev", "document_id": "x", + "question": "q", "paragraphs": [{"title": "T", "text": "alpha beta gamma"}], + } + document = list(documents([record]))[0] + first = chunk_document(document, ChunkingConfig(size=12, overlap=3, token_characters=4))[0] + second = chunk_document(document, ChunkingConfig(size=12, overlap=2, token_characters=4))[0] + self.assertNotEqual(first["chunk_id"], second["chunk_id"]) + self.assertNotEqual(first["policy_version"], second["policy_version"]) + + def test_short_document_identity_depends_on_all_effective_policy_parameters(self): + document = {"document_id": "short", "text": "text", "source_hash": "hash"} + configs = [ChunkingConfig(50, 0, 4), ChunkingConfig(60, 0, 4), + ChunkingConfig(50, 1, 4), ChunkingConfig(50, 0, 5)] + results = [chunk_document(document, config)[0] for config in configs] + self.assertEqual(4, len({item["chunk_id"] for item in results})) + self.assertEqual(4, len({item["policy_version"] for item in results})) + self.assertEqual({"text"}, {item["text"] for item in results}) + if __name__ == "__main__": unittest.main() diff --git a/tools/graphrag/tests/test_pipeline.py b/tools/graphrag/tests/test_pipeline.py index 701dd1a41..298ebd3c8 100644 --- a/tools/graphrag/tests/test_pipeline.py +++ b/tools/graphrag/tests/test_pipeline.py @@ -1,11 +1,12 @@ import hashlib import json import tempfile +import tracemalloc import unittest from pathlib import Path from tools.graphrag.ingest.normalize import ChunkingConfig -from tools.graphrag.ingest.pipeline import ImportFailure, run_import +from tools.graphrag.ingest.pipeline import ImportFailure, JsonlRecords, run_import from tools.graphrag.ingest.sources.source_loader import SourceLoader, SourceManifest @@ -43,6 +44,62 @@ def test_failed_index_is_not_ready(self): lambda output, chunks: (_ for _ in ()).throw(ValueError("bad index")), lambda output, chunks: output / "vector") self.assertEqual(0, len(list((root / "versions").glob("*/READY.json")))) + self.assertEqual([], [path for path in (root / "versions").iterdir() if path.is_dir()]) + self.assertEqual(1, len(list((root / "versions").glob("*.FAILED.json")))) + + def test_release_changes_artifact_version(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source.jsonl" + source.write_text(json.dumps({"id": "d", "question": "q", + "context": [["A", ["A text"]]]}), encoding="utf-8") + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + build = lambda name: lambda output, chunks: _artifact(output, name, chunks) + first = SourceManifest("d", "release-1", "dev", None, str(source), checksum, "p") + second = SourceManifest("d", "release-2", "dev", None, str(source), checksum, "p") + first_result = run_import(first, root / "versions", SourceLoader(), ChunkingConfig(50, 5), + build("bm25"), build("vector")) + second_result = run_import(second, root / "versions", SourceLoader(), ChunkingConfig(50, 5), + build("bm25"), build("vector")) + self.assertNotEqual(first_result["version"], second_result["version"]) + + def test_large_import_spills_repeatable_chunks_and_cleans_transient_database(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "source.jsonl" + with source.open("w", encoding="utf-8") as output: + for index in range(10000): + output.write(json.dumps({"id": str(index), "question": "q", + "context": [["Shared", ["Shared evidence " + str(index)]]] + }) + "\n") + manifest = SourceManifest("d", "r", "dev", None, str(source), + hashlib.sha256(source.read_bytes()).hexdigest(), "p") + passes = [] + + def build(output, chunks): + self.assertIsInstance(chunks, JsonlRecords) + self.assertEqual(10000, len(chunks)) + digest = hashlib.sha256() + for chunk in chunks: + self.assertTrue(chunk["source_uri"].startswith("file://")) + digest.update(chunk["chunk_id"].encode("utf-8")) + passes.append(digest.hexdigest()) + return _artifact(output, "index-" + str(len(passes)), chunks) + + tracemalloc.start() + try: + result = run_import(manifest, root / "versions", SourceLoader(batch_size=32), + ChunkingConfig(100, 0), build, build) + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + self.assertEqual(passes[0], passes[1]) + self.assertEqual(10000, result["documents"]) + self.assertEqual(10000, result["chunks"]) + self.assertLess(peak, 12 * 1024 * 1024) + final = root / "versions" / result["version"] + self.assertTrue((final / "documents.jsonl").is_file()) + self.assertFalse((final / "graph-staging.sqlite").exists()) def _artifact(output, name, chunks): diff --git a/tools/graphrag/tests/test_source_loader.py b/tools/graphrag/tests/test_source_loader.py index 6c28a7683..8dc5d2d47 100644 --- a/tools/graphrag/tests/test_source_loader.py +++ b/tools/graphrag/tests/test_source_loader.py @@ -3,6 +3,7 @@ import json import tempfile import unittest +import zipfile from pathlib import Path from tools.graphrag.ingest.sources.source_loader import SourceError, SourceLoader, SourceManifest @@ -84,6 +85,33 @@ def test_limits_uncompressed_gzip_source(self): with self.assertRaisesRegex(SourceError, "gzip source exceeds"): list(SourceLoader(max_uncompressed_bytes=len(contents) - 1).load(manifest)) + def test_selects_manifest_split_from_multi_file_zip(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source = root / "data.zip" + payload = lambda value: json.dumps({"id": value, "question": "q", + "context": [["T", ["text"]]]}) + with zipfile.ZipFile(source, "w") as archive: + archive.writestr("train.json", payload("train")) + archive.writestr("dev.json", payload("dev")) + archive.writestr("test.json", payload("test")) + checksum = hashlib.sha256(source.read_bytes()).hexdigest() + manifest = SourceManifest("2wikimultihopqa", "r", "dev", None, str(source), checksum, "p") + records = list(SourceLoader().load(manifest)) + self.assertEqual(["dev"], [record["document_id"] for record in records]) + + def test_zip_rejects_missing_or_ambiguous_split(self): + for entries in (("train.json", "test.json"), ("dev.json", "folder/DEV.JSON")): + with self.subTest(entries=entries), tempfile.TemporaryDirectory() as directory: + source = Path(directory) / "data.zip" + with zipfile.ZipFile(source, "w") as archive: + for name in entries: + archive.writestr(name, "[]") + manifest = SourceManifest("2wikimultihopqa", "r", "dev", None, str(source), + hashlib.sha256(source.read_bytes()).hexdigest(), "p") + with self.assertRaisesRegex(SourceError, "matching data files"): + list(SourceLoader().load(manifest)) + if __name__ == "__main__": unittest.main() From e025eb3518da312e0e05034921c693ad1c2a61c8 Mon Sep 17 00:00:00 2001 From: aotenjou Date: Sat, 10 Oct 2026 20:34:40 +0800 Subject: [PATCH 29/30] docs(ai): update MVP review verification --- geaflow-ai/README-hybrid-recall.md | 60 +++++++++++++++++++++++++++++- 1 file changed, 58 insertions(+), 2 deletions(-) diff --git a/geaflow-ai/README-hybrid-recall.md b/geaflow-ai/README-hybrid-recall.md index ccb8d275b..dd2fa7d8a 100644 --- a/geaflow-ai/README-hybrid-recall.md +++ b/geaflow-ai/README-hybrid-recall.md @@ -30,6 +30,10 @@ Build both artifacts with `LuceneBm25IndexBuilder` and `OfflineVectorIndexBuilde `IngestionContext` and canonical chunks. Supply precomputed vectors; no embedding service is called. Load the published artifacts and graph records once: +Each canonical chunk must carry its real `sourceUri` (or use `TextChunk.withSourceUri` before building). +The ingestion pipeline copies this URI from its `SourceDocument` and rejects missing or inconsistent +provenance. The registry rejects chunks without a URI; citations use this URI and the chunk offsets. + ```java RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); registry.load(graphName, graphVersion, indexVersion, bm25Path, vectorPath, @@ -83,7 +87,9 @@ document/chunk identity, uses stable evidence ID for ties, and receives consecut GRAPH_ONLY retains deterministic anchor resolution and one-hop expansion. Both edge scans and unique chunk evidence are capped by `maxCandidates`, including edges that cite several chunks. Paths and anchor provenance survive merging. Hybrid always selects BM25, Vector, and Graph; graph provenance -contains the stable query, graph version, anchor match, path, and supporting chunk IDs. +contains the request ID, graph version, anchor match, path, and supporting chunk IDs. A supplied request +ID is resolved before any channel runs; otherwise one is generated. Every graph evidence `queryId` +matches the response and retrieval trace request ID. The graph evidence trace does not store query text. ## Errors and trace @@ -114,14 +120,64 @@ before vector source/version and the vector payload. Legacy artifacts without th rejected with `INDEX_NOT_READY`; rebuild and publish a fresh graph/index version before switching requests. Existing immutable artifacts are validated rather than silently overwritten. +Chunk fingerprints now include `sourceUri`; artifacts built with the previous fingerprint must be +rebuilt. All three Java artifact builders use a shared publisher with bounded JVM locks and an OS file +lock per canonical target. Publication checks and rename run under that lock; competing builders reuse +an equivalent winner and reject conflicting content. The `.publish.lock` file stays beside the artifact +so every process locks the same inode. Do not delete it while builders are running. + +`MetadataPublisher.begin` returns an `ImportAttempt` ownership token. `publish` and `fail` require that +token. A failed `begin` does not trigger `fail`, and a competing ingestion cannot fail the owner. +Custom publishers must implement the token-aware lifecycle. + +## HTTP access and keyword index lifecycle + +The server binds to `127.0.0.1` by default. Non-loopback binding requires both +`retrieval.remote-access-enabled=true` and a nonblank `retrieval.api-token`; invalid combinations fail +configuration initialization. Environment equivalents are `GEAFLOW_SERVER_HOST`, +`GEAFLOW_REMOTE_ACCESS_ENABLED`, and `GEAFLOW_RETRIEVAL_API_TOKEN`. Remote mode requires +`Authorization: Bearer ` on retrieval and the legacy `/graph/*` and `/query/*` data APIs. +Missing credentials return `401`; invalid credentials return `403`. `/health` remains available. +The single token grants access to the configured server; per-graph authorization is future work. + +Responses use the same configured limits as request validation, including a raised `max-top-k`. +Keyword retrieval builds one complete Lucene store on the first request and reuses it. Candidate +budgets cap search results evaluated rather than the vertices included in the cached index. Graph +updates invalidate the cache; server replacement and shutdown close it. Cold construction and lock +acquisition check the request deadline. Search store closure is idempotent. + +## Python ingestion + +Multi-file ZIP sources select the requested train/dev/test entry by case-insensitive basename; +missing or ambiguous splits fail. Single archives with a generic data filename remain supported. +Artifact identity includes dataset release and the effective chunk policy (size, overlap, and token +estimate settings). The policy fingerprint also participates in chunk IDs, even for short documents. + +Ingestion normalizes one batch at a time, writes canonical `documents.jsonl` and `chunks.jsonl`, and +merges graph records and provenance in a temporary SQLite database. Index builder callbacks receive a +repeatable `JsonlRecords` iterable supporting `len`; callbacks must iterate rather than index a list. +Graph export reads sorted records from disk. Temporary database and staging files are removed on +success or failure; a failed attempt retains only a separate `*.FAILED.json` diagnostic. Canonical +JSONL files remain in the published version for reproducibility. Memory depends on batch size, an +individual document's chunks, and a single exported graph record's provenance, rather than simultaneous +materialization of the entire corpus. This does not bound allocations made by external builder callbacks. + Run the retrieval and HTTP regression suite with JDK 11, which is compatible with the repository's JaCoCo 0.8.8: ```bash -mvn -pl geaflow-ai -Dtest='org.apache.geaflow.ai.retrieval.**.*Test,MemoryServerTest' test +mvn -pl geaflow-ai -am clean test -Drat.skip=true \ + -Dtest='org.apache.geaflow.ai.retrieval.**.*Test,MemoryServerTest,GeaFlowMemoryServerTest,GraphMemoryKeywordSearchTest,SearchStoreTest' \ + -Dsurefire.failIfNoSpecifiedTests=false +python -m pytest tools/graphrag/tests -q ``` The Maven lifecycle also runs Checkstyle and Apache RAT. New retrieval production sources have no blanket Checkstyle suppressions. The suite covers bounded scoring, partial-channel failures, deadlines using a controlled clock, graph limits, version mismatches, JSON trace round trips, concurrent request isolation, and rejection of core modes at the HTTP boundary. +The PR #876 regressions also cover concurrent import ownership, same-target publication for all three +artifacts, competing JVMs, URI citations, request IDs, HTTP authentication and custom limits, cached +keyword retrieval, repeated close, ZIP splits, release/policy identities, and a 10,000-document spill +fixture whose traced Python heap stays below 12 MiB. This fixture uses deterministic builder callbacks; +it is not a full HotpotQA/2Wiki or external embedding benchmark. From a8c3ce8d9be9ad7d9dcf73100f41f7c3545003eb Mon Sep 17 00:00:00 2001 From: aotenjou Date: Mon, 12 Oct 2026 13:00:38 +0800 Subject: [PATCH 30/30] docs(ai): remove obsolete retrieval README files --- geaflow-ai/README-hybrid-recall.md | 183 --------------------------- geaflow-ai/README-week1-retrieval.md | 44 ------- 2 files changed, 227 deletions(-) delete mode 100644 geaflow-ai/README-hybrid-recall.md delete mode 100644 geaflow-ai/README-week1-retrieval.md diff --git a/geaflow-ai/README-hybrid-recall.md b/geaflow-ai/README-hybrid-recall.md deleted file mode 100644 index dd2fa7d8a..000000000 --- a/geaflow-ai/README-hybrid-recall.md +++ /dev/null @@ -1,183 +0,0 @@ - - -# Fixed Hybrid Recall MVP - -`RecallService` is a session-free Java facade over offline, versioned fixtures. Its structured request -supports `BM25_ONLY`, `VECTOR_ONLY`, `GRAPH_ONLY`, and `HYBRID`; `KEYWORD` aliases `BM25_ONLY` in this core -facade. The existing HTTP `/api/v1/retrievals` endpoint supports only `KEYWORD`. Other modes return -`UNSUPPORTED_OPTION` there; HTTP Hybrid serving is not part of this implementation. - -## Load and retrieve - -Build both artifacts with `LuceneBm25IndexBuilder` and `OfflineVectorIndexBuilder`, using the same -`IngestionContext` and canonical chunks. Supply precomputed vectors; no embedding service is called. -Load the published artifacts and graph records once: - -Each canonical chunk must carry its real `sourceUri` (or use `TextChunk.withSourceUri` before building). -The ingestion pipeline copies this URI from its `SourceDocument` and rejects missing or inconsistent -provenance. The registry rejects chunks without a URI; citations use this URI and the chunk offsets. - -```java -RetrievalFixtureRegistry registry = new RetrievalFixtureRegistry(); -registry.load(graphName, graphVersion, indexVersion, bm25Path, vectorPath, - chunks, entities, edges, vectorSource, vectorVersion); -RecallService service = new RecallService(registry, new WeightedRrfConfig(1.0, 1.0, 1.0, 60)); - -RetrievalRequest request = new RetrievalRequest(); -request.setGraphName(graphName); -request.setGraphVersion(graphVersion); -request.setIndexVersion(indexVersion); -request.setMode("HYBRID"); -request.setQuery("Confucius"); -request.setQueryVector(Arrays.asList(1.0, 0.0)); // Must match the artifact's dimensions. -request.setBudget(new RetrievalBudget(10, 3000, 100, 4096)); -request.setAllowPartialResults(true); -RetrievalResponse result = service.retrieve(request); -``` - -`graphVersion` and `indexVersion` are required and must match. The registry validates the persisted -artifact graph name, graph/index versions, and a fingerprint of the complete canonical chunk payload. -Optional request vector source/version constraints must match the loaded vector artifact. Keep the -readers alive while their fixtures are in use; callers own the lifecycle of registered readers. -The legacy overload accepting text, a float vector, and an absolute deadline requires a service -constructed directly with readers, rather than a registry-backed service. - -## Budgets, ranking, and limitations - -Hybrid runs BM25, Vector, and one-hop Graph under one monotonic deadline starting at request entry. It -allocates a deterministic three-way budget: BM25 receives the base plus the first remainder, Vector -receives the base plus the second remainder, and Graph receives the base, where `base = floor(maxCandidates / 3)`. -Hybrid requires `maxCandidates >= 3`; all modes require `1 <= topK <= maxCandidates`. -`topK` is applied only after merging and ranking. Individual retriever convenience methods still -return at most their requested `topK`; the facade uses the full bounded candidate lists for fusion. - -`maxCandidates` caps actual scoring work: matched Lucene documents scored for BM25, vectors evaluated -for Vector, and edges examined for GRAPH_ONLY. Lucene uses a collector that checks the deadline before -scoring each matched document. Parsing, query setup, and sorting are checked at stage boundaries; -deadlines are cooperative, rather than a thread-interruption guarantee. - -Both text and vector channels are bounded approximate recall when their evaluation budget is smaller -than the searchable data. BM25 collects Lucene traversal order; Vector scans stable chunk-ID order and -sorts only the evaluated prefix by cosine similarity. The vector result is **not** a global nearest -neighbor guarantee. Increasing the budget widens the evaluated set. ANN search and adaptive candidate -selection remain future work. - -Weighted RRF uses `sum(weight[channel] / (rankConstant + rank[channel]))`, with defaults `1.0`, `1.0`, -`1.0`, and `60`. Weights must be finite and non-negative, at least one positive; the rank constant is positive. -Zero-weight channels still run and retain their scores/provenance. Evidence merges by canonical -document/chunk identity, uses stable evidence ID for ties, and receives consecutive final ranks. - -GRAPH_ONLY retains deterministic anchor resolution and one-hop expansion. Both edge scans and unique -chunk evidence are capped by `maxCandidates`, including edges that cite several chunks. Paths and -anchor provenance survive merging. Hybrid always selects BM25, Vector, and Graph; graph provenance -contains the request ID, graph version, anchor match, path, and supporting chunk IDs. A supplied request -ID is resolved before any channel runs; otherwise one is generated. Every graph evidence `queryId` -matches the response and retrieval trace request ID. The graph evidence trace does not store query text. - -## Errors and trace - -Missing, empty, non-finite, float-overflowing, zero-norm, or dimension-mismatched query vectors return -`INVALID_REQUEST` for vector modes. Unknown modes and non-sequential execution return -`UNSUPPORTED_OPTION`. Malformed JSON is rejected at the codec boundary. - -The loader records channel-specific missing/corrupt artifact failures as `INDEX_NOT_READY` and only -publishes a READY fixture when at least one channel is usable. With partial results disabled, a failed -selected channel fails the request. With partial results enabled, its evidence is omitted and the -successful channel's evidence remains. If the request deadline expires, strict requests return -`RETRIEVAL_TIMEOUT`; partial requests retain fully evaluated evidence and mark unfinished channels. - -Trace fields include versions, selected channels, effective topK, per-channel budgets, actual evaluated -counts, candidate counts before fusion/topK, graph anchors/edges/neighbors/vertices, validation errors, -elapsed nanoseconds, and the effective three-channel RRF configuration. -`channelStatuses` distinguishes `SUCCESS`, `DEGRADED`, and `NOT_RUN`; `channelStopReasons` carries typed -`COMPLETED`, `CANDIDATE_BUDGET`, `CANDIDATE_LIMIT`, `EDGE_SCAN_LIMIT`, `NO_RELIABLE_ANCHOR`, `DEADLINE`, -or `INDEX_NOT_READY`. `degradationReasons` carries explanatory text separately. Overall `stopReason` -prioritizes deadline, graph trace validation, graph edge/candidate limits, candidate budget, then final topK -truncation, then no results or normal completion. Candidate-budget exhaustion is a valid bounded stage, -not an index failure. A fusion-stage deadline is also reflected in its stage status. - -## Artifact compatibility and verification - -BM25 stores identity in Lucene commit metadata. Vector artifacts use `GEAFLOW-VECTOR-2`, adding identity -before vector source/version and the vector payload. Legacy artifacts without this metadata are -rejected with `INDEX_NOT_READY`; rebuild and publish a fresh graph/index version before switching -requests. Existing immutable artifacts are validated rather than silently overwritten. - -Chunk fingerprints now include `sourceUri`; artifacts built with the previous fingerprint must be -rebuilt. All three Java artifact builders use a shared publisher with bounded JVM locks and an OS file -lock per canonical target. Publication checks and rename run under that lock; competing builders reuse -an equivalent winner and reject conflicting content. The `.publish.lock` file stays beside the artifact -so every process locks the same inode. Do not delete it while builders are running. - -`MetadataPublisher.begin` returns an `ImportAttempt` ownership token. `publish` and `fail` require that -token. A failed `begin` does not trigger `fail`, and a competing ingestion cannot fail the owner. -Custom publishers must implement the token-aware lifecycle. - -## HTTP access and keyword index lifecycle - -The server binds to `127.0.0.1` by default. Non-loopback binding requires both -`retrieval.remote-access-enabled=true` and a nonblank `retrieval.api-token`; invalid combinations fail -configuration initialization. Environment equivalents are `GEAFLOW_SERVER_HOST`, -`GEAFLOW_REMOTE_ACCESS_ENABLED`, and `GEAFLOW_RETRIEVAL_API_TOKEN`. Remote mode requires -`Authorization: Bearer ` on retrieval and the legacy `/graph/*` and `/query/*` data APIs. -Missing credentials return `401`; invalid credentials return `403`. `/health` remains available. -The single token grants access to the configured server; per-graph authorization is future work. - -Responses use the same configured limits as request validation, including a raised `max-top-k`. -Keyword retrieval builds one complete Lucene store on the first request and reuses it. Candidate -budgets cap search results evaluated rather than the vertices included in the cached index. Graph -updates invalidate the cache; server replacement and shutdown close it. Cold construction and lock -acquisition check the request deadline. Search store closure is idempotent. - -## Python ingestion - -Multi-file ZIP sources select the requested train/dev/test entry by case-insensitive basename; -missing or ambiguous splits fail. Single archives with a generic data filename remain supported. -Artifact identity includes dataset release and the effective chunk policy (size, overlap, and token -estimate settings). The policy fingerprint also participates in chunk IDs, even for short documents. - -Ingestion normalizes one batch at a time, writes canonical `documents.jsonl` and `chunks.jsonl`, and -merges graph records and provenance in a temporary SQLite database. Index builder callbacks receive a -repeatable `JsonlRecords` iterable supporting `len`; callbacks must iterate rather than index a list. -Graph export reads sorted records from disk. Temporary database and staging files are removed on -success or failure; a failed attempt retains only a separate `*.FAILED.json` diagnostic. Canonical -JSONL files remain in the published version for reproducibility. Memory depends on batch size, an -individual document's chunks, and a single exported graph record's provenance, rather than simultaneous -materialization of the entire corpus. This does not bound allocations made by external builder callbacks. - -Run the retrieval and HTTP regression suite with JDK 11, which is compatible with the repository's -JaCoCo 0.8.8: - -```bash -mvn -pl geaflow-ai -am clean test -Drat.skip=true \ - -Dtest='org.apache.geaflow.ai.retrieval.**.*Test,MemoryServerTest,GeaFlowMemoryServerTest,GraphMemoryKeywordSearchTest,SearchStoreTest' \ - -Dsurefire.failIfNoSpecifiedTests=false -python -m pytest tools/graphrag/tests -q -``` - -The Maven lifecycle also runs Checkstyle and Apache RAT. New retrieval production sources have no -blanket Checkstyle suppressions. The suite covers bounded scoring, partial-channel failures, -deadlines using a controlled clock, graph limits, version mismatches, JSON trace round trips, -concurrent request isolation, and rejection of core modes at the HTTP boundary. -The PR #876 regressions also cover concurrent import ownership, same-target publication for all three -artifacts, competing JVMs, URI citations, request IDs, HTTP authentication and custom limits, cached -keyword retrieval, repeated close, ZIP splits, release/policy identities, and a 10,000-document spill -fixture whose traced Python heap stays below 12 MiB. This fixture uses deterministic builder callbacks; -it is not a full HotpotQA/2Wiki or external embedding benchmark. diff --git a/geaflow-ai/README-week1-retrieval.md b/geaflow-ai/README-week1-retrieval.md deleted file mode 100644 index 586efddb0..000000000 --- a/geaflow-ai/README-week1-retrieval.md +++ /dev/null @@ -1,44 +0,0 @@ -# Week 1 Retrieval API - -Build the memory server; run its `GeaFlowMemoryServer` main class with the project runtime classpath: - -```bash -GEAFLOW_SERVER_PORT=8080 mvn -pl geaflow-ai -DskipTests package -``` - -Create a small keyword graph and load one vertex: - -```bash -curl -sS -X POST http://localhost:8080/graph/create \ - -H 'Content-Type: application/json' \ - -d '{"graphName":"Confucius","vertexSchemaList":[],"edgeSchemaList":[]}' - -curl -sS -X POST 'http://localhost:8080/graph/addEntitySchema?graphName=Confucius' \ - -H 'Content-Type: application/json' \ - -d '{"label":"chunk","idField":"id","fields":["text"]}' - -curl -sS -X POST 'http://localhost:8080/graph/insertEntity?graphName=Confucius' \ - -H 'Content-Type: application/json' \ - -d '{"label":"chunk","id":"confucius-1","values":["Confucius taught ethics."]}' -``` - -Run bounded keyword retrieval without creating a session: - -```bash -curl -sS -X POST http://localhost:8080/api/v1/retrievals \ - -H 'Content-Type: application/json' \ - -H 'X-Request-Id: curl-example-1' \ - -d '{"graphName":"Confucius","query":"Confucius","mode":"KEYWORD",\ -"budget":{"topK":10,"timeoutMs":3000,"maxCandidates":100,"tokenBudget":4096}}' -``` - -Health and readiness are separate checks: - -```bash -curl -i http://localhost:8080/health -curl -i http://localhost:8080/ready -curl -sS http://localhost:8080/metrics/retrieval -``` - -The HTTP endpoint supports `KEYWORD` only. Offline BM25, vector, graph-only, and Hybrid retrieval -are available through the [session-free core facade](README-hybrid-recall.md).