feat(infra/llm): LLMRegistry with endpoint cache + invalidation

This commit is contained in:
2026-05-04 09:48:03 +08:00
parent a08a7249f0
commit 9bc9873e3c
2 changed files with 162 additions and 0 deletions
+84
View File
@@ -0,0 +1,84 @@
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
}
// 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)
}
}
+78
View File
@@ -0,0 +1,78 @@
package llm
import (
"context"
"errors"
"sync/atomic"
"testing"
"github.com/google/uuid"
"github.com/stretchr/testify/require"
)
type fakeEndpoint struct {
ID uuid.UUID
Provider string
BaseURL string
APIKey string
}
type fakeLoader struct {
endpoints map[uuid.UUID]fakeEndpoint
calls int32
}
func (l *fakeLoader) Load(ctx context.Context, id uuid.UUID) (Endpoint, error) {
atomic.AddInt32(&l.calls, 1)
ep, ok := l.endpoints[id]
if !ok {
return Endpoint{}, errors.New("not found")
}
return Endpoint{ID: ep.ID, Provider: ep.Provider, BaseURL: ep.BaseURL, APIKey: ep.APIKey}, nil
}
func TestRegistry_CachesByEndpointID(t *testing.T) {
id := uuid.New()
loader := &fakeLoader{endpoints: map[uuid.UUID]fakeEndpoint{
id: {ID: id, Provider: "anthropic", BaseURL: "http://localhost", APIKey: "k"},
}}
r := NewRegistry(loader)
c1, err := r.GetClient(context.Background(), id)
require.NoError(t, err)
c2, err := r.GetClient(context.Background(), id)
require.NoError(t, err)
require.Same(t, c1, c2)
require.Equal(t, int32(1), atomic.LoadInt32(&loader.calls))
}
func TestRegistry_InvalidateForcesReload(t *testing.T) {
id := uuid.New()
loader := &fakeLoader{endpoints: map[uuid.UUID]fakeEndpoint{
id: {ID: id, Provider: "anthropic", BaseURL: "http://localhost", APIKey: "k"},
}}
r := NewRegistry(loader)
_, err := r.GetClient(context.Background(), id)
require.NoError(t, err)
r.Invalidate(id)
_, err = r.GetClient(context.Background(), id)
require.NoError(t, err)
require.Equal(t, int32(2), atomic.LoadInt32(&loader.calls))
}
func TestRegistry_UnknownProviderReturnsError(t *testing.T) {
id := uuid.New()
loader := &fakeLoader{endpoints: map[uuid.UUID]fakeEndpoint{
id: {ID: id, Provider: "claude-3-but-typoed"},
}}
r := NewRegistry(loader)
_, err := r.GetClient(context.Background(), id)
require.Error(t, err)
}
func TestRegistry_LoaderErrorPropagates(t *testing.T) {
loader := &fakeLoader{endpoints: map[uuid.UUID]fakeEndpoint{}}
r := NewRegistry(loader)
_, err := r.GetClient(context.Background(), uuid.New())
require.Error(t, err)
}