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

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import com.orgmemory.graphrag.processing.ResolvedDocumentProcessingProfile;
import com.orgmemory.integrations.graphrag.springai.JtokkitTextTokenizer;
import com.orgmemory.integrations.graphrag.springai.SpringAiTextEmbeddingPort;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.Optional;
Expand Down Expand Up @@ -87,7 +88,10 @@ ProcessedSourceDocument process(
DocumentParseRequest request,
EmbeddingModel embeddingModel) {
ParserSpec parser = parsers.route(request.suffix(), properties.parserId());
long parseStartedAt = System.nanoTime();
var parsed = parser.parser().parse(request);
Duration parseDuration =
Duration.ofNanos(Math.max(0, System.nanoTime() - parseStartedAt));
long minimumChunks = Math.ceilDiv(
(long) tokenizer.count(parsed.document().content()),
properties.chunkSize());
Expand All @@ -106,6 +110,7 @@ ProcessedSourceDocument process(
ChunkerOptions options = options(requestedChunker);
Map<String, String> requestedOptions = resolvedOptions(options);
List<ChunkedText> output;
long chunkStartedAt = System.nanoTime();
try {
output = chunkers.execute(
requestedChunker,
Expand All @@ -124,6 +129,11 @@ ProcessedSourceDocument process(
new ChunkingRequest(parsed.document(), tokenizer, Optional.empty()),
options);
}
// Measured before the limit check so a rejected document is still measured work; the
// fallback path above is inside the same window on purpose, because a semantic chunker
// failing over to the recursive one is time this document actually spent chunking.
Duration chunkDuration =
Duration.ofNanos(Math.max(0, System.nanoTime() - chunkStartedAt));
if (output.size() > properties.maximumChunks()) {
throw new RejectedSourceException(
"CHUNK_LIMIT_EXCEEDED",
Expand All @@ -144,7 +154,12 @@ ProcessedSourceDocument process(
: Optional.empty(),
resolvedOptions,
parsed.document().contentSha256());
return new ProcessedSourceDocument(parsed, output, profile);
return new ProcessedSourceDocument(
parsed,
output,
profile,
parseDuration,
chunkDuration);
}

private ChunkerOptions options(String chunkerId) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,27 @@
import com.orgmemory.graphrag.chunking.ChunkedText;
import com.orgmemory.graphrag.parsing.DocumentParseResult;
import com.orgmemory.graphrag.processing.ResolvedDocumentProcessingProfile;
import java.time.Duration;
import java.util.List;
import java.util.Objects;

/**
* @param parseDuration how long the routed parser took
* @param chunkDuration how long chunking took, including a semantic chunker's failover to the
* recursive one
*
* <p>Both are carried out of the engine rather than timed by the caller, because the caller
* makes one {@code process} call and cannot see where inside it parsing ended. Reporting the
* pair as one duration would hide which half a slow document spent its time in, and the two
* have unrelated causes — a parser handed a large scanned file, against a chunker falling
* back to a different algorithm.
*/
record ProcessedSourceDocument(
DocumentParseResult parseResult,
List<ChunkedText> chunks,
ResolvedDocumentProcessingProfile profile) {
ResolvedDocumentProcessingProfile profile,
Duration parseDuration,
Duration chunkDuration) {

ProcessedSourceDocument {
Objects.requireNonNull(parseResult, "parseResult");
Expand All @@ -20,5 +34,10 @@ record ProcessedSourceDocument(
"Processed document requires at least one chunk");
}
Objects.requireNonNull(profile, "profile");
Objects.requireNonNull(parseDuration, "parseDuration");
Objects.requireNonNull(chunkDuration, "chunkDuration");
if (parseDuration.isNegative() || chunkDuration.isNegative()) {
throw new IllegalArgumentException("durations must not be negative");
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import com.orgmemory.core.knowledge.storage.ObjectKey;
import com.orgmemory.core.knowledge.storage.ObjectStoragePort;
import com.orgmemory.core.permission.AccessGate;
import com.orgmemory.graphrag.observability.GraphRagEventSink;
import com.orgmemory.graphrag.parsing.DocumentParseRequest;
import com.orgmemory.graphrag.parsing.DocumentParseResult;
import java.io.IOException;
Expand All @@ -37,15 +38,18 @@
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.time.Duration;
import java.time.Instant;
import java.util.HexFormat;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.TokenCountBatchingStrategy;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

@Component
Expand All @@ -62,7 +66,9 @@ class SourceIngestionProcessor {
private final AiRouteResolver aiRoutes;
private final SourceProcessingProperties properties;
private final DocumentProcessingEngine processingEngine;
private final GraphRagEventSink events;

@Autowired
SourceIngestionProcessor(
SourceIngestionCoordinator coordinator,
KnowledgeIngestionService ingestion,
Expand All @@ -72,7 +78,33 @@ class SourceIngestionProcessor {
ObjectProvider<EmbeddingModel> embeddingModels,
AiRouteResolver aiRoutes,
SourceProcessingProperties properties,
DocumentProcessingEngine processingEngine) {
DocumentProcessingEngine processingEngine,
ObjectProvider<GraphRagEventSink> eventSinks) {
this(
coordinator,
ingestion,
publications,
embeddingProfiles,
objects,
embeddingModels,
aiRoutes,
properties,
processingEngine,
GraphRagEventSink.failureTolerant(
GraphRagEventSink.composite(eventSinks.orderedStream().toList())));
}

SourceIngestionProcessor(
SourceIngestionCoordinator coordinator,
KnowledgeIngestionService ingestion,
KnowledgeAssetPublicationService publications,
EmbeddingProfileRegistry embeddingProfiles,
ObjectStoragePort objects,
ObjectProvider<EmbeddingModel> embeddingModels,
AiRouteResolver aiRoutes,
SourceProcessingProperties properties,
DocumentProcessingEngine processingEngine,
GraphRagEventSink events) {
this.coordinator = coordinator;
this.ingestion = ingestion;
this.publications = publications;
Expand All @@ -82,13 +114,48 @@ class SourceIngestionProcessor {
this.aiRoutes = aiRoutes;
this.properties = properties;
this.processingEngine = processingEngine;
this.events = Objects.requireNonNull(events, "events");
}

void processNext() {
coordinator.claimNext(properties.workerId(), properties.leaseDuration())
.ifPresent(this::process);
}

/**
* Reports the two ingestion stages that had no producer.
*
* <p>The revision status already moved through {@code PARSING} and {@code CHUNKING}, but a
* status says where a job is, not how long it stayed there — so a document that took four
* minutes to parse and one that took four seconds left the same trace. These are emitted
* from the same {@code jobId} the graph indexing stages use, so one upload reads as one
* operation across both processors.
*/
private void emitStage(
ClaimedSourceRevision claim,
GraphRagEventSink.Stage stage,
Duration duration,
int inputCount,
int outputCount) {
try {
events.emit(new GraphRagEventSink.GraphRagEvent(
claim.jobId(),
claim.organizationId(),
stage,
GraphRagEventSink.Outcome.SUCCEEDED,
duration,
inputCount,
outputCount,
null,
null,
null,
null,
Instant.now()));
} catch (RuntimeException ignoredTelemetryFailure) {
// Telemetry must never become an ingestion availability dependency.
}
}

private void process(ClaimedSourceRevision claim) {
String failureStage = "VALIDATION";
Path temporaryFile = null;
Expand Down Expand Up @@ -125,6 +192,18 @@ private void process(ClaimedSourceRevision claim) {
Optional.empty()),
embeddingModel);
DocumentParseResult parsed = processed.parseResult();
emitStage(
claim,
GraphRagEventSink.Stage.PARSE,
processed.parseDuration(),
1,
parsed.document().blocks().size());
emitStage(
claim,
GraphRagEventSink.Stage.CHUNK,
processed.chunkDuration(),
parsed.document().blocks().size(),
processed.chunks().size());
RawSourceRef raw = registerRawSource(claim, parsed);
NormalizedRecordRef normalized = ingestion.normalize(new NormalizeRawSourceCommand(
claim.organizationId(),
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
package com.orgmemory.worker.ingestion;

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 com.orgmemory.graphrag.parsing.DocumentParseRequest;
import java.nio.charset.StandardCharsets;
Expand Down Expand Up @@ -48,6 +50,62 @@ void rejectsDefiniteChunkOverflowBeforeCallingSemanticEmbedding() {
assertEquals("CHUNK_LIMIT_EXCEEDED", failure.code());
}

/**
* The caller makes one {@code process} call and cannot see where parsing ended, so the two
* measurements have to leave the engine separately or the pair collapses into one number
* that hides which half a slow document spent its time in.
*
* <p>This asserts the counts the two stage events report, not the accuracy of the timing.
* A timing assertion was tried and removed: on a fixture this small, starting the chunk
* clock before parsing still passed every bound the test could state, so it proved nothing
* it claimed to. The stage events themselves are proven end to end in
* {@code SourceIngestionPipelineIntegrationTests}.
*/
@Test
void carriesTheParseAndChunkMeasurementsOutSeparately() {
var engine = new DocumentProcessingEngine(
properties("passthrough", "fixed-token"),
new SpringAiDocumentParser());

ProcessedSourceDocument processed = engine.process(
new DocumentParseRequest(
"notes.txt",
"text/plain",
"one two three four five six seven eight"
.getBytes(StandardCharsets.UTF_8),
Optional.empty()),
new FailingEmbeddingModel());

assertFalse(
processed.parseResult().document().blocks().isEmpty(),
"PARSE reports blocks produced, so an empty list would make it report zero work");
assertFalse(
processed.chunks().isEmpty(),
"CHUNK reports chunks produced");
}

private static SourceProcessingProperties properties(
String parserId,
String chunkerId) {
return new SourceProcessingProperties(
false,
Duration.ofSeconds(1),
"test-worker",
Duration.ofMinutes(1),
"test-pipeline",
parserId,
chunkerId,
"o200k_base",
"normalizer",
"fixture",
"fixture-model",
64,
8,
0,
64,
1);
}

private static final class FailingEmbeddingModel implements EmbeddingModel {

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@
import com.orgmemory.core.knowledge.storage.StoredObject;
import com.orgmemory.core.organization.CurrentActor;
import com.orgmemory.core.permission.KnowledgeClassification;
import com.orgmemory.graphrag.observability.GraphRagEventSink;
import java.io.ByteArrayInputStream;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
Expand Down Expand Up @@ -129,6 +130,9 @@ class SourceIngestionPipelineIntegrationTests {
@Autowired
SourceIngestionProcessor processor;

@Autowired
RecordingGraphRagEventSink graphRagEvents;

@Autowired
JdbcTemplate jdbc;

Expand Down Expand Up @@ -414,6 +418,24 @@ INSERT INTO source_acl_snapshot_seals (
1,
staleAcl.evidence().size(),
"The latest complete, sealed ACL generation remains authoritative when sync health is stale");

// Parsing and chunking had no producer at all, so an upload that spent minutes in either
// was indistinguishable from one that spent milliseconds — the revision status says where
// a job is, never how long it stayed there.
var parse = graphRagEvents.require(GraphRagEventSink.Stage.PARSE);
var chunk = graphRagEvents.require(GraphRagEventSink.Stage.CHUNK);
assertEquals(1, parse.inputCount(), "one source document entered parsing");
assertTrue(parse.outputCount() > 0, "parsing reports the blocks it produced");
assertEquals(
parse.outputCount(),
chunk.inputCount(),
"chunking consumes exactly the blocks parsing produced");
assertTrue(chunk.outputCount() > 0, "chunking reports the chunks it produced");
assertEquals(
parse.operationId(),
chunk.operationId(),
"both stages belong to the one job, so an upload reads as one operation");
assertEquals(ORGANIZATION_ID, parse.organizationId());
}

@Test
Expand Down Expand Up @@ -559,6 +581,28 @@ private static float[] embedding() {
return values;
}

static final class RecordingGraphRagEventSink implements GraphRagEventSink {

private final List<GraphRagEvent> received =
java.util.Collections.synchronizedList(new java.util.ArrayList<>());

@Override
public void emit(GraphRagEvent event) {
received.add(event);
}

GraphRagEvent require(Stage stage) {
synchronized (received) {
return received.stream()
.filter(event -> event.stage() == stage)
.findFirst()
.orElseThrow(() -> new AssertionError(
"no " + stage + " event was emitted; saw "
+ received.stream().map(GraphRagEvent::stage).toList()));
}
}
}

@TestConfiguration(proxyBeanMethods = false)
@ComponentScan(
basePackageClasses = SourceUploadService.class,
Expand All @@ -568,6 +612,16 @@ private static float[] embedding() {
pattern = "com\\.orgmemory\\.core\\.knowledge\\.Source(UploadService|UploadRegistrationService|QueryService)"))
static class UploadTestConfiguration {

/**
* Contributed as an ordinary sink so the processor composes it exactly as it composes
* the real backends. A recorder injected any other way would prove the emit call and
* not the wiring, and the wiring is what was missing.
*/
@Bean
RecordingGraphRagEventSink recordingGraphRagEventSink() {
return new RecordingGraphRagEventSink();
}

@Bean
@Primary
AiRouteResolver testAiRouteResolver() {
Expand Down
Loading