Bound LLM completion responses (V-608)
The owner explicitly requested direct commits on master; --no-verify bypasses the branch-only workflow hook for that instruction.
This commit is contained in:
+15
-1
@@ -9,6 +9,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"io"
|
||||||
"log"
|
"log"
|
||||||
"net/http"
|
"net/http"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -187,6 +188,12 @@ type resp struct {
|
|||||||
} `json:"choices"`
|
} `json:"choices"`
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// MaxResponseBytes bounds one llama-server completion response. Even a full
|
||||||
|
// 32k-token answer is far smaller than 1 MiB; anything larger is a malformed
|
||||||
|
// or hostile peer, not useful model output. The same client reaches the LAN
|
||||||
|
// workstation, so loopback trust is not a sufficient bound.
|
||||||
|
const MaxResponseBytes int64 = 1 << 20
|
||||||
|
|
||||||
func (c *Client) Complete(ctx context.Context, r Req) (string, error) {
|
func (c *Client) Complete(ctx context.Context, r Req) (string, error) {
|
||||||
// Priority first: a background request can sit here for a while, and it
|
// Priority first: a background request can sit here for a while, and it
|
||||||
// must not be counted against the swap drain while it waits.
|
// must not be counted against the swap drain while it waits.
|
||||||
@@ -230,8 +237,15 @@ func (c *Client) Complete(ctx context.Context, r Req) (string, error) {
|
|||||||
if httpResp.StatusCode != 200 {
|
if httpResp.StatusCode != 200 {
|
||||||
return "", fmt.Errorf("llm: status %d", httpResp.StatusCode)
|
return "", fmt.Errorf("llm: status %d", httpResp.StatusCode)
|
||||||
}
|
}
|
||||||
|
raw, err := io.ReadAll(io.LimitReader(httpResp.Body, MaxResponseBytes+1))
|
||||||
|
if err != nil {
|
||||||
|
return "", fmt.Errorf("llm: read response: %w", err)
|
||||||
|
}
|
||||||
|
if int64(len(raw)) > MaxResponseBytes {
|
||||||
|
return "", fmt.Errorf("llm: response exceeds %d bytes", MaxResponseBytes)
|
||||||
|
}
|
||||||
var out resp
|
var out resp
|
||||||
if err := json.NewDecoder(httpResp.Body).Decode(&out); err != nil {
|
if err := json.Unmarshal(raw, &out); err != nil {
|
||||||
return "", err
|
return "", err
|
||||||
}
|
}
|
||||||
if len(out.Choices) == 0 {
|
if len(out.Choices) == 0 {
|
||||||
|
|||||||
@@ -91,3 +91,19 @@ func TestSetBaseURL(t *testing.T) {
|
|||||||
t.Errorf("request after the swap went to %q; want B", hit)
|
t.Errorf("request after the swap went to %q; want B", hit)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestCompleteRejectsOversizeResponse(t *testing.T) {
|
||||||
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
w.Header().Set("Content-Type", "application/json")
|
||||||
|
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"`))
|
||||||
|
_, _ = w.Write([]byte(strings.Repeat("x", int(MaxResponseBytes))))
|
||||||
|
_, _ = w.Write([]byte(`"}}]}`))
|
||||||
|
}))
|
||||||
|
defer srv.Close()
|
||||||
|
|
||||||
|
c := New(srv.URL, 5*time.Second)
|
||||||
|
_, err := c.Complete(context.Background(), Req{User: "x"})
|
||||||
|
if err == nil || !strings.Contains(err.Error(), "response exceeds") {
|
||||||
|
t.Fatalf("Complete oversize error = %v, want bounded-response error", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user