You've already forked agentic-coding-workflow
100 lines
2.7 KiB
Go
100 lines
2.7 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
// Endpoint 是 Registry 从外部 loader 拿到的端点描述（已解密 api_key）。
|
|
type Endpoint struct {
|
|
ID uuid.UUID
|
|
Provider string
|
|
BaseURL string
|
|
APIKey string
|
|
}
|
|
|
|
// EndpointLoader 由 chat repository / endpointService 实现，按 id 拉取端点
|
|
// 并完成 AES-GCM 解密。Registry 只看到明文 Endpoint。
|
|
type EndpointLoader interface {
|
|
Load(ctx context.Context, id uuid.UUID) (Endpoint, error)
|
|
}
|
|
|
|
// Registry 缓存 endpoint_id → LLMClient。endpoint mutation 后由 service 层调
|
|
// Invalidate 强制下次 GetClient 重新构造（带新 base_url / api_key）。
|
|
type Registry struct {
|
|
loader EndpointLoader
|
|
mu sync.Mutex
|
|
cache map[uuid.UUID]LLMClient
|
|
}
|
|
|
|
// NewRegistry 构造一个空缓存的 Registry。
|
|
func NewRegistry(loader EndpointLoader) *Registry {
|
|
return &Registry{loader: loader, cache: map[uuid.UUID]LLMClient{}}
|
|
}
|
|
|
|
// GetClient 返回该 endpoint 对应的 LLMClient。命中缓存直接返回，否则 Load + buildClient
|
|
// 后写入缓存。多 goroutine 并发触发 miss 时，仅一个写入生效（double-checked under lock）。
|
|
func (r *Registry) GetClient(ctx context.Context, id uuid.UUID) (LLMClient, error) {
|
|
r.mu.Lock()
|
|
if c, ok := r.cache[id]; ok {
|
|
r.mu.Unlock()
|
|
return c, nil
|
|
}
|
|
r.mu.Unlock()
|
|
|
|
ep, err := r.loader.Load(ctx, id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
c, err := buildClient(ep)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
if existing, ok := r.cache[id]; ok {
|
|
return existing, nil
|
|
}
|
|
r.cache[id] = c
|
|
return c, nil
|
|
}
|
|
|
|
// GetEmbedder 返回该 endpoint 对应的 Embedder（OpenAI / Gemini 实现；Anthropic
|
|
// 客户端虽实现接口但 Embed 返回 ErrEmbeddingsUnsupported）。底层复用 GetClient
|
|
// 的缓存与构造，再做类型断言。endpoint 的 client 未实现 Embedder 时返回错误。
|
|
func (r *Registry) GetEmbedder(ctx context.Context, id uuid.UUID) (Embedder, error) {
|
|
c, err := r.GetClient(ctx, id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
emb, ok := c.(Embedder)
|
|
if !ok {
|
|
return nil, fmt.Errorf("llm registry: endpoint %s provider %q does not support embeddings", id, c.Provider())
|
|
}
|
|
return emb, nil
|
|
}
|
|
|
|
// Invalidate 删除该 endpoint 的缓存条目。endpoint CRUD 后由 service 层调用。
|
|
func (r *Registry) Invalidate(id uuid.UUID) {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
delete(r.cache, id)
|
|
}
|
|
|
|
func buildClient(ep Endpoint) (LLMClient, error) {
|
|
switch ep.Provider {
|
|
case "anthropic":
|
|
return NewAnthropic(ep.APIKey, ep.BaseURL)
|
|
case "openai":
|
|
return NewOpenAI(ep.APIKey, ep.BaseURL)
|
|
case "gemini":
|
|
return NewGemini(ep.APIKey, ep.BaseURL)
|
|
default:
|
|
return nil, fmt.Errorf("llm registry: unknown provider %q", ep.Provider)
|
|
}
|
|
}
|