41 lines
844 B
Go
41 lines
844 B
Go
package capacity
|
|
|
|
import (
|
|
"context"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/llm"
|
|
)
|
|
|
|
type client struct {
|
|
manager *Manager
|
|
next llm.Client
|
|
}
|
|
|
|
// NewClient wraps next with configured active-generation limits. A nil manager
|
|
// leaves next unchanged.
|
|
func NewClient(manager *Manager, next llm.Client) llm.Client {
|
|
if manager == nil {
|
|
return next
|
|
}
|
|
return &client{
|
|
manager: manager,
|
|
next: next,
|
|
}
|
|
}
|
|
|
|
func (c *client) Generate(
|
|
ctx context.Context,
|
|
req domain.GenerateRequest,
|
|
) (*domain.GenerateResponse, error) {
|
|
if c.manager.getPool(req.Target.BackendID) == nil {
|
|
return c.next.Generate(ctx, req)
|
|
}
|
|
release, err := c.manager.Acquire(ctx, req.Target.BackendID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer release()
|
|
return c.next.Generate(ctx, req)
|
|
}
|