Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add integrated inference #181

Open
wants to merge 6 commits into
base: main
Choose a base branch
from
Open

Conversation

rohanshah18
Copy link
Contributor

@rohanshah18 rohanshah18 commented Mar 21, 2025

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:

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)
  • 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.

@rohanshah18 rohanshah18 changed the title Add upsert and search records Add integrated inference Mar 28, 2025
@rohanshah18 rohanshah18 marked this pull request as ready for review March 28, 2025 17:29
@rohanshah18 rohanshah18 requested review from austin-denoble, jhamon and ssmith-pc and removed request for austin-denoble March 28, 2025 17:30
Comment on lines +64 to +77
HashMap<String, String> inputsMap = new HashMap<>();
inputsMap.put("text", "Disease prevention");
SearchRecordsRequestQuery query = new SearchRecordsRequestQuery()
.topK(4)
.inputs(inputsMap);

List<String> fields = new ArrayList<>();
fields.add("category");
fields.add("chunk_text");

// Wait for vectors to be upserted
Thread.sleep(5000);

SearchRecordsResponse recordsResponse = index.searchRecords(namespace, query, fields, null);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we should consider creating some API abstractions over the SearchRecordsRequestQuery interface.

For example, it might feel more natural if the user could query with something like:

RecordsQuery query = new TextRecordsQuery("Disease prevention")
        .topK(4);

index.searchRecords(namespace, query, ["category", "chunk_text"]);

We could also support ID and vector queries in a similar way?

vectorOperations.setCustomBaseUrl(protocol + config.getHost());
}

public void upsertRecords(String namespace, List<UpsertRecord> upsertRecord) throws ApiException {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we add an overload that accepts List<Map<String, Object>>?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This feels kind of cumbersome:

        UpsertRecord record1 = new UpsertRecord();
        record1.id("rec1");
        record1.putAdditionalProperty("category", "digestive system");
        record1.putAdditionalProperty("chunk_text", "Apples are a great source of dietary fiber, which supports digestion and helps maintain a healthy gut.");
        ArrayList<UpsertRecord> records = new ArrayList<>();

        UpsertRecord record2 = new UpsertRecord();
        record2.id("rec2");
        record2.putAdditionalProperty("category", "cultivation");
        record2.putAdditionalProperty("chunk_text", "Apples originated in Central Asia and have been cultivated for thousands of years, with over 7,500 varieties available today.");

Copy link
Contributor

@austin-denoble austin-denoble left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall this looks good, nice work! I've left a few comments around handling of ids in records, and a possibly breaking change to configureServerlessIndex.

@@ -931,6 +931,18 @@ public RequestBody serialize(Object obj, String contentType) throws ApiException
return RequestBody.create((File) obj, MediaType.parse(contentType));
} else if ("text/plain".equals(contentType) && obj instanceof String) {
return RequestBody.create((String) obj, MediaType.parse(contentType));
} else if ("application/x-ndjson".equals(contentType)) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is fine for now, but you'll need to remember to redo this change every time you regenerate.

Alternatively, there's a couple options we've used in other clients:

  • Updating the Java templates to inject this logic at code generation.
  • Using some combination of bash/sed to manipulate the output after generation, as a part of the build script process.

for(Map<String, String> record: upsertRecords) {
UpsertRecord upsertRecord = new UpsertRecord();
for (Map.Entry<String, String> entry : record.entrySet()) {
if(entry.getKey().equals("_id")) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible for an existing id to be overwritten by this logic? The server will respond with a 400 if you try and send both id and _id:

{"error":{"code":"INVALID_ARGUMENT","message":"Error parsing request: Invalid input: Document cannot have both _id and id fields (_id='rec1', id='rec1')"},"status":400}

Maybe instead of mapping one to the other in the client we just pass through what the user has provided, and if they need to clean up ids they can? I just checked for the presence of one or the other on the record in typescript: https://github.com/pinecone-io/pinecone-ts-client/blob/e5b39d7a6c77836b5b0ae1f696ce7e6f6020f82e/src/data/vectors/upsertRecords.ts#L29-L37

* @throws ApiException If there is an issue with the search operation. This could include network errors,
* invalid input data, or issues communicating with the Pinecone service.
*/
public SearchRecordsResponse searchRecordsByText(String text,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice job handling the lack of optionals in all of these methods. 👍

* @return {@link IndexModel} representing the configured index.
* @throws PineconeException if an error occurs during the operation, the index does not exist, or if any of the arguments are invalid.
*/
public IndexModel configureServerlessIndex(String indexName, DeletionProtection deletionProtection, Map<String, String> tags) throws PineconeException {
public IndexModel configureServerlessIndex(String indexName,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would be a breaking change, right? In terms of it altering the method signature for configureServerlessIndex.

If we're intending for these changes to release as a major version bump that's fine, but otherwise we may want to look at overloading this method too to keep things backwards compatible.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants