Skip to content

Commit df2e889

Browse files
krishna-ybkrishna-yb
committed
feat: Add YugabyteDB embedding store examples
- Add four comprehensive examples demonstrating YugabyteDB integration: * Basic embedding store with similarity search * Metadata filtering with JSONB storage * PostgreSQL JDBC driver usage (recommended) * YugabyteDB Smart Driver for distributed deployments - Include helper script (run-example.sh) for easy execution - Use Testcontainers for isolated testing environment - Implement proper resource cleanup to avoid thread warnings - Configure containers matching langchain4j-community test standards
1 parent 0df7c86 commit df2e889

File tree

7 files changed

+563
-0
lines changed

7 files changed

+563
-0
lines changed

‎pom.xml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
<module>helidon-examples</module>
5454
<module>payara-micro-example</module>
5555
<module>watsonx-ai-examples</module>
56+
<module>yugabytedb-example</module>
5657
</modules>
5758

5859
</project>

‎yugabytedb-example/pom.xml‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<project xmlns="http://maven.apache.org/POM/4.0.0"
3+
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
4+
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
5+
<modelVersion>4.0.0</modelVersion>
6+
7+
<groupId>dev.langchain4j</groupId>
8+
<artifactId>yugabytedb-example</artifactId>
9+
<version>1.8.0-beta15-SNAPSHOT</version>
10+
11+
<properties>
12+
<maven.compiler.source>17</maven.compiler.source>
13+
<maven.compiler.target>17</maven.compiler.target>
14+
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
15+
</properties>
16+
17+
<dependencies>
18+
19+
<dependency>
20+
<groupId>dev.langchain4j</groupId>
21+
<artifactId>langchain4j-community-yugabytedb</artifactId>
22+
<version>1.8.0-beta15-SNAPSHOT</version>
23+
</dependency>
24+
25+
<dependency>
26+
<groupId>org.testcontainers</groupId>
27+
<artifactId>testcontainers</artifactId>
28+
<version>1.19.7</version>
29+
</dependency>
30+
31+
<dependency>
32+
<groupId>dev.langchain4j</groupId>
33+
<artifactId>langchain4j-embeddings-all-minilm-l6-v2</artifactId>
34+
<version>1.7.1-beta14</version>
35+
</dependency>
36+
37+
<dependency>
38+
<groupId>org.slf4j</groupId>
39+
<artifactId>slf4j-simple</artifactId>
40+
<version>2.0.12</version>
41+
</dependency>
42+
43+
</dependencies>
44+
45+
</project>
46+

‎yugabytedb-example/run-example.sh‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
#!/bin/bash
2+
3+
# Run YugabyteDB examples as standalone applications
4+
# This avoids exec:java classloader issues with Testcontainers
5+
6+
set -e
7+
8+
EXAMPLE_CLASS=${1:-YugabyteDBEmbeddingStoreExample}
9+
10+
echo "🚀 Building project..."
11+
../mvnw clean compile -q
12+
13+
echo "📦 Running example: $EXAMPLE_CLASS"
14+
echo ""
15+
16+
# Run with proper classpath
17+
../mvnw exec:exec -Dexec.executable="java" \
18+
-Dexec.args="-cp %classpath $EXAMPLE_CLASS" \
19+
-q
20+
21+
echo ""
22+
echo "✅ Example execution completed!"
23+

‎yugabytedb-example/src/main/java/YugabyteDBEmbeddingStoreExample.java‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
import dev.langchain4j.community.store.embedding.yugabytedb.YugabyteDBEmbeddingStore;
2+
import dev.langchain4j.community.store.embedding.yugabytedb.YugabyteDBEngine;
3+
import dev.langchain4j.data.embedding.Embedding;
4+
import dev.langchain4j.data.segment.TextSegment;
5+
import dev.langchain4j.model.embedding.EmbeddingModel;
6+
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
7+
import dev.langchain4j.store.embedding.EmbeddingMatch;
8+
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
9+
import dev.langchain4j.store.embedding.EmbeddingStore;
10+
import org.testcontainers.containers.GenericContainer;
11+
import org.testcontainers.containers.wait.strategy.Wait;
12+
import org.testcontainers.utility.DockerImageName;
13+
14+
import java.time.Duration;
15+
import java.util.List;
16+
17+
public class YugabyteDBEmbeddingStoreExample {
18+
19+
public static void main(String[] args) {
20+
GenericContainer<?> yugabyteContainer = null;
21+
YugabyteDBEngine engine = null;
22+
23+
try {
24+
DockerImageName dockerImageName = DockerImageName.parse("yugabytedb/yugabyte:2025.1.0.1-b3");
25+
yugabyteContainer = new GenericContainer<>(dockerImageName)
26+
.withExposedPorts(5433, 7000, 9000, 15433, 9042)
27+
.withCommand("bin/yugabyted", "start", "--background=false")
28+
.waitingFor(Wait.forListeningPorts(5433).withStartupTimeout(Duration.ofMinutes(5)));
29+
30+
yugabyteContainer.start();
31+
32+
EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
33+
34+
// Create YugabyteDB engine with PostgreSQL driver
35+
engine = YugabyteDBEngine.builder()
36+
.host(yugabyteContainer.getHost())
37+
.port(yugabyteContainer.getMappedPort(5433))
38+
.database("yugabyte")
39+
.username("yugabyte")
40+
.password("yugabyte")
41+
.usePostgreSQLDriver(true) // Use PostgreSQL JDBC driver
42+
.build();
43+
44+
EmbeddingStore<TextSegment> embeddingStore = YugabyteDBEmbeddingStore.builder()
45+
.engine(engine)
46+
.tableName("test_embeddings")
47+
.dimension(embeddingModel.dimension())
48+
.createTableIfNotExists(true)
49+
.build();
50+
51+
TextSegment segment1 = TextSegment.from("I like football.");
52+
Embedding embedding1 = embeddingModel.embed(segment1).content();
53+
embeddingStore.add(embedding1, segment1);
54+
55+
TextSegment segment2 = TextSegment.from("The weather is good today.");
56+
Embedding embedding2 = embeddingModel.embed(segment2).content();
57+
embeddingStore.add(embedding2, segment2);
58+
59+
Embedding queryEmbedding = embeddingModel.embed("What is your favourite sport?").content();
60+
61+
EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
62+
.queryEmbedding(queryEmbedding)
63+
.maxResults(1)
64+
.build();
65+
66+
List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.search(embeddingSearchRequest).matches();
67+
68+
EmbeddingMatch<TextSegment> embeddingMatch = relevant.get(0);
69+
70+
System.out.println(embeddingMatch.score()); // ~0.8144
71+
System.out.println(embeddingMatch.embedded().text()); // I like football.
72+
73+
System.out.println("\n✅ Example completed successfully!");
74+
75+
} catch (Exception e) {
76+
System.err.println("❌ Error running example: " + e.getMessage());
77+
e.printStackTrace();
78+
} finally {
79+
// Give Testcontainers time to cleanup gracefully
80+
try {
81+
Thread.sleep(2000);
82+
} catch (InterruptedException ignored) {
83+
}
84+
85+
// Cleanup resources
86+
System.out.println("🧹 Cleaning up resources...");
87+
if (engine != null) {
88+
try {
89+
engine.close();
90+
} catch (Exception e) {
91+
System.err.println("Error closing engine: " + e.getMessage());
92+
}
93+
}
94+
if (yugabyteContainer != null) {
95+
try {
96+
yugabyteContainer.stop();
97+
} catch (Exception e) {
98+
System.err.println("Error stopping container: " + e.getMessage());
99+
}
100+
}
101+
}
102+
}
103+
}
104+

‎yugabytedb-example/src/main/java/YugabyteDBEmbeddingStoreWithMetadataExample.java‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
import dev.langchain4j.community.store.embedding.yugabytedb.DefaultMetadataStorageConfig;
2+
import dev.langchain4j.community.store.embedding.yugabytedb.MetadataStorageConfig;
3+
import dev.langchain4j.community.store.embedding.yugabytedb.MetadataStorageMode;
4+
import dev.langchain4j.community.store.embedding.yugabytedb.YugabyteDBEmbeddingStore;
5+
import dev.langchain4j.community.store.embedding.yugabytedb.YugabyteDBEngine;
6+
import dev.langchain4j.data.document.Metadata;
7+
import dev.langchain4j.data.embedding.Embedding;
8+
import dev.langchain4j.data.segment.TextSegment;
9+
import dev.langchain4j.model.embedding.EmbeddingModel;
10+
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
11+
import dev.langchain4j.store.embedding.EmbeddingMatch;
12+
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
13+
import dev.langchain4j.store.embedding.EmbeddingStore;
14+
import dev.langchain4j.store.embedding.filter.Filter;
15+
import dev.langchain4j.store.embedding.filter.comparison.IsEqualTo;
16+
import org.testcontainers.containers.GenericContainer;
17+
import org.testcontainers.containers.wait.strategy.Wait;
18+
import org.testcontainers.utility.DockerImageName;
19+
20+
import java.time.Duration;
21+
import java.util.List;
22+
23+
public class YugabyteDBEmbeddingStoreWithMetadataExample {
24+
25+
public static void main(String[] args) {
26+
GenericContainer<?> yugabyteContainer = null;
27+
YugabyteDBEngine engine = null;
28+
29+
try {
30+
DockerImageName dockerImageName = DockerImageName.parse("yugabytedb/yugabyte:2025.1.0.1-b3");
31+
yugabyteContainer = new GenericContainer<>(dockerImageName)
32+
.withExposedPorts(5433, 7000, 9000, 15433, 9042)
33+
.withCommand("bin/yugabyted", "start", "--background=false")
34+
.waitingFor(Wait.forListeningPorts(5433).withStartupTimeout(Duration.ofMinutes(5)));
35+
36+
yugabyteContainer.start();
37+
38+
EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();
39+
40+
// Create YugabyteDB engine with PostgreSQL driver
41+
engine = YugabyteDBEngine.builder()
42+
.host(yugabyteContainer.getHost())
43+
.port(yugabyteContainer.getMappedPort(5433))
44+
.database("yugabyte")
45+
.username("yugabyte")
46+
.password("yugabyte")
47+
.usePostgreSQLDriver(true)
48+
.build();
49+
50+
// Configure metadata storage (JSONB format)
51+
MetadataStorageConfig metadataConfig = DefaultMetadataStorageConfig.builder()
52+
.storageMode(MetadataStorageMode.COMBINED_JSONB)
53+
.build();
54+
55+
EmbeddingStore<TextSegment> embeddingStore = YugabyteDBEmbeddingStore.builder()
56+
.engine(engine)
57+
.tableName("test_embeddings_with_metadata")
58+
.dimension(embeddingModel.dimension())
59+
.metadataStorageConfig(metadataConfig)
60+
.createTableIfNotExists(true)
61+
.build();
62+
63+
// Add embeddings with metadata
64+
TextSegment segment1 = TextSegment.from("I like football.",
65+
Metadata.from("category", "sports").put("user", "john"));
66+
Embedding embedding1 = embeddingModel.embed(segment1).content();
67+
embeddingStore.add(embedding1, segment1);
68+
69+
TextSegment segment2 = TextSegment.from("The weather is good today.",
70+
Metadata.from("category", "weather").put("user", "alice"));
71+
Embedding embedding2 = embeddingModel.embed(segment2).content();
72+
embeddingStore.add(embedding2, segment2);
73+
74+
TextSegment segment3 = TextSegment.from("I love basketball.",
75+
Metadata.from("category", "sports").put("user", "bob"));
76+
Embedding embedding3 = embeddingModel.embed(segment3).content();
77+
embeddingStore.add(embedding3, segment3);
78+
79+
// Search with metadata filter
80+
Embedding queryEmbedding = embeddingModel.embed("What sport do you like?").content();
81+
82+
Filter categoryFilter = new IsEqualTo("category", "sports");
83+
84+
EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest.builder()
85+
.queryEmbedding(queryEmbedding)
86+
.maxResults(5)
87+
.filter(categoryFilter)
88+
.build();
89+
90+
List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.search(searchRequest).matches();
91+
92+
System.out.println("Found " + relevant.size() + " sports-related results:");
93+
for (EmbeddingMatch<TextSegment> match : relevant) {
94+
System.out.println("Score: " + match.score());
95+
System.out.println("Text: " + match.embedded().text());
96+
System.out.println("Metadata: " + match.embedded().metadata());
97+
System.out.println("---");
98+
}
99+
100+
System.out.println("\n✅ Example completed successfully!");
101+
102+
} catch (Exception e) {
103+
System.err.println("❌ Error running example: " + e.getMessage());
104+
e.printStackTrace();
105+
} finally {
106+
// Give Testcontainers time to cleanup gracefully
107+
try {
108+
Thread.sleep(2000);
109+
} catch (InterruptedException ignored) {
110+
}
111+
112+
// Cleanup resources
113+
System.out.println("🧹 Cleaning up resources...");
114+
if (engine != null) {
115+
try {
116+
engine.close();
117+
} catch (Exception e) {
118+
System.err.println("Error closing engine: " + e.getMessage());
119+
}
120+
}
121+
if (yugabyteContainer != null) {
122+
try {
123+
yugabyteContainer.stop();
124+
} catch (Exception e) {
125+
System.err.println("Error stopping container: " + e.getMessage());
126+
}
127+
}
128+
}
129+
}
130+
}
131+

0 commit comments

Comments
 (0)