diff --git a/VectorGraphRag/pom.xml b/VectorGraphRag/pom.xml
new file mode 100644
index 0000000..c28d0b1
--- /dev/null
+++ b/VectorGraphRag/pom.xml
@@ -0,0 +1,103 @@
+
+
+ 4.0.0
+
+ com.zjkl
+ vector-graph-rag
+ 0.1.0
+ jar
+ Vector Graph RAG
+ Graph RAG using pure vector search with Milvus
+
+
+ 21
+ 21
+ 21
+ UTF-8
+ 3.5.11
+ 2.6.18
+ 2.13.2
+ 3.2.0
+ 1.18.44
+ 2.0.17
+
+
+
+
+
+ org.springframework.boot
+ spring-boot-dependencies
+ ${spring-boot.version}
+ pom
+ import
+
+
+
+
+
+
+ org.springframework.boot
+ spring-boot-starter-web
+
+
+
+ org.projectlombok
+ lombok
+ ${lombok.version}
+ provided
+
+
+
+ io.milvus
+ milvus-sdk-java
+ ${milvus-sdk.version}
+
+
+
+ com.google.code.gson
+ gson
+ ${gson.version}
+
+
+
+ com.github.ben-manes.caffeine
+ caffeine
+ ${caffeine.version}
+
+
+
+ org.slf4j
+ slf4j-api
+ ${slf4j.version}
+
+
+
+ org.springframework.boot
+ spring-boot-starter-test
+ test
+
+
+
+
+
+
+ org.apache.maven.plugins
+ maven-compiler-plugin
+ 3.14.0
+
+ 21
+ 21
+
+
+ org.projectlombok
+ lombok
+ ${lombok.version}
+
+
+
+
+
+
+
diff --git a/VectorGraphRag/src/main/java/com/zjkl/vectorgraphrag/VectorGraphRAG.java b/VectorGraphRag/src/main/java/com/zjkl/vectorgraphrag/VectorGraphRAG.java
new file mode 100644
index 0000000..8cba48a
--- /dev/null
+++ b/VectorGraphRag/src/main/java/com/zjkl/vectorgraphrag/VectorGraphRAG.java
@@ -0,0 +1,451 @@
+package com.zjkl.vectorgraphrag;
+
+import com.zjkl.vectorgraphrag.config.VectorGraphRagSettings;
+import com.zjkl.vectorgraphrag.graph.Graph;
+import com.zjkl.vectorgraphrag.graph.GraphBuilder;
+import com.zjkl.vectorgraphrag.graph.GraphRetriever;
+import com.zjkl.vectorgraphrag.graph.SubGraph;
+import com.zjkl.vectorgraphrag.llm.AnswerGenerator;
+import com.zjkl.vectorgraphrag.llm.EntityExtractor;
+import com.zjkl.vectorgraphrag.llm.LLMCache;
+import com.zjkl.vectorgraphrag.llm.LLMReranker;
+import com.zjkl.vectorgraphrag.llm.OpenAiClient;
+import com.zjkl.vectorgraphrag.llm.TripletExtractor;
+import com.zjkl.vectorgraphrag.model.*;
+import com.zjkl.vectorgraphrag.storage.EmbeddingClient;
+import com.zjkl.vectorgraphrag.storage.MilvusStore;
+import lombok.Getter;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.*;
+import java.util.stream.Collectors;
+
+/**
+ * Vector Graph RAG - Graph RAG using pure vector search with Milvus.
+ *
+ * Main entry point for building and querying a Graph RAG system.
+ * Implements the full pipeline: triplet extraction → graph building →
+ * embedding → Milvus indexing → multi-way retrieval → subgraph expansion →
+ * LLM reranking → answer generation.
+ *
+ * Example:
+ * {@code
+ * VectorGraphRagSettings settings = VectorGraphRagSettings.builder()
+ * .openaiApiKey(System.getenv("OPENAI_API_KEY"))
+ * .milvusUri("./my_graph.db")
+ * .build();
+ *
+ * VectorGraphRAG rag = new VectorGraphRAG(settings);
+ *
+ * rag.addTexts(List.of(
+ * "Albert Einstein developed the theory of relativity.",
+ * "The theory of relativity revolutionized physics."
+ * ));
+ *
+ * QueryResult result = rag.query("What did Einstein develop?");
+ * System.out.println(result.getAnswer());
+ * }
+ */
+@Slf4j
+public class VectorGraphRAG {
+
+ private final VectorGraphRagSettings settings;
+
+ private final EmbeddingClient embeddingClient;
+ private final MilvusStore store;
+ private final GraphBuilder graphBuilder;
+ private final TripletExtractor tripletExtractor;
+ private final EntityExtractor entityExtractor;
+ private final LLMReranker reranker;
+ private final AnswerGenerator answerGenerator;
+ private final LLMCache llmCache;
+ private final OpenAiClient openAiClient;
+ private final Graph graph;
+
+ @Getter
+ private GraphRetriever retriever;
+
+ @Getter
+ private ExtractionResult extractionResult;
+
+ /**
+ * Create VectorGraphRAG with custom settings.
+ */
+ public VectorGraphRAG(VectorGraphRagSettings settings) {
+ this.settings = settings;
+ this.llmCache = settings.isUseLlmCache() ? new LLMCache() : null;
+ this.openAiClient = new OpenAiClient(settings, llmCache);
+ this.embeddingClient = new EmbeddingClient(settings, openAiClient);
+ this.store = new MilvusStore(settings, embeddingClient);
+ this.graphBuilder = new GraphBuilder(settings);
+ this.tripletExtractor = new TripletExtractor(settings, openAiClient);
+ this.entityExtractor = new EntityExtractor(settings, openAiClient);
+ this.reranker = new LLMReranker(settings, openAiClient);
+ this.answerGenerator = new AnswerGenerator(settings, openAiClient);
+ this.graph = new Graph(settings, store, embeddingClient);
+
+ // Create collections if they don't exist
+ store.createCollections(false);
+ }
+
+ /**
+ * Create VectorGraphRAG with default settings (reads OPENAI_API_KEY from environment).
+ */
+ public static VectorGraphRAG createDefault() {
+ VectorGraphRagSettings settings = VectorGraphRagSettings.builder()
+ .openaiApiKey(System.getenv().getOrDefault("OPENAI_API_KEY", ""))
+ .build();
+ return new VectorGraphRAG(settings);
+ }
+
+ // ==================== Indexing ====================
+
+ /**
+ * Add text strings to the knowledge base.
+ */
+ public ExtractionResult addTexts(List texts) {
+ return addTexts(texts, null, true);
+ }
+
+ /**
+ * Add text strings with extraction control.
+ */
+ public ExtractionResult addTexts(List texts, List ids, boolean extractTriplets) {
+ List documents = new ArrayList<>();
+ for (int i = 0; i < texts.size(); i++) {
+ String docId = (ids != null && i < ids.size()) ? ids.get(i) : UUID.randomUUID().toString();
+ documents.add(Document.builder()
+ .id(docId)
+ .text(texts.get(i))
+ .build());
+ }
+ return addDocuments(documents, extractTriplets, true);
+ }
+
+ /**
+ * Add Document objects to the knowledge base.
+ * Full indexing pipeline: extract triplets → build graph → embed → index in Milvus.
+ */
+ public ExtractionResult addDocuments(List documents, boolean extractTriplets, boolean showProgress) {
+ // Ensure all documents have IDs
+ for (Document doc : documents) {
+ if (doc.getId() == null) {
+ doc.setId(UUID.randomUUID().toString());
+ }
+ }
+
+ // Extract triplets
+ if (extractTriplets) {
+ documents = tripletExtractor.extractFromDocuments(documents, showProgress);
+ }
+
+ // Build graph
+ this.extractionResult = graphBuilder.buildFromDocuments(documents);
+
+ // Generate embeddings
+ if (showProgress) log.info("Generating embeddings...");
+
+ List entityTexts = graphBuilder.getEntityTexts();
+ List relationTexts = graphBuilder.getRelationTexts();
+ List passageTexts = graphBuilder.getPassageTexts();
+
+ List> entityEmbeddings = entityTexts.isEmpty() ? List.of()
+ : embeddingClient.embedBatch(entityTexts, showProgress);
+ List> relationEmbeddings = relationTexts.isEmpty() ? List.of()
+ : embeddingClient.embedBatch(relationTexts, showProgress);
+ List> passageEmbeddings = passageTexts.isEmpty() ? List.of()
+ : embeddingClient.embedBatch(passageTexts, showProgress);
+
+ // Build metadata for adjacency
+ List> entityMetadatas = new ArrayList<>();
+ for (String eid : graphBuilder.getEntityIds()) {
+ Map meta = new HashMap<>();
+ meta.put("relation_ids", graphBuilder.getEntityToRelationIds()
+ .getOrDefault(eid, List.of()));
+ meta.put("passage_ids", graphBuilder.getEntityToPassageIds()
+ .getOrDefault(eid, List.of()));
+ entityMetadatas.add(meta);
+ }
+
+ List> relationMetadatas = new ArrayList<>();
+ for (String rid : graphBuilder.getRelationIds()) {
+ Map meta = new HashMap<>();
+ meta.put("entity_ids", graphBuilder.getRelationToEntityIds()
+ .getOrDefault(rid, List.of()));
+ meta.put("passage_ids", graphBuilder.getRelationToPassageIds()
+ .getOrDefault(rid, List.of()));
+
+ Triplet triplet = graphBuilder.getRelationIdToTriplet().get(rid);
+ if (triplet != null) {
+ meta.put("subject", triplet.getSubject());
+ meta.put("predicate", triplet.getPredicate());
+ meta.put("object", triplet.getObject());
+ }
+ relationMetadatas.add(meta);
+ }
+
+ List> passageMetadatas = new ArrayList<>();
+ for (String pid : graphBuilder.getPassageIds()) {
+ Map meta = new HashMap<>();
+ meta.put("entity_ids", graphBuilder.getPassageToEntityIds()
+ .getOrDefault(pid, List.of()));
+ meta.put("relation_ids", graphBuilder.getPassageToRelationIds()
+ .getOrDefault(pid, List.of()));
+ passageMetadatas.add(meta);
+ }
+
+ // Drop and recreate for idempotent indexing (matching Python reference behavior)
+ store.dropCollections();
+ store.createCollections(true);
+ if (showProgress) log.info("Inserting into Milvus...");
+
+ store.insertEntities(entityTexts, graphBuilder.getEntityIds(),
+ entityEmbeddings, entityMetadatas, showProgress);
+ store.insertRelations(relationTexts, graphBuilder.getRelationIds(),
+ relationEmbeddings, relationMetadatas, showProgress);
+ store.insertPassages(passageTexts, graphBuilder.getPassageIds(),
+ passageEmbeddings, passageMetadatas, showProgress);
+
+ // Reset retriever
+ this.retriever = null;
+
+ return extractionResult;
+ }
+
+ /**
+ * Add documents with pre-extracted triplets (skip LLM extraction).
+ */
+ public ExtractionResult addDocumentsWithTriplets(List> documents) {
+ List docs = new ArrayList<>();
+ for (Map docData : documents) {
+ String passage = (String) docData.getOrDefault("passage", docData.get("text"));
+ if (passage == null) {
+ throw new IllegalArgumentException("Each document must include 'passage' or 'text'");
+ }
+
+ String docId = (String) docData.getOrDefault("id", UUID.randomUUID().toString());
+
+ @SuppressWarnings("unchecked")
+ List> rawTriplets = (List>) docData.get("triplets");
+ List triplets = new ArrayList<>();
+ if (rawTriplets != null) {
+ for (List raw : rawTriplets) {
+ if (raw.size() >= 3) {
+ triplets.add(new Triplet(raw.get(0), raw.get(1), raw.get(2)));
+ }
+ }
+ }
+
+ Document doc = Document.builder()
+ .id(docId)
+ .text(passage)
+ .triplets(triplets)
+ .build();
+ // Store triplets in metadata for graph builder
+ List> metadataTriplets = triplets.stream()
+ .map(t -> List.of(t.getSubject(), t.getPredicate(), t.getObject()))
+ .collect(Collectors.toList());
+ doc.getMetadata().put("triplets", metadataTriplets);
+
+ docs.add(doc);
+ }
+
+ return addDocuments(docs, false, true);
+ }
+
+ // ==================== Query ====================
+
+ /**
+ * Full Graph RAG query pipeline.
+ */
+ public QueryResult query(String question) {
+ return query(question, true, null, null, null, null, null);
+ }
+
+ /**
+ * Query with custom parameters.
+ */
+ public QueryResult query(String question, boolean useReranking,
+ Integer entityTopK, Integer relationTopK,
+ Float entitySimilarityThreshold, Float relationSimilarityThreshold,
+ String filter) {
+ GraphRetriever ret = getOrCreateRetriever();
+
+ RetrievalResult retrievalResult = ret.retrieve(question, entityTopK, relationTopK,
+ entitySimilarityThreshold, relationSimilarityThreshold,
+ null, filter);
+
+ // Build retrieval detail
+ RetrievalDetail retrievalDetail = RetrievalDetail.builder()
+ .entityIds(retrievalResult.getEntityIds())
+ .entityTexts(retrievalResult.getEntityTexts())
+ .entityScores(retrievalResult.getEntityScores())
+ .relationIds(retrievalResult.getRelationIds())
+ .relationTexts(retrievalResult.getRelationTexts())
+ .relationScores(retrievalResult.getRelationScores())
+ .build();
+
+ // Get candidate relations from subgraph
+ List candidateIds = retrievalResult.getExpandedRelationIds();
+ List candidateTexts = retrievalResult.getExpandedRelationTexts();
+
+ // Rerank
+ RerankResult rerankResult = null;
+ List rerankedIds;
+ List rerankedTexts;
+
+ if (useReranking && !candidateIds.isEmpty()) {
+ Map.Entry, List> result = reranker.rerank(
+ question, candidateIds, candidateTexts);
+ rerankedIds = result.getKey();
+ rerankedTexts = result.getValue();
+ rerankResult = RerankResult.builder()
+ .selectedRelationIds(rerankedIds)
+ .selectedRelationTexts(rerankedTexts)
+ .build();
+ } else {
+ rerankedIds = candidateIds.size() > settings.getFinalTopK()
+ ? candidateIds.subList(0, settings.getFinalTopK())
+ : candidateIds;
+ rerankedTexts = candidateTexts.size() > settings.getFinalTopK()
+ ? candidateTexts.subList(0, settings.getFinalTopK())
+ : candidateTexts;
+ }
+
+ // Get passages from reranked relations
+ List finalPassages = getPassagesFromRelations(rerankedIds, filter);
+
+ // Hybrid fallback: if graph retrieval yields fewer passages than needed,
+ // supplement with naive vector search on passages directly
+ if (finalPassages.size() < settings.getFinalTopK()) {
+ int needed = settings.getFinalTopK() - finalPassages.size();
+ Set existing = new HashSet<>(finalPassages);
+ List naivePassages = ret.retrievePassagesNaive(question, needed, filter);
+ for (String p : naivePassages) {
+ if (existing.add(p)) {
+ finalPassages.add(p);
+ if (finalPassages.size() >= settings.getFinalTopK()) break;
+ }
+ }
+ }
+
+ if (finalPassages.size() > settings.getFinalTopK()) {
+ finalPassages = finalPassages.subList(0, settings.getFinalTopK());
+ }
+
+ // Generate answer
+ String answer = answerGenerator.generate(question, finalPassages);
+
+ return QueryResult.builder()
+ .query(question)
+ .answer(answer)
+ .queryEntities(retrievalResult.getQueryEntities())
+ .retrievedPassages(finalPassages)
+ .retrievedRelations(retrievalResult.getRelationTexts())
+ .expandedRelations(candidateTexts)
+ .rerankedRelations(rerankedTexts)
+ .subgraph(retrievalResult.getSubgraph())
+ .passages(finalPassages)
+ .retrievalDetail(retrievalDetail)
+ .rerankResult(rerankResult)
+ .evictionResult(EvictionResult.builder()
+ .occurred(retrievalResult.isEvictionOccurred())
+ .beforeCount(retrievalResult.getEvictionBeforeCount())
+ .afterCount(retrievalResult.getEvictionAfterCount())
+ .build())
+ .build();
+ }
+
+ /**
+ * Simple query that returns just the answer string.
+ */
+ public String querySimple(String question) {
+ return query(question).getAnswer();
+ }
+
+ /**
+ * Naive RAG query (direct passage retrieval, no graph).
+ */
+ public QueryResult queryNaive(String question, String filter) {
+ GraphRetriever ret = getOrCreateRetriever();
+ List passages = ret.retrievePassagesNaive(question, settings.getFinalTopK(), filter);
+ String answer = answerGenerator.generate(question, passages);
+
+ return QueryResult.builder()
+ .query(question)
+ .answer(answer)
+ .retrievedPassages(passages)
+ .build();
+ }
+
+ // ==================== Utility ====================
+
+ /**
+ * Get passages associated with given relation IDs.
+ */
+ private List getPassagesFromRelations(List relationIds, String filter) {
+ if (relationIds == null || relationIds.isEmpty()) return List.of();
+
+ List> relationData = store.getRelationsByIds(relationIds);
+ Set passageIds = new LinkedHashSet<>();
+ for (Map rel : relationData) {
+ @SuppressWarnings("unchecked")
+ List pids = (List) rel.get("passage_ids");
+ if (pids != null) passageIds.addAll(pids);
+ }
+
+ if (passageIds.isEmpty()) return List.of();
+
+ List> passageData = store.getPassagesByIds(new ArrayList<>(passageIds), filter);
+ Map idToText = new LinkedHashMap<>();
+ for (Map p : passageData) {
+ idToText.put(p.get("id").toString(), p.get("text").toString());
+ }
+
+ return passageIds.stream()
+ .map(idToText::get)
+ .filter(Objects::nonNull)
+ .collect(Collectors.toList());
+ }
+
+ /**
+ * Get knowledge base statistics.
+ */
+ public Map getStats() {
+ if (extractionResult == null) {
+ Map stats = new HashMap<>();
+ stats.put("entities", 0);
+ stats.put("relations", 0);
+ stats.put("passages", 0);
+ return stats;
+ }
+ Map stats = new HashMap<>();
+ stats.put("entities", extractionResult.getEntities().size());
+ stats.put("relations", extractionResult.getRelations().size());
+ stats.put("passages", extractionResult.getDocuments().size());
+ return stats;
+ }
+
+ /**
+ * Reset the knowledge base, removing all data.
+ */
+ public void reset() {
+ store.dropCollections();
+ store.createCollections(true);
+ this.extractionResult = null;
+ this.retriever = null;
+ }
+
+ private GraphRetriever getOrCreateRetriever() {
+ if (retriever == null) {
+ retriever = new GraphRetriever(settings, store, embeddingClient, entityExtractor);
+ }
+ return retriever;
+ }
+
+ // ==================== Accessors ====================
+
+ public Graph getGraph() { return graph; }
+ public MilvusStore getStore() { return store; }
+ public EmbeddingClient getEmbeddingClient() { return embeddingClient; }
+}
diff --git a/VectorGraphRag/src/main/java/com/zjkl/vectorgraphrag/config/VectorGraphRagSettings.java b/VectorGraphRag/src/main/java/com/zjkl/vectorgraphrag/config/VectorGraphRagSettings.java
new file mode 100644
index 0000000..170a334
--- /dev/null
+++ b/VectorGraphRag/src/main/java/com/zjkl/vectorgraphrag/config/VectorGraphRagSettings.java
@@ -0,0 +1,110 @@
+package com.zjkl.vectorgraphrag.config;
+
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.stereotype.Component;
+
+import java.util.Map;
+
+@Data
+@Builder
+@NoArgsConstructor
+@AllArgsConstructor
+@Component
+public class VectorGraphRagSettings {
+ // OpenAI Settings
+ @Builder.Default
+ private String openaiApiKey = System.getenv().getOrDefault("OPENAI_API_KEY", "");
+
+ @Builder.Default
+ private String openaiBaseUrl = "";
+
+ // Model Settings
+ @Builder.Default
+ private String llmModel = "gpt-4o-mini";
+
+ @Builder.Default
+ private String embeddingModel = "text-embedding-3-large";
+
+ @Builder.Default
+ private int embeddingDimension = 3072;
+
+ // Milvus Settings
+ @Builder.Default
+ private String milvusUri = "./vector_graph_rag.db";
+
+ @Builder.Default
+ private String milvusToken = "";
+
+ @Builder.Default
+ private String milvusDb = "";
+
+ // Milvus Index Settings
+ @Builder.Default
+ private String milvusIndexType = "AUTOINDEX";
+
+ @Builder.Default
+ private String milvusMetricType = "IP";
+
+ private Map milvusIndexParams;
+
+ @Builder.Default
+ private String milvusConsistencyLevel = "Bounded";
+
+ // Collection prefix
+ @Builder.Default
+ private String collectionPrefix = "";
+
+ // Collection Names
+ @Builder.Default
+ private String entityCollection = "vgrag_entities";
+
+ @Builder.Default
+ private String relationCollection = "vgrag_relations";
+
+ @Builder.Default
+ private String passageCollection = "vgrag_passages";
+
+ // Retrieval Settings
+ @Builder.Default
+ private int entityTopK = 20;
+
+ @Builder.Default
+ private int relationTopK = 20;
+
+ @Builder.Default
+ private float entitySimilarityThreshold = 0.9f;
+
+ @Builder.Default
+ private float relationSimilarityThreshold = -1.0f;
+
+ @Builder.Default
+ private int expansionDegree = 1;
+
+ @Builder.Default
+ private int relationNumberThreshold = 1000;
+
+ @Builder.Default
+ private int finalTopK = 3;
+
+ // LLM Settings
+ @Builder.Default
+ private double llmTemperature = 0.0;
+
+ @Builder.Default
+ private int llmMaxRetries = 3;
+
+ @Builder.Default
+ private boolean useLlmCache = true;
+
+ // Processing Settings
+ @Builder.Default
+ private int batchSize = 32;
+
+ public boolean hasCollectionPrefix() {
+ return collectionPrefix != null && !collectionPrefix.isEmpty();
+ }
+}
diff --git a/VectorGraphRag/src/main/java/com/zjkl/vectorgraphrag/graph/Graph.java b/VectorGraphRag/src/main/java/com/zjkl/vectorgraphrag/graph/Graph.java
new file mode 100644
index 0000000..477b63c
--- /dev/null
+++ b/VectorGraphRag/src/main/java/com/zjkl/vectorgraphrag/graph/Graph.java
@@ -0,0 +1,258 @@
+package com.zjkl.vectorgraphrag.graph;
+
+import com.zjkl.vectorgraphrag.config.VectorGraphRagSettings;
+import com.zjkl.vectorgraphrag.model.*;
+import com.zjkl.vectorgraphrag.storage.EmbeddingClient;
+import com.zjkl.vectorgraphrag.storage.MilvusStore;
+import lombok.extern.slf4j.Slf4j;
+
+import java.util.*;
+import java.util.stream.Collectors;
+
+import static com.zjkl.vectorgraphrag.graph.GraphBuilder.normalizePhrase;
+
+/**
+ * User-facing graph operations interface.
+ * Encapsulates MilvusStore with CRUD operations and automatic entity/relation linking.
+ */
+@Slf4j
+public class Graph {
+
+ private final VectorGraphRagSettings settings;
+ private final EmbeddingClient embeddingClient;
+ private final MilvusStore store;
+
+ private final Map entityNameToId = new HashMap<>();
+ private final Map relationTextToId = new HashMap<>();
+
+ public Graph(VectorGraphRagSettings settings, MilvusStore store, EmbeddingClient embeddingClient) {
+ this.settings = settings;
+ this.store = store;
+ this.embeddingClient = embeddingClient;
+ }
+
+ // ==================== Passage CRUD ====================
+
+ public String createPassage(String text) {
+ return createPassage(text, null, null);
+ }
+
+ public String createPassage(String text, String id) {
+ return createPassage(text, id, null);
+ }
+
+ public String createPassage(String text, String id, List triplets) {
+ String passageId = id != null ? id : UUID.randomUUID().toString();
+ List embedding = embeddingClient.embed(text);
+
+ List entityIds = new ArrayList<>();
+ List relationIds = new ArrayList<>();
+
+ if (triplets != null) {
+ for (Triplet triplet : triplets) {
+ String relationId = createRelation(triplet.getSubject(), triplet.getPredicate(),
+ triplet.getObject(), null, List.of(passageId));
+ relationIds.add(relationId);
+
+ String subjectId = entityNameToId.get(normalizePhrase(triplet.getSubject()));
+ String objectId = entityNameToId.get(normalizePhrase(triplet.getObject()));
+ if (subjectId != null && !entityIds.contains(subjectId)) entityIds.add(subjectId);
+ if (objectId != null && !entityIds.contains(objectId)) entityIds.add(objectId);
+ }
+ }
+
+ Map metadata = new HashMap<>();
+ if (!entityIds.isEmpty()) metadata.put("entity_ids", entityIds);
+ if (!relationIds.isEmpty()) metadata.put("relation_ids", relationIds);
+
+ List> metadatas = metadata.isEmpty() ? null : List.of(metadata);
+ store.insertPassages(List.of(text), List.of(passageId), List.of(embedding), metadatas, false);
+
+ return passageId;
+ }
+
+ public Passage getPassage(String passageId) {
+ List> results = store.getPassagesByIds(List.of(passageId));
+ if (results.isEmpty()) return null;
+
+ Map data = results.get(0);
+ return Passage.builder()
+ .id(safeGet(data, "id"))
+ .text(safeGet(data, "text"))
+ .entityIds(safeGetList(data, "entity_ids"))
+ .relationIds(safeGetList(data, "relation_ids"))
+ .build();
+ }
+
+ public List searchPassages(String query, int topK) {
+ List queryEmbedding = embeddingClient.embed(query);
+ List> results = store.searchPassages(queryEmbedding, topK, null);
+
+ return results.stream().map(r -> {
+ @SuppressWarnings("unchecked")
+ Map entity = (Map) r.get("entity");
+ return Passage.builder()
+ .id(safeGet(entity, "id"))
+ .text(safeGet(entity, "text"))
+ .entityIds(safeGetList(entity, "entity_ids"))
+ .relationIds(safeGetList(entity, "relation_ids"))
+ .build();
+ }).collect(Collectors.toList());
+ }
+
+ public boolean updatePassage(String passageId, String text, List entityIds, List relationIds) {
+ return store.upsertPassage(passageId, text, null, entityIds, relationIds);
+ }
+
+ public boolean deletePassage(String passageId) {
+ List> passages = store.getPassagesByIds(List.of(passageId));
+ if (passages.isEmpty()) return false;
+
+ Map data = passages.get(0);
+ List entityIds = safeGetList(data, "entity_ids");
+ List relationIds = safeGetList(data, "relation_ids");
+
+ // Cascade: remove passage reference from entities
+ for (String eid : entityIds) {
+ List> entities = store.getEntitiesByIds(List.of(eid));
+ if (!entities.isEmpty()) {
+ List pids = safeGetList(entities.get(0), "passage_ids");
+ pids.remove(passageId);
+ store.upsertEntity(eid, null, null, null, pids);
+ }
+ }
+
+ // Cascade: remove passage reference from relations
+ for (String rid : relationIds) {
+ List> rels = store.getRelationsByIds(List.of(rid));
+ if (!rels.isEmpty()) {
+ List pids = safeGetList(rels.get(0), "passage_ids");
+ pids.remove(passageId);
+ store.upsertRelation(rid, null, null, null, pids, null, null, null);
+ }
+ }
+
+ return store.deletePassage(passageId);
+ }
+
+ // ==================== Private: Entity & Relation CRUD ====================
+
+ private String createRelation(String subject, String predicate, String object,
+ String id, List passageIds) {
+ String normSubject = normalizePhrase(subject);
+ String normPredicate = normalizePhrase(predicate);
+ String normObject = normalizePhrase(object);
+ String relationText = normSubject + " " + normPredicate + " " + normObject;
+
+ if (relationTextToId.containsKey(relationText)) {
+ String existingId = relationTextToId.get(relationText);
+ if (passageIds != null && !passageIds.isEmpty()) {
+ List> existing = store.getRelationsByIds(List.of(existingId));
+ if (!existing.isEmpty()) {
+ List currentPids = safeGetList(existing.get(0), "passage_ids");
+ Set merged = new LinkedHashSet<>(currentPids);
+ merged.addAll(passageIds);
+ store.upsertRelation(existingId, null, null, null, new ArrayList<>(merged),
+ null, null, null);
+ }
+ }
+ return existingId;
+ }
+
+ String relationId = id != null ? id : UUID.randomUUID().toString();
+
+ String subjectId = createEntity(subject, List.of(relationId), passageIds);
+ String objectId = createEntity(object, List.of(relationId), passageIds);
+
+ List embedding = embeddingClient.embed(relationText);
+
+ Map metadata = new HashMap<>();
+ metadata.put("entity_ids", List.of(subjectId, objectId));
+ metadata.put("subject", normSubject);
+ metadata.put("predicate", normPredicate);
+ metadata.put("object", normObject);
+ if (passageIds != null && !passageIds.isEmpty()) {
+ metadata.put("passage_ids", passageIds);
+ }
+
+ store.insertRelations(List.of(relationText), List.of(relationId),
+ List.of(embedding), List.of(metadata), false);
+
+ relationTextToId.put(relationText, relationId);
+ return relationId;
+ }
+
+ private String createEntity(String name, List relationIds, List passageIds) {
+ String normalized = normalizePhrase(name);
+ if (entityNameToId.containsKey(normalized)) {
+ String existingId = entityNameToId.get(normalized);
+ if (relationIds != null || passageIds != null) {
+ List> existing = store.getEntitiesByIds(List.of(existingId));
+ if (!existing.isEmpty()) {
+ List currentRels = safeGetList(existing.get(0), "relation_ids");
+ List currentPids = safeGetList(existing.get(0), "passage_ids");
+ Set mergedRels = new LinkedHashSet<>(currentRels);
+ if (relationIds != null) mergedRels.addAll(relationIds);
+ Set mergedPids = new LinkedHashSet<>(currentPids);
+ if (passageIds != null) mergedPids.addAll(passageIds);
+ store.upsertEntity(existingId, null, null,
+ new ArrayList<>(mergedRels), new ArrayList<>(mergedPids));
+ }
+ }
+ return existingId;
+ }
+
+ String entityId = UUID.randomUUID().toString();
+ List embedding = embeddingClient.embed(normalized);
+
+ Map metadata = new HashMap<>();
+ if (relationIds != null && !relationIds.isEmpty()) metadata.put("relation_ids", relationIds);
+ if (passageIds != null && !passageIds.isEmpty()) metadata.put("passage_ids", passageIds);
+
+ List> metadatas = metadata.isEmpty() ? null : List.of(metadata);
+ store.insertEntities(List.of(normalized), List.of(entityId),
+ List.of(embedding), metadatas, false);
+
+ entityNameToId.put(normalized, entityId);
+ return entityId;
+ }
+
+ // ==================== SubGraph Creation ====================
+
+ public SubGraph createSubGraph() {
+ return new SubGraph(store);
+ }
+
+ // ==================== Collection Management ====================
+
+ public void createCollections(boolean dropExisting) {
+ store.createCollections(dropExisting);
+ }
+
+ public void dropCollections() {
+ store.dropCollections();
+ entityNameToId.clear();
+ relationTextToId.clear();
+ }
+
+ public void reset() {
+ dropCollections();
+ store.createCollections(true);
+ }
+
+ // ==================== Helpers ====================
+
+ @SuppressWarnings("unchecked")
+ private List safeGetList(Map map, String key) {
+ Object val = map.get(key);
+ if (val instanceof List) {
+ return ((List