From ea6350af1674b75ac0770880f13f35c27a5ac31f Mon Sep 17 00:00:00 2001 From: Antoine Popineau Date: Fri, 18 Jul 2025 19:21:41 +0200 Subject: [PATCH 1/2] Implement sending batch requests. --- adapter.go | 27 +++++ batches.go | 57 +++++++++ examples/batch/main.go | 47 ++++++++ go.mod | 67 ++++++++--- go.sum | 221 ++++++++++++++++++---------------- internal/interface.go | 4 +- llms/aistudio/adapter_test.go | 3 +- llms/aistudio/aistudio.go | 31 +++-- llms/aistudio/batches.go | 218 +++++++++++++++++++++++++++++++++ llms/aistudio/options.go | 6 + llms/openai/openai.go | 7 ++ mock_provider_test.go | 3 + options.go | 2 + request.go | 12 ++ 14 files changed, 571 insertions(+), 134 deletions(-) create mode 100644 batches.go create mode 100644 examples/batch/main.go create mode 100644 llms/aistudio/batches.go diff --git a/adapter.go b/adapter.go index 5fbbada..a13943a 100644 --- a/adapter.go +++ b/adapter.go @@ -17,6 +17,7 @@ const ( // It provides a contract for initializing, managing context, // and performing chat completions with different language models. type Llm interface { + SetName(string) // Init initializes the LLM provider with the given adapter configuration. // It is called once when the provider is added to the adapter. Init(llm internal.Adapter) error @@ -35,6 +36,10 @@ type Llm interface { // request options struct. This is used for type checking and reflection // when processing custom request options. RequestOptionsType() reflect.Type + + SubmitBatch(context.Context, internal.Adapter, ...Requester) (*BatchPromise, error) + Check(context.Context, *BatchPromise) (BatchStatus, error) + Wait(ctx context.Context, pr *BatchPromise) <-chan BatchWaitResponse } // LlmAdapter is the main entrypoint for interacting with different LLM providers. @@ -109,6 +114,28 @@ func (llm *LlmAdapter) GetProvider(requestProvider *string) (Llm, error) { return provider, nil } +func (llm *LlmAdapter) SubmitBatch(ctx context.Context, providerName string, reqs ...Requester) (*BatchPromise, error) { + p, ok := llm.providers[providerName] + if !ok { + return nil, errors.Newf("unknown provider '%s'", providerName) + } + + return p.SubmitBatch(ctx, llm, reqs...) +} + +func (llm *LlmAdapter) BatchPromise(providerName string, id string) (*BatchPromise, error) { + provider, ok := llm.providers[providerName] + if !ok { + return nil, errors.New("cannot find the provider that created this promise") + } + + return &BatchPromise{ + ProviderName: providerName, + Provider: provider, + Id: id, + }, nil +} + // LlmAdapter implementation of Adapter func (llm LlmAdapter) DefaultModel() string { diff --git a/batches.go b/batches.go new file mode 100644 index 0000000..4b4b526 --- /dev/null +++ b/batches.go @@ -0,0 +1,57 @@ +package llmadapter + +import ( + "context" + + "github.com/checkmarble/marble-llm-adapter/internal" + "github.com/cockroachdb/errors" +) + +type ( + BatchStatus int +) + +const ( + BatchPending BatchStatus = iota + BatchRunning + BatchFinished + BatchError +) + +type BatchUnsupported struct{} + +func (BatchUnsupported) SubmitBatch(ctx context.Context, llm internal.Adapter, reqs ...Requester) (*BatchPromise, error) { + return nil, errors.New("provider does not support batch mode") +} + +func (BatchUnsupported) Check(context.Context, *BatchPromise) (BatchStatus, error) { + return BatchError, errors.New("provider does not support batch mode") +} + +func (BatchUnsupported) Wait(ctx context.Context, pr *BatchPromise) <-chan BatchWaitResponse { + return nil +} + +type BatchRequest struct { + Provider Llm + Filename string +} + +type BatchPromise struct { + Provider Llm + ProviderName string + Id string +} + +func (p *BatchPromise) Check(ctx context.Context) (BatchStatus, error) { + return p.Provider.Check(ctx, p) +} + +func (p *BatchPromise) Wait(ctx context.Context) <-chan BatchWaitResponse { + return p.Provider.Wait(ctx, p) +} + +type BatchWaitResponse struct { + Status BatchStatus + Error error +} diff --git a/examples/batch/main.go b/examples/batch/main.go new file mode 100644 index 0000000..0342e72 --- /dev/null +++ b/examples/batch/main.go @@ -0,0 +1,47 @@ +package main + +import ( + "context" + "fmt" + "log" + "os" + + llmadapter "github.com/checkmarble/marble-llm-adapter" + "github.com/checkmarble/marble-llm-adapter/llms/aistudio" + "google.golang.org/genai" +) + +func main() { + ctx := context.Background() + + provider, _ := aistudio.New( + aistudio.WithBackend(genai.BackendVertexAI), + aistudio.WithProject(os.Getenv("GOOGLE_CLOUD_PROJECT")), + aistudio.WithLocation("europe-west1"), + aistudio.WithApiKey(os.Getenv("LLM_API_KEY")), + aistudio.WithBucket(os.Getenv("LLM_BATCH_BUCKET")), + ) + + llm, _ := llmadapter.New( + llmadapter.WithProvider("vertex", provider), + llmadapter.WithDefaultModel("gemini-2.5-flash"), + ) + + reqs := []llmadapter.Requester{ + llmadapter.NewUntypedRequest().WithProvider("vertex").WithId("how").WithText(llmadapter.RoleUser, "How are you?"), + llmadapter.NewUntypedRequest().WithProvider("vertex").WithId("addition").WithText(llmadapter.RoleUser, "What is 1 + 1?"), + } + + promise, err := llm.SubmitBatch(ctx, "vertex", reqs...) + if err != nil { + log.Fatal(err) + } + + result := <-promise.Wait(ctx) + + if result.Error != nil { + log.Fatal(err) + } + + fmt.Printf("%#v\n", result) +} diff --git a/go.mod b/go.mod index caf46f9..1d5d80c 100644 --- a/go.mod +++ b/go.mod @@ -3,52 +3,87 @@ module github.com/checkmarble/marble-llm-adapter go 1.24.4 require ( + cloud.google.com/go/storage v1.55.0 github.com/cockroachdb/errors v1.12.0 + github.com/fatih/structs v1.1.0 github.com/h2non/gock v1.2.0 github.com/invopop/jsonschema v0.13.0 github.com/openai/openai-go v1.9.0 github.com/samber/lo v1.51.0 + github.com/simonfrey/jsonl v0.0.0-20240904112901-935399b9a740 github.com/stretchr/testify v1.10.0 github.com/tidwall/gjson v1.18.0 google.golang.org/genai v1.15.0 ) require ( - cloud.google.com/go v0.116.0 // indirect - cloud.google.com/go/auth v0.9.3 // indirect - cloud.google.com/go/compute/metadata v0.5.0 // indirect + cel.dev/expr v0.20.0 // indirect + cloud.google.com/go v0.121.1 // indirect + cloud.google.com/go/auth v0.16.1 // indirect + cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect + cloud.google.com/go/compute/metadata v0.7.0 // indirect + cloud.google.com/go/iam v1.5.2 // indirect + cloud.google.com/go/monitoring v1.24.2 // indirect + github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.27.0 // indirect + github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.51.0 // indirect + github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.51.0 // indirect github.com/bahlo/generic-list-go v0.2.0 // indirect github.com/buger/jsonparser v1.1.1 // indirect + github.com/cespare/xxhash/v2 v2.3.0 // indirect + github.com/cncf/xds/go v0.0.0-20250121191232-2f005788dc42 // indirect github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b // indirect github.com/cockroachdb/redact v1.1.5 // indirect github.com/davecgh/go-spew v1.1.1 // indirect - github.com/fatih/structs v1.1.0 // indirect + github.com/envoyproxy/go-control-plane/envoy v1.32.4 // indirect + github.com/envoyproxy/protoc-gen-validate v1.2.1 // indirect + github.com/felixge/httpsnoop v1.0.4 // indirect github.com/getsentry/sentry-go v0.27.0 // indirect + github.com/go-jose/go-jose/v4 v4.0.4 // indirect + github.com/go-logr/logr v1.4.2 // indirect + github.com/go-logr/stdr v1.2.2 // indirect github.com/gogo/protobuf v1.3.2 // indirect - github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect - github.com/google/go-cmp v0.6.0 // indirect - github.com/google/s2a-go v0.1.8 // indirect - github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect + github.com/google/go-cmp v0.7.0 // indirect + github.com/google/s2a-go v0.1.9 // indirect + github.com/google/uuid v1.6.0 // indirect + github.com/googleapis/enterprise-certificate-proxy v0.3.6 // indirect + github.com/googleapis/gax-go/v2 v2.14.2 // indirect github.com/gorilla/websocket v1.5.3 // indirect github.com/h2non/parth v0.0.0-20190131123155-b4df798d6542 // indirect github.com/kr/pretty v0.3.1 // indirect github.com/kr/text v0.2.0 // indirect github.com/mailru/easyjson v0.7.7 // indirect github.com/pkg/errors v0.9.1 // indirect + github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect - github.com/rogpeppe/go-internal v1.9.0 // indirect + github.com/rogpeppe/go-internal v1.13.1 // indirect + github.com/spiffe/go-spiffe/v2 v2.5.0 // indirect github.com/stretchr/objx v0.5.2 // indirect github.com/tidwall/match v1.1.1 // indirect github.com/tidwall/pretty v1.2.1 // indirect github.com/tidwall/sjson v1.2.5 // indirect github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect - go.opencensus.io v0.24.0 // indirect - golang.org/x/crypto v0.36.0 // indirect - golang.org/x/net v0.38.0 // indirect - golang.org/x/sys v0.31.0 // indirect + github.com/zeebo/errs v1.4.0 // indirect + go.opentelemetry.io/auto/sdk v1.1.0 // indirect + go.opentelemetry.io/contrib/detectors/gcp v1.36.0 // indirect + go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0 // indirect + go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 // indirect + go.opentelemetry.io/otel v1.36.0 // indirect + go.opentelemetry.io/otel/metric v1.36.0 // indirect + go.opentelemetry.io/otel/sdk v1.36.0 // indirect + go.opentelemetry.io/otel/sdk/metric v1.36.0 // indirect + go.opentelemetry.io/otel/trace v1.36.0 // indirect + golang.org/x/crypto v0.38.0 // indirect + golang.org/x/net v0.40.0 // indirect + golang.org/x/oauth2 v0.30.0 // indirect + golang.org/x/sync v0.15.0 // indirect + golang.org/x/sys v0.33.0 // indirect golang.org/x/text v0.26.0 // indirect - google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 // indirect - google.golang.org/grpc v1.66.2 // indirect - google.golang.org/protobuf v1.34.2 // indirect + golang.org/x/time v0.11.0 // indirect + google.golang.org/api v0.235.0 // indirect + google.golang.org/genproto v0.0.0-20250505200425-f936aa4a68b2 // indirect + google.golang.org/genproto/googleapis/api v0.0.0-20250512202823-5a2f75b736a9 // indirect + google.golang.org/genproto/googleapis/rpc v0.0.0-20250512202823-5a2f75b736a9 // indirect + google.golang.org/grpc v1.72.1 // indirect + google.golang.org/protobuf v1.36.6 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index d7803ad..a019914 100644 --- a/go.sum +++ b/go.sum @@ -1,18 +1,41 @@ -cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw= -cloud.google.com/go v0.116.0 h1:B3fRrSDkLRt5qSHWe40ERJvhvnQwdZiHu0bJOpldweE= -cloud.google.com/go v0.116.0/go.mod h1:cEPSRWPzZEswwdr9BxE6ChEn01dWlTaF05LiC2Xs70U= -cloud.google.com/go/auth v0.9.3 h1:VOEUIAADkkLtyfr3BLa3R8Ed/j6w1jTBmARx+wb5w5U= -cloud.google.com/go/auth v0.9.3/go.mod h1:7z6VY+7h3KUdRov5F1i8NDP5ZzWKYmEPO842BgCsmTk= -cloud.google.com/go/compute/metadata v0.5.0 h1:Zr0eK8JbFv6+Wi4ilXAR8FJ3wyNdpxHKJNPos6LTZOY= -cloud.google.com/go/compute/metadata v0.5.0/go.mod h1:aHnloV2TPI38yx4s9+wAZhHykWvVCfu7hQbF+9CWoiY= -github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= +cel.dev/expr v0.20.0 h1:OunBvVCfvpWlt4dN7zg3FM6TDkzOePe1+foGJ9AXeeI= +cel.dev/expr v0.20.0/go.mod h1:MrpN08Q+lEBs+bGYdLxxHkZoUSsCp0nSKTs0nTymJgw= +cloud.google.com/go v0.121.1 h1:S3kTQSydxmu1JfLRLpKtxRPA7rSrYPRPEUmL/PavVUw= +cloud.google.com/go v0.121.1/go.mod h1:nRFlrHq39MNVWu+zESP2PosMWA0ryJw8KUBZ2iZpxbw= +cloud.google.com/go/auth v0.16.1 h1:XrXauHMd30LhQYVRHLGvJiYeczweKQXZxsTbV9TiguU= +cloud.google.com/go/auth v0.16.1/go.mod h1:1howDHJ5IETh/LwYs3ZxvlkXF48aSqqJUM+5o02dNOI= +cloud.google.com/go/auth/oauth2adapt v0.2.8 h1:keo8NaayQZ6wimpNSmW5OPc283g65QNIiLpZnkHRbnc= +cloud.google.com/go/auth/oauth2adapt v0.2.8/go.mod h1:XQ9y31RkqZCcwJWNSx2Xvric3RrU88hAYYbjDWYDL+c= +cloud.google.com/go/compute/metadata v0.7.0 h1:PBWF+iiAerVNe8UCHxdOt6eHLVc3ydFeOCw78U8ytSU= +cloud.google.com/go/compute/metadata v0.7.0/go.mod h1:j5MvL9PprKL39t166CoB1uVHfQMs4tFQZZcKwksXUjo= +cloud.google.com/go/iam v1.5.2 h1:qgFRAGEmd8z6dJ/qyEchAuL9jpswyODjA2lS+w234g8= +cloud.google.com/go/iam v1.5.2/go.mod h1:SE1vg0N81zQqLzQEwxL2WI6yhetBdbNQuTvIKCSkUHE= +cloud.google.com/go/logging v1.13.0 h1:7j0HgAp0B94o1YRDqiqm26w4q1rDMH7XNRU34lJXHYc= +cloud.google.com/go/logging v1.13.0/go.mod h1:36CoKh6KA/M0PbhPKMq6/qety2DCAErbhXT62TuXALA= +cloud.google.com/go/longrunning v0.6.7 h1:IGtfDWHhQCgCjwQjV9iiLnUta9LBCo8R9QmAFsS/PrE= +cloud.google.com/go/longrunning v0.6.7/go.mod h1:EAFV3IZAKmM56TyiE6VAP3VoTzhZzySwI/YI1s/nRsY= +cloud.google.com/go/monitoring v1.24.2 h1:5OTsoJ1dXYIiMiuL+sYscLc9BumrL3CarVLL7dd7lHM= +cloud.google.com/go/monitoring v1.24.2/go.mod h1:x7yzPWcgDRnPEv3sI+jJGBkwl5qINf+6qY4eq0I9B4U= +cloud.google.com/go/storage v1.55.0 h1:NESjdAToN9u1tmhVqhXCaCwYBuvEhZLLv0gBr+2znf0= +cloud.google.com/go/storage v1.55.0/go.mod h1:ztSmTTwzsdXe5syLVS0YsbFxXuvEmEyZj7v7zChEmuY= +cloud.google.com/go/trace v1.11.6 h1:2O2zjPzqPYAHrn3OKl029qlqG6W8ZdYaOWRyr8NgMT4= +cloud.google.com/go/trace v1.11.6/go.mod h1:GA855OeDEBiBMzcckLPE2kDunIpC72N+Pq8WFieFjnI= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.27.0 h1:ErKg/3iS1AKcTkf3yixlZ54f9U1rljCkQyEXWUnIUxc= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.27.0/go.mod h1:yAZHSGnqScoU556rBOVkwLze6WP5N+U11RHuWaGVxwY= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.51.0 h1:fYE9p3esPxA/C0rQ0AHhP0drtPXDRhaWiwg1DPqO7IU= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.51.0/go.mod h1:BnBReJLvVYx2CS/UHOgVz2BXKXD9wsQPxZug20nZhd0= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.51.0 h1:OqVGm6Ei3x5+yZmSJG1Mh2NwHvpVmZ08CB5qJhT9Nuk= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/cloudmock v0.51.0/go.mod h1:SZiPHWGOOk3bl8tkevxkoiwPgsIl6CwrWcbwjfHZpdM= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.51.0 h1:6/0iUd0xrnX7qt+mLNRwg5c0PGv8wpE8K90ryANQwMI= +github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.51.0/go.mod h1:otE2jQekW/PqXk1Awf5lmfokJx4uwuqcj1ab5SpGeW0= github.com/bahlo/generic-list-go v0.2.0 h1:5sz/EEAK+ls5wF+NeqDpk5+iNdMDXrh3z3nPnH1Wvgk= github.com/bahlo/generic-list-go v0.2.0/go.mod h1:2KvAjgMlE5NNynlg/5iLrrCCZ2+5xWbdbCW3pNTGyYg= github.com/buger/jsonparser v1.1.1 h1:2PnMjfWD7wBILjqQbt530v576A/cAbQvEW9gGIpYMUs= github.com/buger/jsonparser v1.1.1/go.mod h1:6RYKKt7H4d4+iWqouImQ9R2FZql3VbhNgx27UK13J/0= -github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU= -github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= -github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc= +github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs= +github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= +github.com/cncf/xds/go v0.0.0-20250121191232-2f005788dc42 h1:Om6kYQYDUk5wWbT0t0q6pvyM49i9XZAv9dDrkDA7gjk= +github.com/cncf/xds/go v0.0.0-20250121191232-2f005788dc42/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8= github.com/cockroachdb/errors v1.12.0 h1:d7oCs6vuIMUQRVbi6jWWWEJZahLCfJpnJSVobd1/sUo= github.com/cockroachdb/errors v1.12.0/go.mod h1:SvzfYNNBshAVbZ8wzNc/UPK3w1vf0dKDUP41ucAIf7g= github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b h1:r6VH0faHjZeQy818SGhaone5OnYfxFR/+AzdY3sf5aE= @@ -20,48 +43,47 @@ github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b/go.mod h1:Vz9D github.com/cockroachdb/redact v1.1.5 h1:u1PMllDkdFfPWaNGMyLD1+so+aq3uUItthCFqzwPJ30= github.com/cockroachdb/redact v1.1.5/go.mod h1:BVNblN9mBWFyMyqK1k3AAiSxhvhfK2oOZZ2lK+dpvRg= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= -github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= -github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= -github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= -github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98= -github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= +github.com/envoyproxy/go-control-plane v0.13.4 h1:zEqyPVyku6IvWCFwux4x9RxkLOMUL+1vC9xUFv5l2/M= +github.com/envoyproxy/go-control-plane v0.13.4/go.mod h1:kDfuBlDVsSj2MjrLEtRWtHlsWIFcGyB2RMO44Dc5GZA= +github.com/envoyproxy/go-control-plane/envoy v1.32.4 h1:jb83lalDRZSpPWW2Z7Mck/8kXZ5CQAFYVjQcdVIr83A= +github.com/envoyproxy/go-control-plane/envoy v1.32.4/go.mod h1:Gzjc5k8JcJswLjAx1Zm+wSYE20UrLtt7JZMWiWQXQEw= +github.com/envoyproxy/go-control-plane/ratelimit v0.1.0 h1:/G9QYbddjL25KvtKTv3an9lx6VBE2cnb8wp1vEGNYGI= +github.com/envoyproxy/go-control-plane/ratelimit v0.1.0/go.mod h1:Wk+tMFAFbCXaJPzVVHnPgRKdUdwW/KdbRt94AzgRee4= +github.com/envoyproxy/protoc-gen-validate v1.2.1 h1:DEo3O99U8j4hBFwbJfrz9VtgcDfUKS7KJ7spH3d86P8= +github.com/envoyproxy/protoc-gen-validate v1.2.1/go.mod h1:d/C80l/jxXLdfEIhX1W2TmLfsJ31lvEjwamM4DxlWXU= github.com/fatih/structs v1.1.0 h1:Q7juDM0QtcnhCpeyLGQKyg4TOIghuNXrkL32pHAUMxo= github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga6PJ7M= +github.com/felixge/httpsnoop v1.0.4 h1:NFTV2Zj1bL4mc9sqWACXbQFVBBg2W3GPvqp8/ESS2Wg= +github.com/felixge/httpsnoop v1.0.4/go.mod h1:m8KPJKqk1gH5J9DgRY2ASl2lWCfGKXixSwevea8zH2U= github.com/getsentry/sentry-go v0.27.0 h1:Pv98CIbtB3LkMWmXi4Joa5OOcwbmnX88sF5qbK3r3Ps= github.com/getsentry/sentry-go v0.27.0/go.mod h1:lc76E2QywIyW8WuBnwl8Lc4bkmQH4+w1gwTf25trprY= github.com/go-errors/errors v1.4.2 h1:J6MZopCL4uSllY1OfXM374weqZFFItUbrImctkmUxIA= github.com/go-errors/errors v1.4.2/go.mod h1:sIVyrIiJhuEF+Pj9Ebtd6P/rEYROXFi3BopGUQ5a5Og= +github.com/go-jose/go-jose/v4 v4.0.4 h1:VsjPI33J0SB9vQM6PLmNjoHqMQNGPiZ0rHL7Ni7Q6/E= +github.com/go-jose/go-jose/v4 v4.0.4/go.mod h1:NKb5HO1EZccyMpiZNbdUw/14tiXNyUJh188dfnMCAfc= +github.com/go-logr/logr v1.2.2/go.mod h1:jdQByPbusPIv2/zmleS9BjJVeZ6kBagPoEUsqbVz/1A= +github.com/go-logr/logr v1.4.2 h1:6pFjapn8bFcIbiKo3XT4j/BhANplGihG6tvd+8rYgrY= +github.com/go-logr/logr v1.4.2/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY= +github.com/go-logr/stdr v1.2.2 h1:hSWxHoqTgW2S2qGc0LTAI563KZ5YKYRhT3MFKZMbjag= +github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre4VKE= github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q= github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q= -github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= -github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc= -github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE= -github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc= -github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A= -github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= -github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= -github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8= -github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA= -github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs= -github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:WU3c8KckQ9AFe+yFwt9sWVRKCVIyN9cPHBJSNnbL67w= -github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= -github.com/golang/protobuf v1.4.1/go.mod h1:U8fpvMrcmy5pZrNK1lt4xCsGvpyWQ/VVv6QDs8UjoX8= -github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= -github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= -github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= -github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= -github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= -github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= -github.com/google/go-cmp v0.5.3/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= -github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= -github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= -github.com/google/s2a-go v0.1.8 h1:zZDs9gcbt9ZPLV0ndSyQk6Kacx2g/X+SKYovpnz3SMM= -github.com/google/s2a-go v0.1.8/go.mod h1:6iNWHTpQ+nfNRN5E00MSdfDwVesa8hhS32PhPO8deJA= -github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/googleapis/enterprise-certificate-proxy v0.3.4 h1:XYIDZApgAnrN1c855gTgghdIA6Stxb52D5RnLI1SLyw= -github.com/googleapis/enterprise-certificate-proxy v0.3.4/go.mod h1:YKe7cfqYXjKGpGvmSg28/fFvhNzinZQm8DGnaburhGA= +github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek= +github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/martian/v3 v3.3.3 h1:DIhPTQrbPkgs2yJYdXU/eNACCG5DVQjySNRNlflZ9Fc= +github.com/google/martian/v3 v3.3.3/go.mod h1:iEPrYcgCF7jA9OtScMFQyAlZZ4YXTKEtJ1E6RWzmBA0= +github.com/google/s2a-go v0.1.9 h1:LGD7gtMgezd8a/Xak7mEWL0PjoTQFvpRudN895yqKW0= +github.com/google/s2a-go v0.1.9/go.mod h1:YA0Ei2ZQL3acow2O62kdp9UlnvMmU7kA6Eutn0dXayM= +github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= +github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/googleapis/enterprise-certificate-proxy v0.3.6 h1:GW/XbdyBFQ8Qe+YAmFU9uHLo7OnF5tL52HFAgMmyrf4= +github.com/googleapis/enterprise-certificate-proxy v0.3.6/go.mod h1:MkHOF77EYAE7qfSuSS9PU6g4Nt4e11cnsDUowfwewLA= +github.com/googleapis/gax-go/v2 v2.14.2 h1:eBLnkZ9635krYIPD+ag1USrOAI0Nr0QYF3+/3GqO0k0= +github.com/googleapis/gax-go/v2 v2.14.2/go.mod h1:ON64QhlJkhVtSqp4v1uaK92VyZ2gmvDQsweuyLV+8+w= github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/h2non/gock v1.2.0 h1:K6ol8rfrRkUOefooBC8elXoaNGYkpp7y2qcxGG6BzUE= @@ -88,21 +110,21 @@ github.com/pingcap/errors v0.11.4/go.mod h1:Oi8TUi2kEtXXLMJk9l1cGmz20kV3TaQ0usTw github.com/pkg/diff v0.0.0-20210226163009-20ebb0f2a09e/go.mod h1:pJLUxLENpZxwdsKMEsNbx1VGcRFpLqf3715MtcvvzbA= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= +github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 h1:GFCKgmp0tecUJ0sJuv4pzYCqS9+RGSn52M3FUwPs+uo= +github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10/go.mod h1:t/avpk3KcrXxUnYOhZhMXJlSEyie6gQbtLq5NM3loB8= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= -github.com/rogpeppe/go-internal v1.9.0 h1:73kH8U+JUqXU8lRuOHeVHaa/SZPifC7BkcraZVejAe8= github.com/rogpeppe/go-internal v1.9.0/go.mod h1:WtVeX8xhTBvf0smdhujwtBcq4Qrzq/fJaraNFVN+nFs= +github.com/rogpeppe/go-internal v1.13.1 h1:KvO1DLK/DRN07sQ1LQKScxyZJuNnedQ5/wKSR38lUII= +github.com/rogpeppe/go-internal v1.13.1/go.mod h1:uMEvuHeurkdAXX61udpOXGD/AzZDWNMNyH2VO9fmH0o= github.com/samber/lo v1.51.0 h1:kysRYLbHy/MB7kQZf5DSN50JHmMsNEdeY24VzJFu7wI= github.com/samber/lo v1.51.0/go.mod h1:4+MXEGsJzbKGaUEQFKBq2xtfuznW9oz/WrgyzMzRoM0= -github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= -github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= -github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= +github.com/simonfrey/jsonl v0.0.0-20240904112901-935399b9a740 h1:CXJI+lliMiiEwzfgE8yt/38K0heYDgQ0L3f/3fxRnQU= +github.com/simonfrey/jsonl v0.0.0-20240904112901-935399b9a740/go.mod h1:G4w16caPmc6at7u4fmkj/8OAoOnM9mkmJr2fvL0vhaw= +github.com/spiffe/go-spiffe/v2 v2.5.0 h1:N2I01KCUkv1FAjZXJMwh95KK1ZIQLYbPfhaxw8WS0hE= +github.com/spiffe/go-spiffe/v2 v2.5.0/go.mod h1:P+NxobPc6wXhVtINNtFjNWGBTreew1GBUCwT2wPmb7g= github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY= github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA= -github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= -github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= -github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= @@ -119,53 +141,60 @@ github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/ github.com/wk8/go-ordered-map/v2 v2.1.8/go.mod h1:5nJHM5DyteebpVlHnWMV0rPz6Zp7+xBAnxjb1X5vnTw= github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= -go.opencensus.io v0.24.0 h1:y73uSU6J157QMP2kn2r30vwW1A2W2WFwSCGnAVxeaD0= -go.opencensus.io v0.24.0/go.mod h1:vNK8G9p7aAivkbmorf4v+7Hgx+Zs0yY+0fOtgBfjQKo= +github.com/zeebo/errs v1.4.0 h1:XNdoD/RRMKP7HD0UhJnIzUy74ISdGGxURlYG8HSWSfM= +github.com/zeebo/errs v1.4.0/go.mod h1:sgbWHsvVuTPHcqJJGQ1WhI5KbWlHYz+2+2C/LSEtCw4= +go.opentelemetry.io/auto/sdk v1.1.0 h1:cH53jehLUN6UFLY71z+NDOiNJqDdPRaXzTel0sJySYA= +go.opentelemetry.io/auto/sdk v1.1.0/go.mod h1:3wSPjt5PWp2RhlCcmmOial7AvC4DQqZb7a7wCow3W8A= +go.opentelemetry.io/contrib/detectors/gcp v1.36.0 h1:F7q2tNlCaHY9nMKHR6XH9/qkp8FktLnIcy6jJNyOCQw= +go.opentelemetry.io/contrib/detectors/gcp v1.36.0/go.mod h1:IbBN8uAIIx734PTonTPxAxnjc2pQTxWNkwfstZ+6H2k= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0 h1:x7wzEgXfnzJcHDwStJT+mxOz4etr2EcexjqhBvmoakw= +go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.60.0/go.mod h1:rg+RlpR5dKwaS95IyyZqj5Wd4E13lk/msnTS0Xl9lJM= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0 h1:sbiXRNDSWJOTobXh5HyQKjq6wUC5tNybqjIqDpAY4CU= +go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.60.0/go.mod h1:69uWxva0WgAA/4bu2Yy70SLDBwZXuQ6PbBpbsa5iZrQ= +go.opentelemetry.io/otel v1.36.0 h1:UumtzIklRBY6cI/lllNZlALOF5nNIzJVb16APdvgTXg= +go.opentelemetry.io/otel v1.36.0/go.mod h1:/TcFMXYjyRNh8khOAO9ybYkqaDBb/70aVwkNML4pP8E= +go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.36.0 h1:rixTyDGXFxRy1xzhKrotaHy3/KXdPhlWARrCgK+eqUY= +go.opentelemetry.io/otel/exporters/stdout/stdoutmetric v1.36.0/go.mod h1:dowW6UsM9MKbJq5JTz2AMVp3/5iW5I/TStsk8S+CfHw= +go.opentelemetry.io/otel/metric v1.36.0 h1:MoWPKVhQvJ+eeXWHFBOPoBOi20jh6Iq2CcCREuTYufE= +go.opentelemetry.io/otel/metric v1.36.0/go.mod h1:zC7Ks+yeyJt4xig9DEw9kuUFe5C3zLbVjV2PzT6qzbs= +go.opentelemetry.io/otel/sdk v1.36.0 h1:b6SYIuLRs88ztox4EyrvRti80uXIFy+Sqzoh9kFULbs= +go.opentelemetry.io/otel/sdk v1.36.0/go.mod h1:+lC+mTgD+MUWfjJubi2vvXWcVxyr9rmlshZni72pXeY= +go.opentelemetry.io/otel/sdk/metric v1.36.0 h1:r0ntwwGosWGaa0CrSt8cuNuTcccMXERFwHX4dThiPis= +go.opentelemetry.io/otel/sdk/metric v1.36.0/go.mod h1:qTNOhFDfKRwX0yXOqJYegL5WRaW376QbB7P4Pb0qva4= +go.opentelemetry.io/otel/trace v1.36.0 h1:ahxWNuqZjpdiFAyrIoQ4GIiAIhxAunQR6MUoKrsNd4w= +go.opentelemetry.io/otel/trace v1.36.0/go.mod h1:gQ+OnDZzrybY4k4seLzPAWNwVBBVlF2szhehOBB/tGA= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20191011191535-87dc89f01550/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto= -golang.org/x/crypto v0.36.0 h1:AnAEvhDddvBdpY+uR+MyHmuZzzNqXSe/GvuDeob5L34= -golang.org/x/crypto v0.36.0/go.mod h1:Y4J0ReaxCR1IMaabaSMugxJES1EpwhBHhv2bDHklZvc= -golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= -golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= -golang.org/x/lint v0.0.0-20190227174305-5b3e6a55c961/go.mod h1:wehouNa3lNwaWXcvxsM5YxQ5yQlVC4a0KAMCusXpPoU= -golang.org/x/lint v0.0.0-20190313153728-d0100b6bd8b3/go.mod h1:6SW0HCj/g11FgYtHlgUYUwCkIfeOF89ocIRzGO/8vkc= +golang.org/x/crypto v0.38.0 h1:jt+WWG8IZlBnVbomuhg2Mdq0+BBQaHbtqHEFEigjUV8= +golang.org/x/crypto v0.38.0/go.mod h1:MvrbAqul58NNYPKnOra203SB9vpuZW0e+RRZV+Ggqjw= golang.org/x/mod v0.2.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= golang.org/x/mod v0.3.0/go.mod h1:s0Qsj1ACt9ePp/hMypM3fl4fZqREWJwdYDEqhRiZZUA= -golang.org/x/net v0.0.0-20180724234803-3673e40ba225/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= -golang.org/x/net v0.0.0-20180826012351-8a410e7b638d/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= -golang.org/x/net v0.0.0-20190213061140-3a22650c66bd/go.mod h1:mL1N/T3taQHkDXs73rZJwtUhF3w3ftmwwsq0BUmARs4= -golang.org/x/net v0.0.0-20190311183353-d8887717615a/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20200226121028-0de0cce0169b/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s= golang.org/x/net v0.0.0-20201021035429-f5854403a974/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= -golang.org/x/net v0.0.0-20201110031124-69a78807bb2b/go.mod h1:sp8m0HH+o8qH0wwXwYZr8TS3Oi6o0r6Gce1SSxlDquU= -golang.org/x/net v0.38.0 h1:vRMAPTMaeGqVhG5QyLJHqNDwecKTomGeqbnfZyKlBI8= -golang.org/x/net v0.38.0/go.mod h1:ivrbrMbzFq5J41QOQh0siUuly180yBYtLp+CKbEaFx8= -golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U= -golang.org/x/sync v0.0.0-20180314180146-1d60e4601c6f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= -golang.org/x/sync v0.0.0-20181108010431-42b317875d0f/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= +golang.org/x/net v0.40.0 h1:79Xs7wF06Gbdcg4kdCCIQArK11Z1hr5POQ6+fIYHNuY= +golang.org/x/net v0.40.0/go.mod h1:y0hY0exeL2Pku80/zKK7tpntoX23cqL3Oa6njdgRtds= +golang.org/x/oauth2 v0.30.0 h1:dnDm7JmhM45NNpd8FDDeLhK6FwqbOf4MLCM9zb1BOHI= +golang.org/x/oauth2 v0.30.0/go.mod h1:B++QgG3ZKulg6sRPGD/mqlHQs5rB3Ml9erfeDY7xKlU= golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20190911185100-cd5d95a43a6e/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sync v0.15.0 h1:KWH3jNZsfyT6xfAfKiz6MRNmd46ByHDYaZ7KSkCtdW8= golang.org/x/sync v0.15.0/go.mod h1:1dzgHSNfp02xaA81J2MS99Qcpr2w7fw1gpm99rleRqA= -golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= -golang.org/x/sys v0.31.0 h1:ioabZlmFYtWhL+TRYpcnNlLwhyxaM9kWTDEmfnprqik= -golang.org/x/sys v0.31.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= +golang.org/x/sys v0.33.0 h1:q3i8TbbEz+JRD9ywIRlyRAQbM0qF7hu24q3teo2hbuw= +golang.org/x/sys v0.33.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k= golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ= golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.26.0 h1:P42AVeLghgTYr4+xUnTRKDMqpar+PtX7KWuNQL21L8M= golang.org/x/text v0.26.0/go.mod h1:QK15LZJUUQVJxhz7wXgxSy/CJaTFjd0G+YLonydOVQA= +golang.org/x/time v0.11.0 h1:/bpjEDfN9tkoN/ryeYHnv5hcMlc8ncjMcM4XBk5NWV0= +golang.org/x/time v0.11.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= -golang.org/x/tools v0.0.0-20190114222345-bf090417da8b/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= -golang.org/x/tools v0.0.0-20190226205152-f727befe758c/go.mod h1:9Yl7xja0Znq3iFh3HoIrodX9oNMXvdceNzlUR8zjMvY= -golang.org/x/tools v0.0.0-20190311212946-11955173bddd/go.mod h1:LCzVGOaR6xXOjkQ3onu1FJEFr0SW1gC7cKk1uF8kGRs= -golang.org/x/tools v0.0.0-20190524140312-2c0ae7006135/go.mod h1:RgjU9mgBXZiqYHBnxXauZ1Gv1EHHAz9KjViQ78xBX0Q= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE= golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= @@ -173,38 +202,22 @@ golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8T golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= -google.golang.org/appengine v1.1.0/go.mod h1:EbEs0AVv82hx2wNQdGPgUI5lhzA/G0D9YwlJXL52JkM= -google.golang.org/appengine v1.4.0/go.mod h1:xpcJRLb0r/rnEns0DIKYYv+WjYCduHsrkT7/EB5XEv4= +google.golang.org/api v0.235.0 h1:C3MkpQSRxS1Jy6AkzTGKKrpSCOd2WOGrezZ+icKSkKo= +google.golang.org/api v0.235.0/go.mod h1:QpeJkemzkFKe5VCE/PMv7GsUfn9ZF+u+q1Q7w6ckxTg= google.golang.org/genai v1.15.0 h1:zFaM+1JfGa0KCGDqrZdwVMucEu9n5AJEKkWcSPw0qro= google.golang.org/genai v1.15.0/go.mod h1:QPj5NGJw+3wEOHg+PrsWwJKvG6UC84ex5FR7qAYsN/M= -google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc= -google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc= -google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo= -google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1 h1:pPJltXNxVzT4pK9yD8vR9X75DaWYYmLGMsEvBfFQZzQ= -google.golang.org/genproto/googleapis/rpc v0.0.0-20240903143218-8af14fe29dc1/go.mod h1:UqMtugtsSgubUsoxbuAoiCXvqvErP7Gf0so0mK9tHxU= -google.golang.org/grpc v1.19.0/go.mod h1:mqu4LbDTu4XGKhr4mRzUsmM4RtVoemTSY81AxZiDr8c= -google.golang.org/grpc v1.23.0/go.mod h1:Y5yQAOtifL1yxbo5wqy6BxZv8vAUGQwXBOALyacEbxg= -google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQciAY= -google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk= -google.golang.org/grpc v1.33.2/go.mod h1:JMHMWHQWaTccqQQlmk3MJZS+GWXOdAesneDmEnv2fbc= -google.golang.org/grpc v1.66.2 h1:3QdXkuq3Bkh7w+ywLdLvM56cmGvQHUMZpiCzt6Rqaoo= -google.golang.org/grpc v1.66.2/go.mod h1:s3/l6xSSCURdVfAnL+TqCNMyTDAGN6+lZeVxnZR128Y= -google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8= -google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0= -google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM= -google.golang.org/protobuf v1.20.1-0.20200309200217-e05f789c0967/go.mod h1:A+miEFZTKqfCUM6K7xSMQL9OKL/b6hQv+e19PK+JZNE= -google.golang.org/protobuf v1.21.0/go.mod h1:47Nbq4nVaFHyn7ilMalzfO3qCViNmqZ2kzikPIcrTAo= -google.golang.org/protobuf v1.22.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= -google.golang.org/protobuf v1.23.0/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= -google.golang.org/protobuf v1.23.1-0.20200526195155-81db48ad09cc/go.mod h1:EGpADcykh3NcUnDUJcl1+ZksZNG86OlYog2l/sGQquU= -google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlbajtzgsN7c= -google.golang.org/protobuf v1.34.2 h1:6xV6lTsCfpGD21XK49h7MhtcApnLqkfYgPcdHftf6hg= -google.golang.org/protobuf v1.34.2/go.mod h1:qYOHts0dSfpeUzUFpOMr/WGzszTmLH+DiWniOlNbLDw= +google.golang.org/genproto v0.0.0-20250505200425-f936aa4a68b2 h1:1tXaIXCracvtsRxSBsYDiSBN0cuJvM7QYW+MrpIRY78= +google.golang.org/genproto v0.0.0-20250505200425-f936aa4a68b2/go.mod h1:49MsLSx0oWMOZqcpB3uL8ZOkAh1+TndpJ8ONoCBWiZk= +google.golang.org/genproto/googleapis/api v0.0.0-20250512202823-5a2f75b736a9 h1:WvBuA5rjZx9SNIzgcU53OohgZy6lKSus++uY4xLaWKc= +google.golang.org/genproto/googleapis/api v0.0.0-20250512202823-5a2f75b736a9/go.mod h1:W3S/3np0/dPWsWLi1h/UymYctGXaGBM2StwzD0y140U= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250512202823-5a2f75b736a9 h1:IkAfh6J/yllPtpYFU0zZN1hUPYdT0ogkBT/9hMxHjvg= +google.golang.org/genproto/googleapis/rpc v0.0.0-20250512202823-5a2f75b736a9/go.mod h1:qQ0YXyHHx3XkvlzUtpXDkS29lDSafHMZBAZDc03LQ3A= +google.golang.org/grpc v1.72.1 h1:HR03wO6eyZ7lknl75XlxABNVLLFc2PAb6mHlYh756mA= +google.golang.org/grpc v1.72.1/go.mod h1:wH5Aktxcg25y1I3w7H69nHfXdOG3UiadoBtjh3izSDM= +google.golang.org/protobuf v1.36.6 h1:z1NpPI8ku2WgiWnf+t9wTPsn6eP1L7ksHUlkfLvd9xY= +google.golang.org/protobuf v1.36.6/go.mod h1:jduwjTPXsFjZGTmRluh+L6NjiWu7pchiJ2/5YcXBHnY= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c h1:Hei/4ADfdWqJk1ZMxUNpqntNwaWcugrBjAiHlqqRiVk= gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c/go.mod h1:JHkPIbrfpd72SG/EVd6muEfDQjcINNoR0C8j2r3qZ4Q= -gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= -honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= -honnef.co/go/tools v0.0.0-20190523083050-ea95bdfd59fc/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4= diff --git a/internal/interface.go b/internal/interface.go index 61f1ead..8312bd6 100644 --- a/internal/interface.go +++ b/internal/interface.go @@ -1,6 +1,8 @@ package internal -import "net/http" +import ( + "net/http" +) // Adapter defines the interface for internal configuration and utility methods // that LLM providers can access from the main LlmAdapter. diff --git a/llms/aistudio/adapter_test.go b/llms/aistudio/adapter_test.go index 8665bb2..9ec64bc 100644 --- a/llms/aistudio/adapter_test.go +++ b/llms/aistudio/adapter_test.go @@ -38,10 +38,9 @@ func TestRequestAdapter(t *testing.T) { req := llmadapter.NewUntypedRequest(). WithModel("themodel") - contents, cfg, err := p.adaptRequest(llm, req, lo.FromPtr[RequestOptions](nil)) + contents, _, err := p.adaptRequest(llm, req, lo.FromPtr[RequestOptions](nil)) assert.Nil(t, err) - assert.Nil(t, cfg.SystemInstruction) assert.Len(t, contents, 0) }) diff --git a/llms/aistudio/aistudio.go b/llms/aistudio/aistudio.go index 0fd095d..7a73b30 100644 --- a/llms/aistudio/aistudio.go +++ b/llms/aistudio/aistudio.go @@ -14,6 +14,7 @@ import ( ) type AiStudio struct { + name string client *genai.Client history llmadapter.History[*genai.Content] @@ -22,6 +23,8 @@ type AiStudio struct { project string location string model *string + + bucket string } func (*AiStudio) RequestOptionsType() reflect.Type { @@ -40,6 +43,10 @@ func New(opts ...opt) (*AiStudio, error) { return &llm, nil } +func (p *AiStudio) SetName(name string) { + p.name = name +} + func (p *AiStudio) Init(adapter internal.Adapter) error { cfg := genai.ClientConfig{ Backend: genai.BackendGeminiAPI, @@ -109,7 +116,7 @@ func (p *AiStudio) adaptRequest(llm internal.Adapter, requester llmadapter.Reque } cfg := genai.GenerateContentConfig{ - CandidateCount: int32(lo.FromPtr(r.MaxCandidates)), + CandidateCount: int32(lo.CoalesceOrEmpty(lo.FromPtr(r.MaxCandidates), 1)), MaxOutputTokens: int32(lo.FromPtr(r.MaxTokens)), Temperature: internal.MaybeF64ToF32(r.Temperature), TopP: internal.MaybeF64ToF32(r.TopP), @@ -133,17 +140,19 @@ func (p *AiStudio) adaptRequest(llm internal.Adapter, requester llmadapter.Reque cfg.ResponseJsonSchema = lo.CoalesceOrEmpty(r.SchemaOverride, r.ResponseSchema) } - cfg.Tools = append(cfg.Tools, lo.MapToSlice(r.Tools, func(_ string, t internal.Tool) *genai.Tool { - return &genai.Tool{ - FunctionDeclarations: []*genai.FunctionDeclaration{ - { - Name: t.Name, - Description: t.Description, - ParametersJsonSchema: t.Parameters, + if len(r.Tools) > 0 { + cfg.Tools = append(cfg.Tools, lo.MapToSlice(r.Tools, func(_ string, t internal.Tool) *genai.Tool { + return &genai.Tool{ + FunctionDeclarations: []*genai.FunctionDeclaration{ + { + Name: t.Name, + Description: t.Description, + ParametersJsonSchema: t.Parameters, + }, }, - }, - } - })...) + } + })...) + } Messages: for _, msg := range r.Messages { diff --git a/llms/aistudio/batches.go b/llms/aistudio/batches.go new file mode 100644 index 0000000..e8f4cdb --- /dev/null +++ b/llms/aistudio/batches.go @@ -0,0 +1,218 @@ +package aistudio + +import ( + "bytes" + "context" + "fmt" + "io" + "time" + + "cloud.google.com/go/storage" + llmadapter "github.com/checkmarble/marble-llm-adapter" + "github.com/checkmarble/marble-llm-adapter/internal" + "github.com/cockroachdb/errors" + "github.com/samber/lo" + "github.com/simonfrey/jsonl" + "google.golang.org/genai" +) + +type BatchPayload struct { + Key string `json:"key"` + Request genai.InlinedRequest `json:"request"` +} + +func (p *AiStudio) SubmitBatch(ctx context.Context, llm internal.Adapter, reqs ...llmadapter.Requester) (*llmadapter.BatchPromise, error) { + if p.name == "" { + return nil, errors.New("batches can only be created on named providers") + } + + payload, err := p.createBatchInput(llm, reqs...) + if err != nil { + return nil, err + } + + req, err := p.uploadFile(ctx, payload) + if err != nil { + return nil, err + } + + src := genai.BatchJobSource{} + + switch p.backend { + case genai.BackendGeminiAPI: + src.FileName = req.Filename + case genai.BackendVertexAI: + src.Format = "jsonl" + src.GCSURI = []string{req.Filename} + } + + job, err := p.client.Batches.Create( + ctx, lo.FromPtr(lo.CoalesceOrEmpty(p.model, lo.ToPtr(llm.DefaultModel()))), + &src, + &genai.CreateBatchJobConfig{ + DisplayName: "Created on " + time.Now().Format(time.RFC3339), + Dest: &genai.BatchJobDestination{ + Format: "jsonl", + GCSURI: fmt.Sprintf("gs://%s/llm/outputs", p.bucket), + }, + }) + + if err != nil { + return nil, err + } + + return &llmadapter.BatchPromise{ + Provider: p, + ProviderName: p.name, + Id: job.Name, + }, nil +} + +func (p *AiStudio) createBatchInput(llm internal.Adapter, requesters ...llmadapter.Requester) (io.Reader, error) { + if p.name == "" { + return nil, errors.New("batches can only be created on named providers") + } + + var buf bytes.Buffer + + w := jsonl.NewWriter(&buf) + + for _, requester := range requesters { + id := requester.ToRequest().Id + + if requester.Error() != nil { + return nil, requester.Error() + } + + if id == "" { + return nil, errors.New("all requests in a batch must have an ID") + } + + model, ok := lo.Coalesce(requester.ToRequest().Model, p.model, lo.ToPtr(llm.DefaultModel())) + if !ok { + return nil, errors.New("no model was configured") + } + + opts := internal.CastProviderOptions[RequestOptions](requester.ProviderRequestOptions(p)) + + contents, cfg, err := p.adaptRequest(llm, requester, opts) + if err != nil { + return nil, err + } + + payload := BatchPayload{ + Key: id, + Request: genai.InlinedRequest{ + Model: *model, + Config: cfg, + Contents: contents, + }, + } + + if err := w.Write(payload); err != nil { + return nil, err + } + } + + return &buf, nil +} + +func (p *AiStudio) Check(ctx context.Context, pr *llmadapter.BatchPromise) (llmadapter.BatchStatus, error) { + job, err := p.client.Batches.Get(ctx, pr.Id, nil) + if err != nil { + return llmadapter.BatchError, err + } + + return adaptJobState(job.State), nil +} + +func (p *AiStudio) Wait(ctx context.Context, pr *llmadapter.BatchPromise) <-chan llmadapter.BatchWaitResponse { + ch := make(chan llmadapter.BatchWaitResponse) + + go func() { + for { + select { + case <-ctx.Done(): + close(ch) + return + + default: + job, err := p.client.Batches.Get(ctx, pr.Id, nil) + if err != nil { + ch <- llmadapter.BatchWaitResponse{Error: err} + close(ch) + return + } + + if !job.EndTime.IsZero() { + ch <- llmadapter.BatchWaitResponse{Status: adaptJobState(job.State)} + close(ch) + return + } + } + + time.Sleep(30 * time.Second) + } + }() + + return ch +} +func (p *AiStudio) uploadFile(ctx context.Context, r io.Reader) (*llmadapter.BatchRequest, error) { + if p.name == "" { + return nil, errors.New("batches can only be created on named providers") + } + + switch p.backend { + case genai.BackendGeminiAPI: + file, err := p.client.Files.Upload(ctx, r, &genai.UploadFileConfig{ + MIMEType: "jsonl", + }) + + if err != nil { + return nil, err + } + + return &llmadapter.BatchRequest{ + Provider: p, + Filename: file.Name, + }, nil + + case genai.BackendVertexAI: + client, err := storage.NewClient(ctx) + if err != nil { + return nil, err + } + + filename := fmt.Sprintf("llm/inputs/%s", "input.jsonl") + object := client.Bucket(p.bucket).Object(filename) + wr := object.NewWriter(ctx) + + if _, err := io.Copy(wr, r); err != nil { + return nil, err + } + if err := wr.Close(); err != nil { + return nil, err + } + + return &llmadapter.BatchRequest{ + Provider: p, + Filename: fmt.Sprintf("gs://%s/%s", p.bucket, filename), + }, nil + + default: + return nil, errors.New("invalid backend") + } +} + +func adaptJobState(state genai.JobState) llmadapter.BatchStatus { + switch state { + case genai.JobStatePending: + return llmadapter.BatchPending + case genai.JobStateCancelled, genai.JobStateSucceeded, genai.JobStatePartiallySucceeded: + return llmadapter.BatchFinished + case genai.JobStateRunning: + return llmadapter.BatchRunning + default: + return llmadapter.BatchPending + } +} diff --git a/llms/aistudio/options.go b/llms/aistudio/options.go index 8ab6c31..47ceebe 100644 --- a/llms/aistudio/options.go +++ b/llms/aistudio/options.go @@ -42,3 +42,9 @@ func WithDefaultModel(model string) opt { p.model = &model } } + +func WithBucket(bucket string) opt { + return func(p *AiStudio) { + p.bucket = bucket + } +} diff --git a/llms/openai/openai.go b/llms/openai/openai.go index 869b1d4..73705e2 100644 --- a/llms/openai/openai.go +++ b/llms/openai/openai.go @@ -16,6 +16,9 @@ import ( ) type OpenAi struct { + llmadapter.BatchUnsupported + + name string client openai.Client history llmadapter.History[openai.ChatCompletionMessageParamUnion] @@ -40,6 +43,10 @@ func New(opts ...Opt) (*OpenAi, error) { return &llm, nil } +func (p *OpenAi) SetName(name string) { + p.name = name +} + func (p *OpenAi) Init(llm internal.Adapter) error { opts := []option.RequestOption{ option.WithAPIKey(p.apiKey), diff --git a/mock_provider_test.go b/mock_provider_test.go index 5ba14eb..6c66fde 100644 --- a/mock_provider_test.go +++ b/mock_provider_test.go @@ -15,6 +15,7 @@ type MockMessage struct { type MockProvider struct { mock.Mock + BatchUnsupported History History[MockMessage] } @@ -31,6 +32,8 @@ func NewMockProvider() *MockProvider { } } +func (p *MockProvider) SetName(name string) {} + func (p *MockProvider) Init(llm internal.Adapter) error { return p.Called(llm).Error(0) } diff --git a/options.go b/options.go index 03421b4..4ece224 100644 --- a/options.go +++ b/options.go @@ -17,6 +17,8 @@ func WithDefaultProvider(provider Llm) llmOption { // The first one to be registered will become the default, unless a default was // already or is defined later with `SetDefaultProvider`. func WithProvider(name string, provider Llm) llmOption { + provider.SetName(name) + return func(llm *LlmAdapter) { llm.providers[name] = provider diff --git a/request.go b/request.go index b2f8fb4..46bed38 100644 --- a/request.go +++ b/request.go @@ -35,6 +35,7 @@ const ( // // Used internally to abstract over request types across packages. type Requester interface { + Error() error // ToRequest unwraps the actual request. ToRequest() innerRequest // ProviderRequestOptions extracts the provider-specific configuration @@ -62,6 +63,7 @@ type Message struct { // innerRequest represents the actual request to be sent to the provider, before // being adapted for it. type innerRequest struct { + Id string ThreadId *ThreadId SkipSaveInput bool SkipSaveOutput bool @@ -185,6 +187,12 @@ func (r Request[T]) Do(ctx context.Context, llm *LlmAdapter) (*Response[T], erro }, nil } +func (r Request[T]) WithId(id string) Request[T] { + r.innerRequest.Id = id + + return r +} + func (r Request[T]) WithProvider(name string) Request[T] { r.provider = &name @@ -503,6 +511,10 @@ func (r Request[T]) WithTopP(topp float64) Request[T] { // Request[T] implementation of Requester. +func (r Request[T]) Error() error { + return r.err +} + func (r Request[T]) ToRequest() innerRequest { return r.innerRequest } From 692da1caa8e113ee9db75adbb12115493973e9bb Mon Sep 17 00:00:00 2001 From: Antoine Popineau Date: Mon, 21 Jul 2025 09:09:45 +0200 Subject: [PATCH 2/2] Move things around to make them generic. --- adapter.go | 12 ++++---- batches.go | 62 +++++++++++++++++++++++++++++++--------- examples/batch/main.go | 15 ++++++---- llms/aistudio/batches.go | 42 ++++++++++++--------------- 4 files changed, 83 insertions(+), 48 deletions(-) diff --git a/adapter.go b/adapter.go index a13943a..8a19023 100644 --- a/adapter.go +++ b/adapter.go @@ -37,9 +37,9 @@ type Llm interface { // when processing custom request options. RequestOptionsType() reflect.Type - SubmitBatch(context.Context, internal.Adapter, ...Requester) (*BatchPromise, error) - Check(context.Context, *BatchPromise) (BatchStatus, error) - Wait(ctx context.Context, pr *BatchPromise) <-chan BatchWaitResponse + SubmitBatch(context.Context, internal.Adapter, ...Requester) (*UntypedBatchPromise, error) + Check(context.Context, *UntypedBatchPromise) (BatchStatus, error) + Wait(ctx context.Context, pr *UntypedBatchPromise) <-chan BatchWaitResponse } // LlmAdapter is the main entrypoint for interacting with different LLM providers. @@ -114,7 +114,7 @@ func (llm *LlmAdapter) GetProvider(requestProvider *string) (Llm, error) { return provider, nil } -func (llm *LlmAdapter) SubmitBatch(ctx context.Context, providerName string, reqs ...Requester) (*BatchPromise, error) { +func (llm *LlmAdapter) SubmitBatch(ctx context.Context, providerName string, reqs ...Requester) (*UntypedBatchPromise, error) { p, ok := llm.providers[providerName] if !ok { return nil, errors.Newf("unknown provider '%s'", providerName) @@ -123,13 +123,13 @@ func (llm *LlmAdapter) SubmitBatch(ctx context.Context, providerName string, req return p.SubmitBatch(ctx, llm, reqs...) } -func (llm *LlmAdapter) BatchPromise(providerName string, id string) (*BatchPromise, error) { +func (llm *LlmAdapter) BatchPromise(providerName string, id string) (*UntypedBatchPromise, error) { provider, ok := llm.providers[providerName] if !ok { return nil, errors.New("cannot find the provider that created this promise") } - return &BatchPromise{ + return &UntypedBatchPromise{ ProviderName: providerName, Provider: provider, Id: id, diff --git a/batches.go b/batches.go index 4b4b526..88001dd 100644 --- a/batches.go +++ b/batches.go @@ -20,38 +20,74 @@ const ( type BatchUnsupported struct{} -func (BatchUnsupported) SubmitBatch(ctx context.Context, llm internal.Adapter, reqs ...Requester) (*BatchPromise, error) { +func (BatchUnsupported) SubmitBatch(ctx context.Context, llm internal.Adapter, reqs ...Requester) (*UntypedBatchPromise, error) { return nil, errors.New("provider does not support batch mode") } -func (BatchUnsupported) Check(context.Context, *BatchPromise) (BatchStatus, error) { +func (BatchUnsupported) Check(context.Context, *UntypedBatchPromise) (BatchStatus, error) { return BatchError, errors.New("provider does not support batch mode") } -func (BatchUnsupported) Wait(ctx context.Context, pr *BatchPromise) <-chan BatchWaitResponse { +func (BatchUnsupported) Wait(ctx context.Context, pr *UntypedBatchPromise) <-chan BatchWaitResponse { return nil } -type BatchRequest struct { - Provider Llm - Filename string +type Batch[T any] struct { + Requests []Request[T] +} + +func (b Batch[T]) Batch(ctx context.Context, llm *LlmAdapter, providerName string) (*BatchPromise[T], error) { + requesters := make([]Requester, len(b.Requests)) + + for idx, r := range b.Requests { + requesters[idx] = Requester(r) + } + + promise, err := llm.SubmitBatch(ctx, providerName, requesters...) + + if err != nil { + return nil, err + } + + return &BatchPromise[T]{promise}, nil } -type BatchPromise struct { +type UntypedBatchPromise struct { Provider Llm ProviderName string Id string } -func (p *BatchPromise) Check(ctx context.Context) (BatchStatus, error) { - return p.Provider.Check(ctx, p) +type BatchPromise[T any] struct { + *UntypedBatchPromise } -func (p *BatchPromise) Wait(ctx context.Context) <-chan BatchWaitResponse { - return p.Provider.Wait(ctx, p) +func (p BatchPromise[T]) Check(ctx context.Context) (BatchStatus, error) { + return p.Provider.Check(ctx, p.UntypedBatchPromise) +} + +func (p BatchPromise[T]) Wait(ctx context.Context) (map[string]Response[T], error) { + inners := <-p.Provider.Wait(ctx, p.UntypedBatchPromise) + + if inners.Error != nil { + return nil, inners.Error + } + + responses := make(map[string]Response[T], len(inners.Responses)) + + for id, resp := range inners.Responses { + responses[id] = Response[T]{ + InnerResponse: resp, + } + } + + return responses, nil } type BatchWaitResponse struct { - Status BatchStatus - Error error + Status BatchStatus + Filename string + Error error + + Responses map[string]InnerResponse } diff --git a/examples/batch/main.go b/examples/batch/main.go index 0342e72..93b2518 100644 --- a/examples/batch/main.go +++ b/examples/batch/main.go @@ -27,19 +27,22 @@ func main() { llmadapter.WithDefaultModel("gemini-2.5-flash"), ) - reqs := []llmadapter.Requester{ - llmadapter.NewUntypedRequest().WithProvider("vertex").WithId("how").WithText(llmadapter.RoleUser, "How are you?"), - llmadapter.NewUntypedRequest().WithProvider("vertex").WithId("addition").WithText(llmadapter.RoleUser, "What is 1 + 1?"), + reqs := llmadapter.Batch[string]{ + Requests: []llmadapter.Request[string]{ + llmadapter.NewUntypedRequest().WithProvider("vertex").WithId("how").WithText(llmadapter.RoleUser, "How are you?"), + llmadapter.NewUntypedRequest().WithProvider("vertex").WithId("addition").WithText(llmadapter.RoleUser, "What is 1 + 1?"), + }, } - promise, err := llm.SubmitBatch(ctx, "vertex", reqs...) + // promise, err := llm.SubmitBatch(ctx, "vertex", []llmadapter.Requester(reqs)...) + promise, err := reqs.Batch(ctx, llm, "vertex") if err != nil { log.Fatal(err) } - result := <-promise.Wait(ctx) + result, err := promise.Wait(ctx) - if result.Error != nil { + if err != nil { log.Fatal(err) } diff --git a/llms/aistudio/batches.go b/llms/aistudio/batches.go index e8f4cdb..ce7c9b0 100644 --- a/llms/aistudio/batches.go +++ b/llms/aistudio/batches.go @@ -21,7 +21,7 @@ type BatchPayload struct { Request genai.InlinedRequest `json:"request"` } -func (p *AiStudio) SubmitBatch(ctx context.Context, llm internal.Adapter, reqs ...llmadapter.Requester) (*llmadapter.BatchPromise, error) { +func (p *AiStudio) SubmitBatch(ctx context.Context, llm internal.Adapter, reqs ...llmadapter.Requester) (*llmadapter.UntypedBatchPromise, error) { if p.name == "" { return nil, errors.New("batches can only be created on named providers") } @@ -31,7 +31,7 @@ func (p *AiStudio) SubmitBatch(ctx context.Context, llm internal.Adapter, reqs . return nil, err } - req, err := p.uploadFile(ctx, payload) + filename, err := p.uploadFile(ctx, payload) if err != nil { return nil, err } @@ -40,10 +40,10 @@ func (p *AiStudio) SubmitBatch(ctx context.Context, llm internal.Adapter, reqs . switch p.backend { case genai.BackendGeminiAPI: - src.FileName = req.Filename + src.FileName = filename case genai.BackendVertexAI: src.Format = "jsonl" - src.GCSURI = []string{req.Filename} + src.GCSURI = []string{filename} } job, err := p.client.Batches.Create( @@ -61,7 +61,7 @@ func (p *AiStudio) SubmitBatch(ctx context.Context, llm internal.Adapter, reqs . return nil, err } - return &llmadapter.BatchPromise{ + return &llmadapter.UntypedBatchPromise{ Provider: p, ProviderName: p.name, Id: job.Name, @@ -117,7 +117,7 @@ func (p *AiStudio) createBatchInput(llm internal.Adapter, requesters ...llmadapt return &buf, nil } -func (p *AiStudio) Check(ctx context.Context, pr *llmadapter.BatchPromise) (llmadapter.BatchStatus, error) { +func (p *AiStudio) Check(ctx context.Context, pr *llmadapter.UntypedBatchPromise) (llmadapter.BatchStatus, error) { job, err := p.client.Batches.Get(ctx, pr.Id, nil) if err != nil { return llmadapter.BatchError, err @@ -126,7 +126,7 @@ func (p *AiStudio) Check(ctx context.Context, pr *llmadapter.BatchPromise) (llma return adaptJobState(job.State), nil } -func (p *AiStudio) Wait(ctx context.Context, pr *llmadapter.BatchPromise) <-chan llmadapter.BatchWaitResponse { +func (p *AiStudio) Wait(ctx context.Context, pr *llmadapter.UntypedBatchPromise) <-chan llmadapter.BatchWaitResponse { ch := make(chan llmadapter.BatchWaitResponse) go func() { @@ -145,7 +145,9 @@ func (p *AiStudio) Wait(ctx context.Context, pr *llmadapter.BatchPromise) <-chan } if !job.EndTime.IsZero() { - ch <- llmadapter.BatchWaitResponse{Status: adaptJobState(job.State)} + // TODO: get the file from GCS / Files API and build InnerResponses with it + + ch <- llmadapter.BatchWaitResponse{Status: adaptJobState(job.State), Filename: job.Dest.GCSURI} close(ch) return } @@ -157,9 +159,9 @@ func (p *AiStudio) Wait(ctx context.Context, pr *llmadapter.BatchPromise) <-chan return ch } -func (p *AiStudio) uploadFile(ctx context.Context, r io.Reader) (*llmadapter.BatchRequest, error) { +func (p *AiStudio) uploadFile(ctx context.Context, r io.Reader) (string, error) { if p.name == "" { - return nil, errors.New("batches can only be created on named providers") + return "", errors.New("batches can only be created on named providers") } switch p.backend { @@ -169,18 +171,15 @@ func (p *AiStudio) uploadFile(ctx context.Context, r io.Reader) (*llmadapter.Bat }) if err != nil { - return nil, err + return "", err } - return &llmadapter.BatchRequest{ - Provider: p, - Filename: file.Name, - }, nil + return file.Name, nil case genai.BackendVertexAI: client, err := storage.NewClient(ctx) if err != nil { - return nil, err + return "", err } filename := fmt.Sprintf("llm/inputs/%s", "input.jsonl") @@ -188,19 +187,16 @@ func (p *AiStudio) uploadFile(ctx context.Context, r io.Reader) (*llmadapter.Bat wr := object.NewWriter(ctx) if _, err := io.Copy(wr, r); err != nil { - return nil, err + return "", err } if err := wr.Close(); err != nil { - return nil, err + return "", err } - return &llmadapter.BatchRequest{ - Provider: p, - Filename: fmt.Sprintf("gs://%s/%s", p.bucket, filename), - }, nil + return fmt.Sprintf("gs://%s/%s", p.bucket, filename), nil default: - return nil, errors.New("invalid backend") + return "", errors.New("invalid backend") } }