diff --git a/internal/infra/llm/registry.go b/internal/infra/llm/registry.go new file mode 100644 index 0000000..27cabc7 --- /dev/null +++ b/internal/infra/llm/registry.go @@ -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) + } +} diff --git a/internal/infra/llm/registry_test.go b/internal/infra/llm/registry_test.go new file mode 100644 index 0000000..405c5ae --- /dev/null +++ b/internal/infra/llm/registry_test.go @@ -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) +}