77 lines
2.3 KiB
Go
77 lines
2.3 KiB
Go
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/all"
|
|
"github.com/bornholm/genai/llm/provider/openrouter"
|
|
)
|
|
|
|
func NewForgeManagerFromConfig(ctx context.Context, conf *config.Config) (*service.ForgeManager, error) {
|
|
client, err := provider.Create(ctx,
|
|
provider.WithChatCompletionOptions(provider.ClientOptions{
|
|
Provider: provider.Name(conf.LLM.Provider.Name),
|
|
BaseURL: conf.LLM.Provider.BaseURL,
|
|
APIKey: conf.LLM.Provider.Key,
|
|
Model: conf.LLM.Provider.Model,
|
|
}))
|
|
if err != nil {
|
|
return nil, errors.Wrapf(err, "could not create llm client '%s'", conf.LLM.Provider.Name)
|
|
}
|
|
|
|
if conf.LLM.Provider.Name == string(openrouter.Name) {
|
|
client = &extendedContextClient{
|
|
client: client,
|
|
extend: func(ctx context.Context) context.Context {
|
|
// Automatically "compress" prompts when using openrouter
|
|
// See https://openrouter.ai/docs/features/message-transforms
|
|
ctx = openrouter.WithTransforms(ctx, []string{"middle-out"})
|
|
return ctx
|
|
},
|
|
}
|
|
}
|
|
|
|
client = retry.Wrap(client, time.Second, 5)
|
|
|
|
forgeFactories, err := getForgeFactories(conf)
|
|
if err != nil {
|
|
return nil, errors.Wrap(err, "could not get forge factories")
|
|
}
|
|
|
|
forgeManager := service.NewForgeManager(client, forgeFactories...)
|
|
|
|
return forgeManager, nil
|
|
}
|
|
|
|
type extendedContextClient struct {
|
|
client llm.Client
|
|
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.ChatCompletionResponse, error) {
|
|
ctx = c.extend(ctx)
|
|
return c.client.ChatCompletion(ctx, funcs...)
|
|
}
|
|
|
|
// Model implements llm.Client.
|
|
func (c *extendedContextClient) Model() string {
|
|
return c.Model()
|
|
}
|
|
|
|
var _ llm.Client = &extendedContextClient{}
|