Skip to content

Commit ee5d882

Browse files
authored
Add integrated inference (#181)
## Problem Add integrated inference to Java SDK. ## Solution The code was already generated since integrated inference was a part of 2025-01 api spec. So as a part of this PR, I have added the following features: 1. Create index for model i.e. create an index with an associated embedding model 2. Configure an existing index to associate it with an embedding model 3. Upsert records 4. Search records by id 5. Search records by vector 6. Search records by text Example: ```java import io.pinecone.clients.Index; import io.pinecone.clients.Pinecone; import io.pinecone.helpers.RandomStringBuilder; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.openapitools.db_control.client.model.CreateIndexForModelRequest; import org.openapitools.db_control.client.model.CreateIndexForModelRequestEmbed; import org.openapitools.db_control.client.model.DeletionProtection; import org.openapitools.db_data.client.ApiException; import org.openapitools.db_data.client.model.SearchRecordsRequestQuery; import org.openapitools.db_data.client.model.SearchRecordsResponse; import org.openapitools.db_data.client.model.UpsertRecord; import java.util.ArrayList; import java.util.HashMap; import java.util.List; ... Pinecone pinecone = new Pinecone.Builder(System.getenv("PINECONE_API_KEY")).build(); String indexName = RandomStringBuilder.build("inf", 8); // Create index associated with a model HashMap<String, String> fieldMap = new HashMap<>(); fieldMap.put("text", "chunk_text"); CreateIndexForModelRequestEmbed embed = new CreateIndexForModelRequestEmbed() .model("multilingual-e5-large") .fieldMap(fieldMap); pinecone.createIndexForModel(indexName, CreateIndexForModelRequest.CloudEnum.AWS, "us-west-2", embed, DeletionProtection.DISABLED, new HashMap<>()); // Wait for index to be created Thread.sleep(10000); Index index = pinecone.getIndexConnection(indexName); // Upsert records HashMap<String, String> record1 = new HashMap<>(); record1.put("_id", "rec1"); record1.put("category", "digestive system"); record1.put("chunk_text", "Apples are a great source of dietary fiber, which supports digestion and helps maintain a healthy gut."); HashMap<String, String> record2 = new HashMap<>(); record2.put("_id", "rec2"); record2.put("category", "cultivation"); record2.put("chunk_text", "Apples originated in Central Asia and have been cultivated for thousands of years, with over 7,500 varieties available today."); HashMap<String, String> record3 = new HashMap<>(); record3.put("_id", "rec3"); record3.put("category", "immune system"); record3.put("chunk_text", "Rich in vitamin C and other antioxidants, apples contribute to immune health and may reduce the risk of chronic diseases."); HashMap<String, String> record4 = new HashMap<>(); record4.put("_id", "rec4"); record4.put("category", "endocrine system"); record4.put("chunk_text", "The high fiber content in apples can also help regulate blood sugar levels, making them a favorable snack for people with diabetes."); upsertRecords.add(record1); upsertRecords.add(record2); upsertRecords.add(record3); upsertRecords.add(record4); index.upsertRecords("example-namespace", upsertRecords); // Wait for vectors to be upserted Thread.sleep(5000); String namespace = "example-namespace"; List<String> fields = new ArrayList<>(); fields.add("category"); fields.add("chunk_text"); SearchRecordsRequestRerank rerank = new SearchRecordsRequestRerank() .model("bge-reranker-v2-m3") .topN(2) .rankFields(Arrays.asList("chunk_text")); // Search records SearchRecordsResponse recordsResponse = index.searchRecordsByText("Disease prevention", namespace, fields, 4, null, rerank); System.out.println(recordsResponse); ``` ## Type of Change - [ ] Bug fix (non-breaking change which fixes an issue) - [X] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected) - [ ] This change requires a documentation update - [ ] Infrastructure change (CI configs, etc) - [ ] Non-code change (docs, etc) - [ ] None of the above: (explain here) ## Test Plan Added integration test that creates an index associated with a model, upserts and queries records.
1 parent c83ea24 commit ee5d882

File tree

11 files changed

+522
-41
lines changed

11 files changed

+522
-41
lines changed

src/integration/java/io/pinecone/clients/ConnectionsMapTest.java

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package io.pinecone.clients;
22

3+
import io.pinecone.configs.PineconeConfig;
34
import io.pinecone.configs.PineconeConnection;
45
import io.pinecone.exceptions.PineconeNotFoundException;
56
import io.pinecone.helpers.RandomStringBuilder;
@@ -57,6 +58,10 @@ public void testMultipleIndexesWithMultipleClients() throws InterruptedException
5758
// Get index1's host
5859
String host1 = indexModel1.getHost();
5960

61+
// Create config1 for getting index connection and set the host
62+
PineconeConfig config1 = new PineconeConfig(System.getenv("PINECONE_API_KEY"));
63+
config1.setHost(host1);
64+
6065
// Create index-2
6166
pinecone1.createServerlessIndex(indexName2,
6267
null,
@@ -72,6 +77,11 @@ public void testMultipleIndexesWithMultipleClients() throws InterruptedException
7277
// Get index2's host
7378
String host2 = indexModel2.getHost();
7479

80+
// Create config2 for getting index connection and set the host
81+
PineconeConfig config2 = new PineconeConfig(System.getenv("PINECONE_API_KEY"));
82+
config1.setHost(host2);
83+
84+
7585
// Establish grpc connection for index-1
7686
Index index1_1 = pinecone1.getIndexConnection(indexName1);
7787
// Get connections map
@@ -94,7 +104,7 @@ public void testMultipleIndexesWithMultipleClients() throws InterruptedException
94104
assertEquals(host2, connectionsMap1_2.get(indexName2).toString());
95105

96106
// Establishing connections with index1 and index2 using another pinecone client
97-
pinecone2.getConnection(indexName1);
107+
pinecone2.getConnection(indexName1, config1);
98108
ConcurrentHashMap<String, PineconeConnection> connectionsMap2_1 = pinecone1.getConnectionsMap();
99109
// Verify the new connections map is pointing to the same reference
100110
assert connectionsMap2_1 == connectionsMap1_2;
@@ -103,7 +113,7 @@ public void testMultipleIndexesWithMultipleClients() throws InterruptedException
103113
// Verify the connection value for index1 is host1
104114
assertEquals(host1, connectionsMap2_1.get(indexName1).toString());
105115

106-
pinecone2.getConnection(indexName2);
116+
pinecone2.getConnection(indexName2, config2);
107117
ConcurrentHashMap<String, PineconeConnection> connectionsMap2_2 = pinecone1.getConnectionsMap();
108118
// Verify the new connections map is pointing to the same reference
109119
assert connectionsMap2_1 == connectionsMap2_2;

src/integration/java/io/pinecone/integration/controlPlane/serverless/DeletionProtectionTest.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,12 +44,12 @@ public void createPodIndexWithDeletionProtectionDisabled() {
4444
Map<String, String> actualTags = indexModel.getTags();
4545
Assertions.assertEquals(expectedTags, actualTags);
4646
// Configure index to enable deletionProtection
47-
controlPlaneClient.configureServerlessIndex(indexName, DeletionProtection.ENABLED, expectedTags);
47+
controlPlaneClient.configureServerlessIndex(indexName, DeletionProtection.ENABLED, expectedTags, null);
4848
indexModel = controlPlaneClient.describeIndex(indexName);
4949
deletionProtection = indexModel.getDeletionProtection();
5050
Assertions.assertEquals(deletionProtection, DeletionProtection.ENABLED);
5151
// Configure index to disable deletionProtection
52-
controlPlaneClient.configureServerlessIndex(indexName, DeletionProtection.DISABLED, expectedTags);
52+
controlPlaneClient.configureServerlessIndex(indexName, DeletionProtection.DISABLED, expectedTags, null);
5353
// Delete index
5454
controlPlaneClient.deleteIndex(indexName);
5555
}

src/integration/java/io/pinecone/integration/controlPlane/serverless/SparseIndexTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ public void configureSparseIndex() throws InterruptedException {
6363
waitUntilIndexIsReady(pinecone, indexName, 200000);
6464

6565
// Disable deletion protection and add more index tags
66-
pinecone.configureServerlessIndex(indexName, DeletionProtection.DISABLED, tags);
66+
pinecone.configureServerlessIndex(indexName, DeletionProtection.DISABLED, tags, null);
6767
Thread.sleep(7000);
6868

6969
// Describe index to confirm deletion protection is disabled

src/integration/java/io/pinecone/integration/dataPlane/QueryErrorTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ public static void setUp() throws IOException, InterruptedException {
3434
when(connectionMock.getBlockingStub()).thenReturn(stubMock);
3535
when(connectionMock.getAsyncStub()).thenReturn(asyncStubMock);
3636

37-
index = new Index(connectionMock, "some-index-name");
37+
index = new Index(config, connectionMock, "some-index-name");
3838
asyncIndex = new AsyncIndex(config, connectionMock, "some-index-name");
3939
}
4040

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package io.pinecone.integration.dataPlane;
2+
3+
import io.pinecone.clients.Index;
4+
import io.pinecone.clients.Pinecone;
5+
import io.pinecone.helpers.RandomStringBuilder;
6+
import org.junit.jupiter.api.Assertions;
7+
import org.junit.jupiter.api.Test;
8+
import org.openapitools.db_control.client.model.CreateIndexForModelRequest;
9+
import org.openapitools.db_control.client.model.CreateIndexForModelRequestEmbed;
10+
import org.openapitools.db_control.client.model.DeletionProtection;
11+
import org.openapitools.db_data.client.ApiException;
12+
import org.openapitools.db_data.client.model.SearchRecordsRequestQuery;
13+
import org.openapitools.db_data.client.model.SearchRecordsRequestRerank;
14+
import org.openapitools.db_data.client.model.SearchRecordsResponse;
15+
16+
import java.util.*;
17+
18+
public class UpsertAndSearchRecordsTest {
19+
@Test
20+
public void upsertAndSearchRecordsTest() throws ApiException, org.openapitools.db_control.client.ApiException, InterruptedException {
21+
Pinecone pinecone = new Pinecone.Builder(System.getenv("PINECONE_API_KEY")).build();
22+
String indexName = RandomStringBuilder.build("inf", 8);
23+
HashMap<String, String> fieldMap = new HashMap<>();
24+
fieldMap.put("text", "chunk_text");
25+
CreateIndexForModelRequestEmbed embed = new CreateIndexForModelRequestEmbed()
26+
.model("multilingual-e5-large")
27+
.fieldMap(fieldMap);
28+
pinecone.createIndexForModel(indexName, CreateIndexForModelRequest.CloudEnum.AWS, "us-west-2", embed, DeletionProtection.DISABLED, new HashMap<>());
29+
30+
// Wait for index to be created
31+
Thread.sleep(10000);
32+
33+
Index index = pinecone.getIndexConnection(indexName);
34+
ArrayList<Map<String, String>> upsertRecords = new ArrayList<>();
35+
36+
HashMap<String, String> record1 = new HashMap<>();
37+
record1.put("_id", "rec1");
38+
record1.put("category", "digestive system");
39+
record1.put("chunk_text", "Apples are a great source of dietary fiber, which supports digestion and helps maintain a healthy gut.");
40+
41+
HashMap<String, String> record2 = new HashMap<>();
42+
record2.put("_id", "rec2");
43+
record2.put("category", "cultivation");
44+
record2.put("chunk_text", "Apples originated in Central Asia and have been cultivated for thousands of years, with over 7,500 varieties available today.");
45+
46+
HashMap<String, String> record3 = new HashMap<>();
47+
record3.put("_id", "rec3");
48+
record3.put("category", "immune system");
49+
record3.put("chunk_text", "Rich in vitamin C and other antioxidants, apples contribute to immune health and may reduce the risk of chronic diseases.");
50+
51+
HashMap<String, String> record4 = new HashMap<>();
52+
record4.put("_id", "rec4");
53+
record4.put("category", "endocrine system");
54+
record4.put("chunk_text", "The high fiber content in apples can also help regulate blood sugar levels, making them a favorable snack for people with diabetes.");
55+
56+
upsertRecords.add(record1);
57+
upsertRecords.add(record2);
58+
upsertRecords.add(record3);
59+
upsertRecords.add(record4);
60+
61+
index.upsertRecords("example-namespace", upsertRecords);
62+
63+
String namespace = "example-namespace";
64+
HashMap<String, String> inputsMap = new HashMap<>();
65+
inputsMap.put("text", "Disease prevention");
66+
SearchRecordsRequestQuery query = new SearchRecordsRequestQuery()
67+
.topK(4)
68+
.inputs(inputsMap);
69+
70+
List<String> fields = new ArrayList<>();
71+
fields.add("category");
72+
fields.add("chunk_text");
73+
74+
// Wait for vectors to be upserted
75+
Thread.sleep(5000);
76+
77+
SearchRecordsResponse recordsResponse = index.searchRecords(namespace, query, fields, null);
78+
Assertions.assertEquals(upsertRecords.size(), recordsResponse.getResult().getHits().size());
79+
Assertions.assertEquals(record3.get("_id"), recordsResponse.getResult().getHits().get(0).getId());
80+
81+
recordsResponse = index.searchRecordsById(record1.get("_id"), namespace, fields, 1, null, null);
82+
Assertions.assertEquals(1, recordsResponse.getResult().getHits().size());
83+
Assertions.assertEquals(record1.get("_id"), recordsResponse.getResult().getHits().get(0).getId());
84+
85+
SearchRecordsRequestRerank rerank = new SearchRecordsRequestRerank()
86+
.model("bge-reranker-v2-m3")
87+
.topN(2)
88+
.rankFields(Arrays.asList("chunk_text"));
89+
90+
recordsResponse = index.searchRecordsByText("Disease prevention", namespace, fields, 4, null, rerank);
91+
Assertions.assertEquals(record3.get("_id"), recordsResponse.getResult().getHits().get(0).getId());
92+
93+
pinecone.deleteIndex(indexName);
94+
}
95+
}

src/integration/java/io/pinecone/integration/dataPlane/UpsertErrorTest.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ public static void setUp() throws IOException, InterruptedException {
3737
when(connectionMock.getBlockingStub()).thenReturn(stubMock);
3838
when(connectionMock.getAsyncStub()).thenReturn(asyncStubMock);
3939

40-
index = new Index(connectionMock, "some-index-name");
40+
index = new Index(config, connectionMock, "some-index-name");
4141
asyncIndex = new AsyncIndex(config, connectionMock, "some-index-name");
4242
}
4343

src/main/java/io/pinecone/clients/AsyncIndex.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,9 @@
3131

3232
import static io.pinecone.clients.Pinecone.buildOkHttpClient;
3333

34-
3534
/**
36-
* A client for interacting with a Pinecone index via GRPC asynchronously. Allows for upserting, querying, fetching, updating, and deleting vectors.
35+
* A client for interacting with a Pinecone index asynchronously. Allows for vector operations such as upserting,
36+
* querying, fetching, updating, and deleting vectors along with records operations such as upsert and search records.
3737
* This class provides a direct interface to interact with a specific index, encapsulating network communication and request validation.
3838
* <p>
3939
* Example:
@@ -70,6 +70,7 @@ public class AsyncIndex implements IndexInterface<ListenableFuture<UpsertRespons
7070
* AsyncIndex asyncIndex = client.getAsyncIndexConnection("my-index");
7171
* }</pre>
7272
*
73+
* @param config The {@link PineconeConfig} configuration of the index.
7374
* @param connection The {@link PineconeConnection} configuration to be used for this index.
7475
* @param indexName The name of the index to interact with. The index host will be automatically resolved.
7576
* @throws PineconeValidationException if the connection object is null.

0 commit comments

Comments
 (0)