feat: add retry to chat completion client
This commit is contained in:
@ -2,24 +2,25 @@ package setup
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"forge.cadoles.com/wpetit/clearcase/internal/config"
|
||||
"forge.cadoles.com/wpetit/clearcase/internal/core/service"
|
||||
"github.com/bornholm/genai/llm"
|
||||
"github.com/bornholm/genai/llm/provider"
|
||||
"github.com/bornholm/genai/llm/retry"
|
||||
"github.com/pkg/errors"
|
||||
|
||||
_ "github.com/bornholm/genai/llm/provider/openai"
|
||||
_ "github.com/bornholm/genai/llm/provider/all"
|
||||
"github.com/bornholm/genai/llm/provider/openrouter"
|
||||
_ "github.com/bornholm/genai/llm/provider/openrouter"
|
||||
)
|
||||
|
||||
func NewForgeManagerFromConfig(ctx context.Context, conf *config.Config) (*service.ForgeManager, error) {
|
||||
client, err := provider.Create(ctx,
|
||||
provider.WithConfig(&provider.Config{
|
||||
provider.WithChatCompletionOptions(provider.ClientOptions{
|
||||
Provider: provider.Name(conf.LLM.Provider.Name),
|
||||
BaseURL: conf.LLM.Provider.BaseURL,
|
||||
Key: conf.LLM.Provider.Key,
|
||||
APIKey: conf.LLM.Provider.Key,
|
||||
Model: conf.LLM.Provider.Model,
|
||||
}))
|
||||
if err != nil {
|
||||
@ -38,6 +39,8 @@ func NewForgeManagerFromConfig(ctx context.Context, conf *config.Config) (*servi
|
||||
}
|
||||
}
|
||||
|
||||
client = retry.Wrap(client, time.Second, 5)
|
||||
|
||||
forgeFactories, err := getForgeFactories(conf)
|
||||
if err != nil {
|
||||
return nil, errors.Wrap(err, "could not get forge factories")
|
||||
@ -53,8 +56,14 @@ type extendedContextClient struct {
|
||||
extend func(ctx context.Context) context.Context
|
||||
}
|
||||
|
||||
// Embeddings implements llm.Client.
|
||||
func (c *extendedContextClient) Embeddings(ctx context.Context, input string, funcs ...llm.EmbeddingsOptionFunc) (llm.EmbeddingsResponse, error) {
|
||||
ctx = c.extend(ctx)
|
||||
return c.client.Embeddings(ctx, input, funcs...)
|
||||
}
|
||||
|
||||
// ChatCompletion implements llm.Client.
|
||||
func (c *extendedContextClient) ChatCompletion(ctx context.Context, funcs ...llm.ChatCompletionOptionFunc) (llm.CompletionResponse, error) {
|
||||
func (c *extendedContextClient) ChatCompletion(ctx context.Context, funcs ...llm.ChatCompletionOptionFunc) (llm.ChatCompletionResponse, error) {
|
||||
ctx = c.extend(ctx)
|
||||
return c.client.ChatCompletion(ctx, funcs...)
|
||||
}
|
||||
|
Reference in New Issue
Block a user