feat: add retry to chat completion client

This commit is contained in:
2025-08-13 11:17:03 +02:00
parent d68562b631
commit 3a7ff257b0
5 changed files with 73 additions and 47 deletions

View File

@ -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...)
}