fix(llm): read a bounded response body, 8 MiB on success and the error bound on failure (#53)
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
@@ -37,6 +37,21 @@ type Client struct {
|
||||
temp float64
|
||||
httpClient *http.Client
|
||||
usageHook func(model string, prompt, completion int)
|
||||
maxBody int64
|
||||
}
|
||||
|
||||
// defaultMaxResponseBytes bounds a success body. A completion is a few KB of text; this
|
||||
// leaves room for a long one and none for a gateway that streams forever (#53).
|
||||
const defaultMaxResponseBytes = 8 << 20
|
||||
|
||||
// WithMaxResponseBytes overrides the bound on a success response body. A non-positive n is
|
||||
// ignored (keeps the default).
|
||||
func WithMaxResponseBytes(n int64) Option {
|
||||
return func(c *Client) {
|
||||
if n > 0 {
|
||||
c.maxBody = n
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Option configures a Client at construction. Variadic so the existing 4-arg
|
||||
@@ -101,6 +116,7 @@ func New(baseURL, apiKey, model string, timeout time.Duration, opts ...Option) *
|
||||
maxTokens: defaultMaxTokens,
|
||||
temp: defaultTemperature,
|
||||
httpClient: &http.Client{Timeout: timeout},
|
||||
maxBody: defaultMaxResponseBytes,
|
||||
}
|
||||
for _, opt := range opts {
|
||||
opt(c)
|
||||
@@ -184,12 +200,17 @@ func (c *Client) Complete(ctx context.Context, system, user string) (string, err
|
||||
}
|
||||
defer resp.Body.Close() //nolint:errcheck
|
||||
|
||||
out, err := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
// Only the bounded head is ever used (HTTPError.Error truncates to it), so read no more.
|
||||
head, _ := io.ReadAll(io.LimitReader(resp.Body, maxErrorBody+1))
|
||||
return "", &HTTPError{Code: resp.StatusCode, Body: string(head)}
|
||||
}
|
||||
out, err := io.ReadAll(io.LimitReader(resp.Body, c.maxBody+1))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("read response: %w", err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", &HTTPError{Code: resp.StatusCode, Body: string(out)}
|
||||
if int64(len(out)) > c.maxBody {
|
||||
return "", fmt.Errorf("LLM response exceeds %d bytes", c.maxBody)
|
||||
}
|
||||
|
||||
var cr chatResponse
|
||||
|
||||
Reference in New Issue
Block a user