Google Vertex for ITextEmbeddingGenerator and ITextGenerator

[cisionmarkwalls]

For those of us on the GCP stack we would like to use Vertex for embeddings and text generation.

Something like this with configuration like is used for OpenAI (project names etc...):

public class VertexTextGenerator : ITextGenerator
{
    string projectId = "usa-dev-composer-001";
    string locationId = "us-central1";
    string publisher = "google";
    string model = "gemini-1.0-pro-002";
    private readonly Google.Cloud.AIPlatform.V1.PredictionServiceClient predictionClient;
    private readonly LlmUtilityServiceClient llmUtilityServiceClient;
    private readonly EndpointName predictionEndpoint;

    public int MaxTokenTotal => 10000;

    public VertexTextGenerator()
    {
        predictionClient = new PredictionServiceClientBuilder
        {
            Endpoint = $"{locationId}-aiplatform.googleapis.com"
        }.Build();
        llmUtilityServiceClient = new LlmUtilityServiceClientBuilder
        {
            Endpoint = $"{locationId}-aiplatform.googleapis.com"
        }.Build();
        predictionEndpoint = EndpointName.FromProjectLocationPublisherModel(projectId, locationId, publisher, model);
    }

    public int CountTokens(string text)
    {
        var request = new CountTokensRequest
        {
            Endpoint = $"projects/{projectId}/locations/{locationId}/publishers/{publisher}/models/{model}",
            Model = $"projects/{projectId}/locations/{locationId}/publishers/{publisher}/models/{model}",
            Contents =
            {
                new Content
                {
                    Role = "USER",
                    Parts = { new Part { Text = text } }
                }
            }
        };
        var response = llmUtilityServiceClient.CountTokensAsync(request).Result;
        int tokenCount = response.TotalTokens;
        return tokenCount;
    }

    public async IAsyncEnumerable<string> GenerateTextAsync(string prompt, TextGenerationOptions options, CancellationToken cancellationToken = default)
    {
        var generateContentRequest = new GenerateContentRequest
        {
            Model = $"projects/{projectId}/locations/{locationId}/publishers/{publisher}/models/{model}",
            Contents =
            {
                new Content
                {
                    Role = "USER",
                    Parts =
                    {
                        new Part { Text = prompt }
                    }
                }
            }
        };

        var stream = predictionClient.StreamGenerateContent(generateContentRequest)
            .GetResponseStream();
        await foreach (var response in EnumerateResponseStreamAsync(stream, cancellationToken))
        {
            foreach (var candidate in response.Candidates)
            {
                if (candidate != null)
                {
                    yield return candidate.Content.Parts[0].Text;
                }
            }
        }
    }

    private static async IAsyncEnumerable<TResponse> EnumerateResponseStreamAsync<TResponse>(AsyncResponseStream<TResponse> responseStream, [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        while (await responseStream.MoveNextAsync(cancellationToken))
        {
            yield return responseStream.Current;
        }
    }
}

Embeddings are pretty much the same with slightly different plumbing:

 public async Task<Embedding> GenerateEmbeddingAsync(string text, CancellationToken cancellationToken = default)
 {
     var instances = new List<Value>
     {
         Value.ForStruct(new()
         {
             Fields =
             {
                 ["content"] = Value.ForString(text),
             }
         })
     };

     var response = predictionClient.Predict(predictionEndpoint, instances, null);

     // Parse and return the embedding vector count.
     var values = ConvertToFloatList(response.Predictions.First().StructValue.Fields["embeddings"].StructValue.Fields["values"].ListValue.Values.ToList());
    
     //Works
     return new Embedding(values.ToArray());
 }

 private List<float> ConvertToFloatList(IEnumerable<Value> values)
 {
     var floatList = new List<float>();
     foreach (var item in values)
     {
         if (item.KindCase == Value.KindOneofCase.NumberValue)
         {
             floatList.Add((float)item.NumberValue);
         }
     }
     return floatList;
 }

If I productionalize that code (make it match the OpenAI patterns, add tests etc...) are you open to a PR for it?

[dluc]

if the code works and follows the existing KM design patterns (from similar classes in the project) it should be an easy merge. We have limited bandwidth so it might just take a while to go through the reviews.