Is Text Embedding supported?

[mshytikov]

Hi,

I'm having trouble generating embedding using Vertex AI API via SDK Using the gemini-embedding-001 model.

Details
The embedContent endpoint is supported by SDK  but it just does not work with the model. 
rpc method: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/projects.locations.publishers.models/embedContent
sdk support: 

google-cloud-rust/src/generated/cloud/aiplatform/v1/src/stub.rs

Lines 6036 to 6037 in 509556c

	/// Implements [super::client::PredictionService::embed_content].
	fn embed_content(

When called via SDK the server gives BadReqeust with strange and misleading information. 

Error { kind: Service(ServiceDetails { status_code: Some(400), ... status: Status { code: InvalidArgument, message: "Invalid value (oneof), oneof field '_model' is already set. Cannot set 'model'", details: [BadRequest(BadRequest { field_violations: [FieldViolation { field: "", description: "Invalid value (oneof), oneof field '_model' is already set. Cannot set 'model'", reason: "", localized_message: None }] })] } }), source: None })

When called via  The REST API is saying 'model not supported' with

  "error": {
    "code": 400,
    "message": "Publisher Model `projects/.../locations/europe-west4/publishers/google/models/gemini-embedding-001` is not supported in the embedContent API.",
    "status": "INVALID_ARGUMENT"
  }
}

The official documentations here https://docs.cloud.google.com/vertex-ai/generative-ai/docs/embeddings/get-text-embeddings#get_text_embeddings_for_a_snippet_of_text
Gives examples where they use :predict rpc method https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1/projects.locations.publishers.models/predict
And it works via the REST API. 

But when I want to use it via Rust SDK I can't make it work.
The TextEmbeddingPredictionInstance should be passed as part of instances parameters to google_cloud_aiplatform_v1::model::PredictRequest ( rpc method and rpc type )
But it looks like it is not supported by Rust SDK (the TextEmbeddingPredictionInstance is not defined) https://raw.githubusercontent.com/googleapis/google-cloud-rust/509556cf9b529712e81e0310f8ff167c145d315c/src/generated/cloud/aiplatform/v1/src/model.rs)

Qustion
Am I missing something or is it just a Rust SDK limitation? 
Could you please help with some example how to use the :predict endpoint for embedding use cases via SDK (for gemini-embedding-001 model).    

Thank you,
Max

P.S.
REST Examples for the reference:

#!/bin/bash

# gives Model is not supported
ping-embedContent() {
  local project_id=$1
  curl -X POST \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json" \
    "https://europe-west4-aiplatform.googleapis.com/v1/projects/${project_id}/locations/europe-west4/publishers/google/models/gemini-embedding-001:embedContent" \
    -d '{
      "content": {
        "parts": [
          {
            "text": "I would like embeddings for this text!"
          }
        ],
        "role": "USER"
      },
      "taskType": "RETRIEVAL_DOCUMENT"
    }'
}

# works
ping-predict() {
  local project_id=$1
  curl -X POST \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json" \
    "https://europe-west4-aiplatform.googleapis.com/v1/projects/${project_id}/locations/europe-west4/publishers/google/models/gemini-embedding-001:predict" \
    -d '{
      "instances": [
        {
          "task_type": "RETRIEVAL_DOCUMENT",
          "title": "document title",
          "content": "I would like embeddings for this text!"
        }
      ]
    }'
}

[coryan]

This looks more like a question about the service than a question about the Rust SDK. The fact you were able to reproduce with plain REST suggests that is the case. I am no expert in all the Google Cloud APIs, so I apologize in advance if I get something wrong. The service resources page has some links to get help directly from the service folks.

With that out of the way. The list of supported services is here:

https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models

The model IDs are a bit hard to find, but are included with each model sheet, e.g.:

The gemini-embedding-001 model is part of the Gemini API:

https://ai.google.dev/gemini-api/docs/embeddings

I find it confusing myself, but the Gemini API and the Vertex AI APIs are different things, with different models. They use similar methods and messages, but they use different endpoints and models. The Rust SDK targets the Vertex AI API.

[mshytikov]

Hi @coryan, thanks for the answer. Much appreciated.

I fully agree that the endpoint vs. supported model question is a server-side concern, not an SDK issue.

However, as mentioned in my problem explanation, I suspect that the Rust SDK does not currently cover the Vertex AI API fully. Specifically, I expected the SDK to include TextEmbeddingPredictionInstance since it is a part of the Vertex AI API
See this part from above (* it has wrong link to REST docs but the same type definition exists in gRPC docs)

But when I want to use it via Rust SDK I can't make it work.
The TextEmbeddingPredictionInstance should be passed as part of instances parameters to google_cloud_aiplatform_v1::model::PredictRequest ( rpc method and rpc type )
But it looks like it is not supported by Rust SDK (the TextEmbeddingPredictionInstance is not defined) https://raw.githubusercontent.com/googleapis/google-cloud-rust/509556cf9b529712e81e0310f8ff167c145d315c/src/generated/cloud/aiplatform/v1/src/model.rs)

Speaking in terms of Rust SDK, this is how I would expect it to look like:

        use google_cloud_aiplatform_v1 as vertexai;

        let endpoint = "https://europe-west4-aiplatform.googleapis.com";
        let project = "my-gcp-project";
        let model = "gemini-embedding-001";
        let location = "europe-west4";

        let client = vertexai::client::PredictionService::builder()
            .with_endpoint(endpoint)
            .build()
            .await
            .expect("create Vertex AI client");

        let response = client
            .predict()
            // should be set_model but based on server naming it is 'endpoint' Param so set_endpoint
            .set_endpoint(format!(
                "projects/{project}/locations/{location}/publishers/google/models/{model}"
            ))
            .set_instances([vertexai::model::TextEmbeddingPredictionInstance::new()
                // optional
                // .set_task_type(...)
                // .set_title(...)
                .set_content("embedd this text")])
            .send()
            .await
            .expect("send embedding request");

The TextEmbeddingPredictionInstance does exists on server side in gRPC spec here
https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rpc/google.cloud.aiplatform.v1/schema/predict.instance#google.cloud.aiplatform.v1.schema.predict.instance.TextEmbeddingPredictionInstance
But the it does not exists in Rust SDK.
(note: I also can't find the definition of that TextEmbeddingPredictionInstance in googleapis repo )

So I got confused what is supported and what is not. Since this 'predict' functionality is v1 and the embedding model is GA on Vertex AI

Shall I just use raw wkt::Value in set_instances instead of waiting for TextEmbeddingPredictionInstance to be added to SDK?

[coryan]

However, as mentioned in my problem explanation, I suspect that the Rust SDK does not currently cover the Vertex AI API fully. Specifically, I expected the SDK to include TextEmbeddingPredictionInstance since it is a part of the Vertex AI API See this part from above (* it has wrong link to REST docs but the same type definition exists in gRPC docs)

I missed that part of your question. I read your question again, and I do not think I missed anything else, let me know if I did.

Speaking in terms of Rust SDK, this is how I would expect it to look like:

        use google_cloud_aiplatform_v1 as vertexai;

        let endpoint = "https://europe-west4-aiplatform.googleapis.com";
        let project = "my-gcp-project";
        let model = "gemini-embedding-001";
        let location = "europe-west4";

        let client = vertexai::client::PredictionService::builder()
            .with_endpoint(endpoint)
            .build()
            .await
            .expect("create Vertex AI client");

        let response = client
            .predict()
            // should be set_model but based on server naming it is 'endpoint' Param so set_endpoint
            .set_endpoint(format!(
                "projects/{project}/locations/{location}/publishers/google/models/{model}"
            ))
            .set_instances([vertexai::model::TextEmbeddingPredictionInstance::new()

This consumes a serde_json::Value:

https://docs.rs/google-cloud-aiplatform-v1/1.1.0/google_cloud_aiplatform_v1/model/struct.PredictRequest.html#method.set_instances

that is how the Rust SDK represents an arbitrary JSON object (or google.protobuf.Struct). You could change this to:

use serde_json::json;

    .set_instances([
      json!({
          "content": "...",
          "tittle": "....",
           "taskType": "RETRIEVAL_DOCUMENT" // or any of the other "enum" values.
      })
    ])

The TextEmbeddingPredictionInstance does exists on server side in gRPC spec here https://docs.cloud.google.com/vertex-ai/generative-ai/docs/reference/rpc/google.cloud.aiplatform.v1/schema/predict.instance#google.cloud.aiplatform.v1.schema.predict.instance.TextEmbeddingPredictionInstance But the it does not exists in Rust SDK. (note: I also can't find the definition of that TextEmbeddingPredictionInstance in googleapis repo )

Yes. That seems like a bug, I have filed a report internally, hopefully it would be addressed soon and newer versions of the client library will pick the new messages.

Shall I just use raw wkt::Value in set_instances

Yes. Even if we had TextEmbeddingPredictionInstance available you would need to serialize that to JSON. That would give you some type safety, but you can make it work without it.

[mshytikov]

Hey @coryan, perfect!

I just tried your suggestion to use serde_json::Value, and it all works now.
Thanks for reporting it internally.

To move forward, I will use the SDK with JSON and monitor for new releases. 👍

Thank you for the help and the SDK! 🙌