From ea5209cc18c168bef082d35dc21d733367faa980 Mon Sep 17 00:00:00 2001 From: Antonio Calabrese Date: Tue, 19 Aug 2025 17:33:46 +0100 Subject: [PATCH] googleai: fix user provided embedding model getting overridden Fixes issue where user-provided embedding model was being overridden by default model --- .../internal/palmclient/palm_client_option.go | 64 ++++++++++ .../internal/palmclient/palmclient.go | 38 +++--- llms/googleai/option.go | 10 ++ llms/googleai/palm/palm_llm.go | 6 +- llms/googleai/shared_test/shared_test.go | 111 ++++++++++++++++++ llms/googleai/vertex/new.go | 6 +- 6 files changed, 218 insertions(+), 17 deletions(-) create mode 100644 llms/googleai/internal/palmclient/palm_client_option.go diff --git a/llms/googleai/internal/palmclient/palm_client_option.go b/llms/googleai/internal/palmclient/palm_client_option.go new file mode 100644 index 00000000..cf234646 --- /dev/null +++ b/llms/googleai/internal/palmclient/palm_client_option.go @@ -0,0 +1,64 @@ +package palmclient + +import ( + "google.golang.org/api/option" +) + +const ( + embeddingModelName = "text-embedding-005" + TextModelName = "text-bison" + ChatModelName = "chat-bison" +) + +// Options are the Palm client options +type Options struct { + EmbeddingModelName string + TextModelName string + ChatModelName string + ClientOptions []option.ClientOption +} + +// Option is an option +type Option func(*Options) + +// WithEmbeddingModelName sets the default embedding model +func WithEmbeddingModelName(modelName string) Option { + return func(o *Options) { + if modelName != "" { + o.EmbeddingModelName = modelName + } + } +} + +// WithTextModelName sets the default text model +func WithTextModelName(modelName string) Option { + return func(o *Options) { + if modelName != "" { + o.TextModelName = modelName + } + } +} + +// WithChatModelName sets the default chat model +func WithChatModelName(modelName string) Option { + return func(o *Options) { + if modelName != "" { + o.ChatModelName = modelName + } + } +} + +// WithClientOptions sets the client options for the Google API client +func WithClientOptions(opts ...option.ClientOption) Option { + return func(o *Options) { + o.ClientOptions = append(o.ClientOptions, opts...) + } +} + +func defaultOptions() Options { + return Options{ + EmbeddingModelName: embeddingModelName, + TextModelName: TextModelName, + ChatModelName: ChatModelName, + } +} diff --git a/llms/googleai/internal/palmclient/palmclient.go b/llms/googleai/internal/palmclient/palmclient.go index 9b14cb9e..4bed8e9c 100644 --- a/llms/googleai/internal/palmclient/palmclient.go +++ b/llms/googleai/internal/palmclient/palmclient.go @@ -28,40 +28,48 @@ var defaultParameters = map[string]interface{}{ //nolint:gochecknoglobals } const ( - embeddingModelName = "text-embedding-005" - TextModelName = "text-bison" - ChatModelName = "chat-bison" - defaultMaxConns = 4 ) // PaLMClient represents a Vertex AI based PaLM API client. type PaLMClient struct { - client *aiplatform.PredictionClient - projectID string + client *aiplatform.PredictionClient + projectID string + embeddingModelName string + textModelName string + chatModelName string } // New returns a new Vertex AI based PaLM API client. -func New(ctx context.Context, projectID, location string, opts ...option.ClientOption) (*PaLMClient, error) { +func New(ctx context.Context, projectID, location string, opts ...Option) (*PaLMClient, error) { numConns := runtime.GOMAXPROCS(0) if numConns > defaultMaxConns { numConns = defaultMaxConns } + + pOpt := defaultOptions() + for _, o := range opts { + o(&pOpt) + } + o := []option.ClientOption{ option.WithGRPCConnectionPool(numConns), option.WithEndpoint(fmt.Sprintf("%s-aiplatform.googleapis.com:443", location)), } - opts = append(o, opts...) + pOpt.ClientOptions = append(o, pOpt.ClientOptions...) // PredictionClient only support GRPC. - opts = append(opts, option.WithHTTPClient(nil)) + pOpt.ClientOptions = append(pOpt.ClientOptions, option.WithHTTPClient(nil)) - client, err := aiplatform.NewPredictionClient(ctx, opts...) + client, err := aiplatform.NewPredictionClient(ctx, pOpt.ClientOptions...) if err != nil { return nil, err } return &PaLMClient{ - client: client, - projectID: projectID, + client: client, + projectID: projectID, + embeddingModelName: pOpt.EmbeddingModelName, + textModelName: pOpt.TextModelName, + chatModelName: pOpt.ChatModelName, }, nil } @@ -92,7 +100,7 @@ func (c *PaLMClient) CreateCompletion(ctx context.Context, r *CompletionRequest) "top_k": r.TopK, "stopSequences": convertArray(r.StopSequences), } - predictions, err := c.batchPredict(ctx, TextModelName, r.Prompts, params) + predictions, err := c.batchPredict(ctx, c.textModelName, r.Prompts, params) if err != nil { return nil, err } @@ -118,7 +126,7 @@ type EmbeddingRequest struct { // CreateEmbedding creates embeddings. func (c *PaLMClient) CreateEmbedding(ctx context.Context, r *EmbeddingRequest) ([][]float32, error) { params := map[string]interface{}{} - responses, err := c.batchPredict(ctx, embeddingModelName, r.Input, params) + responses, err := c.batchPredict(ctx, c.embeddingModelName, r.Input, params) if err != nil { return nil, err } @@ -327,7 +335,7 @@ func (c *PaLMClient) chat(ctx context.Context, r *ChatRequest) ([]*structpb.Valu structpb.NewStructValue(instance), } resp, err := c.client.Predict(ctx, &aiplatformpb.PredictRequest{ - Endpoint: c.projectLocationPublisherModelPath(c.projectID, "us-central1", "google", ChatModelName), + Endpoint: c.projectLocationPublisherModelPath(c.projectID, "us-central1", "google", c.chatModelName), Instances: instances, Parameters: structpb.NewStructValue(mergedParams), }) diff --git a/llms/googleai/option.go b/llms/googleai/option.go index 05a79721..ef61f029 100644 --- a/llms/googleai/option.go +++ b/llms/googleai/option.go @@ -7,6 +7,7 @@ import ( "cloud.google.com/go/vertexai/genai" "google.golang.org/api/option" + "google.golang.org/grpc" ) // Options is a set of options for GoogleAI and Vertex clients. @@ -99,6 +100,15 @@ func WithHTTPClient(httpClient *http.Client) Option { } } +// WithGRPCConn appends a ClientOption that uses the provided gRPC client connection to +// make requests. +// This is useful for testing embeddings in vertex clients. +func WithGRPCConn(conn *grpc.ClientConn) Option { + return func(opts *Options) { + opts.ClientOptions = append(opts.ClientOptions, option.WithGRPCConn(conn)) + } +} + // WithCloudProject passes the GCP cloud project name to the client. This is // useful for vertex clients. func WithCloudProject(p string) Option { diff --git a/llms/googleai/palm/palm_llm.go b/llms/googleai/palm/palm_llm.go index 54d14263..ecba910a 100644 --- a/llms/googleai/palm/palm_llm.go +++ b/llms/googleai/palm/palm_llm.go @@ -116,5 +116,9 @@ func newClient(opts ...Option) (*palmclient.PaLMClient, error) { return nil, ErrMissingLocation } - return palmclient.New(context.TODO(), options.projectID, options.location, options.clientOptions...) + palmOptions := []palmclient.Option{ + palmclient.WithClientOptions(options.clientOptions...), + } + + return palmclient.New(context.TODO(), options.location, options.projectID, palmOptions...) } diff --git a/llms/googleai/shared_test/shared_test.go b/llms/googleai/shared_test/shared_test.go index 2751b1c1..f8e12dae 100644 --- a/llms/googleai/shared_test/shared_test.go +++ b/llms/googleai/shared_test/shared_test.go @@ -7,6 +7,7 @@ import ( "encoding/json" "fmt" "io" + "net" "net/http" "os" "path/filepath" @@ -16,6 +17,7 @@ import ( "strings" "testing" + "cloud.google.com/go/aiplatform/apiv1/aiplatformpb" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/tmc/langchaingo/embeddings" @@ -24,6 +26,11 @@ import ( "github.com/tmc/langchaingo/llms" "github.com/tmc/langchaingo/llms/googleai" "github.com/tmc/langchaingo/llms/googleai/vertex" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "google.golang.org/grpc/test/bufconn" + "google.golang.org/protobuf/types/known/structpb" ) func newGoogleAIClient(t *testing.T, opts ...googleai.Option) *googleai.GoogleAI { @@ -160,6 +167,21 @@ func TestVertexShared(t *testing.T) { } } +// TestVertex_WithCustomEmbeddingModel tests custom embedding models passed as an option. +// TODO: refactor testConfig to have a opts provider func so it this can be moved to a test config. +func TestVertex_WithCustomEmbeddingModel(t *testing.T) { + t.Helper() + t.Parallel() + const modelName = "custom-embedding-model" + opts := getCustomEmbeddingModelTestOptionsWithGRPC(t, modelName) + + llm, err := vertex.New(context.Background(), opts...) + require.NoError(t, err) + + _, err = llm.CreateEmbedding(context.Background(), []string{"test"}) + require.NoError(t, err) +} + func testMultiContentText(t *testing.T, llm llms.Model) { t.Helper() t.Parallel() @@ -617,6 +639,41 @@ func getHTTPTestClientOptions() []googleai.Option { return []googleai.Option{googleai.WithRest(), googleai.WithHTTPClient(client)} } +// getCustomEmbeddingModelTestOptionsWithGRPC creates options to connect to a fake gRPC server. +func getCustomEmbeddingModelTestOptionsWithGRPC(t *testing.T, model string) []googleai.Option { + t.Helper() + + // Create an in-memory "network connection" + lis := bufconn.Listen(1024 * 1024) + // Create the mock gRPC server and register the fake prediction service + grpcServer := grpc.NewServer() + mockPredictionServer := &mockPredictionServer{ + predictFunc: getPredictHandlerFuncWithCustomEmbeddingModel(model)} + aiplatformpb.RegisterPredictionServiceServer(grpcServer, mockPredictionServer) + + go func() { + if err := grpcServer.Serve(lis); err != nil { + t.Logf("gRPC server exited with error: %v", err) + } + }() + + t.Cleanup(func() { grpcServer.Stop() }) + + // Create a client connection to the fake server + conn, err := grpc.DialContext(context.Background(), "bufnet", + grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) { + return lis.Dial() + }), + grpc.WithInsecure(), + ) + require.NoError(t, err) + + return []googleai.Option{ + googleai.WithDefaultEmbeddingModel(model), + googleai.WithGRPCConn(conn), + } +} + type testRequestInterceptor struct{} func (i *testRequestInterceptor) RoundTrip(req *http.Request) (*http.Response, error) { @@ -666,3 +723,57 @@ func checkMatch(t *testing.T, got string, wants ...string) { } } } + +// PredictHandlerFunc is a handler func that matches the gRPC Predict method. +type PredictHandlerFunc func(context.Context, *aiplatformpb.PredictRequest) (*aiplatformpb.PredictResponse, error) + +// mockPredictionServer is a mock gRPC prediction server. +type mockPredictionServer struct { + aiplatformpb.UnimplementedPredictionServiceServer + predictFunc PredictHandlerFunc +} + +// NewMockPredictionServer creates a new server with a custom Predict handler. +func NewMockPredictionServer(handler PredictHandlerFunc) *mockPredictionServer { + return &mockPredictionServer{ + predictFunc: handler, + } +} + +// Predict implements the UnimplementedPredictionServiceServer. It calls predictFunc. +func (s *mockPredictionServer) Predict(ctx context.Context, req *aiplatformpb.PredictRequest) (*aiplatformpb.PredictResponse, error) { + if s.predictFunc == nil { + return nil, status.Error(codes.Unimplemented, "Predict handler was not provided") + } + + return s.predictFunc(ctx, req) +} + +// getPredictHandlerFuncWithCustomEmbeddingModel returns a predictFunc which checks that the embedding request has been made +// with the custom provided embedding model. +func getPredictHandlerFuncWithCustomEmbeddingModel(embeddingModel string) func(ctx context.Context, req *aiplatformpb.PredictRequest) (*aiplatformpb.PredictResponse, error) { + return func(ctx context.Context, req *aiplatformpb.PredictRequest) (*aiplatformpb.PredictResponse, error) { + expectedEndpointSuffix := "/models/" + embeddingModel + if !strings.HasSuffix(req.Endpoint, expectedEndpointSuffix) { + return nil, status.Errorf(codes.InvalidArgument, "model name mismatch, expected suffix '%s'", expectedEndpointSuffix) + } + + // Create a dummy embedding response that the client code will accept. + // The client expects a structure like: { "embeddings": { "values": [0.1, 0.2, ...] } } + embeddingStruct, err := structpb.NewStruct(map[string]interface{}{ + "embeddings": map[string]interface{}{ + "values": []interface{}{0.1, 0.2, 0.3, 0.4}, + }, + }) + if err != nil { + return nil, status.Errorf(codes.Internal, "failed to create embedding struct: %v", err) + } + + predictionValue := structpb.NewStructValue(embeddingStruct) + response := &aiplatformpb.PredictResponse{ + Predictions: []*structpb.Value{predictionValue}, + } + + return response, nil + } +} diff --git a/llms/googleai/vertex/new.go b/llms/googleai/vertex/new.go index f0c3e64c..33beb5aa 100644 --- a/llms/googleai/vertex/new.go +++ b/llms/googleai/vertex/new.go @@ -42,11 +42,15 @@ func New(ctx context.Context, opts ...googleai.Option) (*Vertex, error) { return nil, err } + palmOpts := []palmclient.Option{ + palmclient.WithEmbeddingModelName(clientOptions.DefaultEmbeddingModel), + palmclient.WithClientOptions(clientOptions.ClientOptions...), + } palmClient, err := palmclient.New( ctx, clientOptions.CloudProject, clientOptions.CloudLocation, - clientOptions.ClientOptions...) + palmOpts...) if err != nil { return nil, err }