41 lines
833 B
Go
41 lines
833 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) {
|
|
pool := c.manager.getPool(req.Target.BackendID)
|
|
if pool == nil {
|
|
return c.next.Generate(ctx, req)
|
|
}
|
|
if err := pool.acquire(ctx); err != nil {
|
|
return nil, err
|
|
}
|
|
defer pool.releaseActive()
|
|
return c.next.Generate(ctx, req)
|
|
}
|