Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
518e78db54 | ||
|
|
cb925618bc |
@@ -0,0 +1,103 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"time"
|
||||||
|
"unicode/utf8"
|
||||||
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
// ErrAuth means the vendor rejected the credentials.
|
||||||
|
ErrAuth = errors.New("llm: authentication failed")
|
||||||
|
// ErrRateLimited means the vendor throttled the request. The concrete
|
||||||
|
// error is a *RateLimitedError carrying the requested wait.
|
||||||
|
ErrRateLimited = errors.New("llm: rate limited")
|
||||||
|
// ErrTimeout means the request did not finish within its timeout.
|
||||||
|
ErrTimeout = errors.New("llm: request timed out")
|
||||||
|
// ErrBadOutput means the model replied but the reply is unusable, e.g.
|
||||||
|
// structured output that does not parse.
|
||||||
|
ErrBadOutput = errors.New("llm: unusable model output")
|
||||||
|
// ErrProvider means the vendor returned an unexpected HTTP status. The
|
||||||
|
// concrete error is a *ProviderError carrying the status and a redacted
|
||||||
|
// body excerpt.
|
||||||
|
ErrProvider = errors.New("llm: provider error")
|
||||||
|
)
|
||||||
|
|
||||||
|
// RateLimitedError is a throttled request. It matches ErrRateLimited.
|
||||||
|
type RateLimitedError struct {
|
||||||
|
// RetryAfter is how long the vendor asked callers to wait; 0 means it
|
||||||
|
// did not say.
|
||||||
|
RetryAfter time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
func (e *RateLimitedError) Error() string {
|
||||||
|
if e.RetryAfter == 0 {
|
||||||
|
return ErrRateLimited.Error()
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%s: retry after %v", ErrRateLimited, e.RetryAfter)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Is reports whether target is ErrRateLimited.
|
||||||
|
func (e *RateLimitedError) Is(target error) bool { return target == ErrRateLimited }
|
||||||
|
|
||||||
|
// MaxBodyExcerpt is the longest response-body excerpt, in bytes, a
|
||||||
|
// ProviderError keeps.
|
||||||
|
const MaxBodyExcerpt = 512
|
||||||
|
|
||||||
|
// redacted replaces each secret in a ProviderError body.
|
||||||
|
const redacted = "[REDACTED]"
|
||||||
|
|
||||||
|
// ProviderError is an unexpected HTTP status from the vendor. It matches
|
||||||
|
// ErrProvider. Build it with NewProviderError, which redacts the body.
|
||||||
|
type ProviderError struct {
|
||||||
|
status int
|
||||||
|
body string
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewProviderError returns the error for an unexpected HTTP status.
|
||||||
|
// secrets must hold the API key (and any other credential sent with the
|
||||||
|
// request). Every occurrence of each non-empty secret in body is replaced
|
||||||
|
// before the body is cut to MaxBodyExcerpt bytes, so no part of a secret
|
||||||
|
// survives. If secrets holds no non-empty value the body cannot be
|
||||||
|
// redacted, so no excerpt is kept at all.
|
||||||
|
func NewProviderError(status int, body []byte, secrets []string) *ProviderError {
|
||||||
|
s := string(body)
|
||||||
|
redactedAny := false
|
||||||
|
for _, secret := range secrets {
|
||||||
|
if secret != "" {
|
||||||
|
s = strings.ReplaceAll(s, secret, redacted)
|
||||||
|
redactedAny = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !redactedAny {
|
||||||
|
return &ProviderError{status: status}
|
||||||
|
}
|
||||||
|
s = strings.TrimSpace(s)
|
||||||
|
if len(s) > MaxBodyExcerpt {
|
||||||
|
cut := MaxBodyExcerpt
|
||||||
|
for cut > 0 && !utf8.RuneStart(s[cut]) {
|
||||||
|
cut--
|
||||||
|
}
|
||||||
|
s = s[:cut] + "…"
|
||||||
|
}
|
||||||
|
return &ProviderError{status: status, body: s}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Status returns the HTTP status code.
|
||||||
|
func (e *ProviderError) Status() int { return e.status }
|
||||||
|
|
||||||
|
// Body returns the redacted, truncated response-body excerpt; it is empty
|
||||||
|
// when the vendor sent no body or no secret was given to redact it with.
|
||||||
|
func (e *ProviderError) Body() string { return e.body }
|
||||||
|
|
||||||
|
func (e *ProviderError) Error() string {
|
||||||
|
if e.body == "" {
|
||||||
|
return fmt.Sprintf("%s: HTTP %d", ErrProvider, e.status)
|
||||||
|
}
|
||||||
|
return fmt.Sprintf("%s: HTTP %d: %s", ErrProvider, e.status, e.body)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Is reports whether target is ErrProvider.
|
||||||
|
func (e *ProviderError) Is(target error) bool { return target == ErrProvider }
|
||||||
@@ -0,0 +1,121 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSentinelErrorsMatchWhenWrapped(t *testing.T) {
|
||||||
|
for _, sentinel := range []error{ErrAuth, ErrTimeout, ErrBadOutput} {
|
||||||
|
wrapped := fmt.Errorf("openai: generate: %w", sentinel)
|
||||||
|
if !errors.Is(wrapped, sentinel) {
|
||||||
|
t.Errorf("errors.Is(%v, %v) = false", wrapped, sentinel)
|
||||||
|
}
|
||||||
|
for _, other := range []error{ErrAuth, ErrTimeout, ErrBadOutput, ErrRateLimited, ErrProvider} {
|
||||||
|
if other != sentinel && errors.Is(wrapped, other) {
|
||||||
|
t.Errorf("errors.Is(%v, %v) = true", wrapped, other)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRateLimitedError(t *testing.T) {
|
||||||
|
err := fmt.Errorf("anthropic: %w", &RateLimitedError{RetryAfter: 30 * time.Second})
|
||||||
|
if !errors.Is(err, ErrRateLimited) {
|
||||||
|
t.Fatal("errors.Is(err, ErrRateLimited) = false")
|
||||||
|
}
|
||||||
|
if errors.Is(err, ErrProvider) || errors.Is(err, ErrAuth) {
|
||||||
|
t.Fatal("rate-limited error matches an unrelated sentinel")
|
||||||
|
}
|
||||||
|
var rl *RateLimitedError
|
||||||
|
if !errors.As(err, &rl) {
|
||||||
|
t.Fatal("errors.As(err, *RateLimitedError) = false")
|
||||||
|
}
|
||||||
|
if rl.RetryAfter != 30*time.Second {
|
||||||
|
t.Fatalf("RetryAfter = %v, want 30s", rl.RetryAfter)
|
||||||
|
}
|
||||||
|
if got := err.Error(); got != "anthropic: llm: rate limited: retry after 30s" {
|
||||||
|
t.Fatalf("Error() = %q", got)
|
||||||
|
}
|
||||||
|
if got := (&RateLimitedError{}).Error(); got != "llm: rate limited" {
|
||||||
|
t.Fatalf("Error() without RetryAfter = %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderError(t *testing.T) {
|
||||||
|
err := fmt.Errorf("gemini: %w", NewProviderError(503, []byte(` {"error":"overloaded"} `), []string{"sk-test"}))
|
||||||
|
if !errors.Is(err, ErrProvider) {
|
||||||
|
t.Fatal("errors.Is(err, ErrProvider) = false")
|
||||||
|
}
|
||||||
|
if errors.Is(err, ErrRateLimited) || errors.Is(err, ErrAuth) {
|
||||||
|
t.Fatal("provider error matches an unrelated sentinel")
|
||||||
|
}
|
||||||
|
var pe *ProviderError
|
||||||
|
if !errors.As(err, &pe) {
|
||||||
|
t.Fatal("errors.As(err, *ProviderError) = false")
|
||||||
|
}
|
||||||
|
if pe.Status() != 503 || pe.Body() != `{"error":"overloaded"}` {
|
||||||
|
t.Fatalf("Status, Body = %d, %q", pe.Status(), pe.Body())
|
||||||
|
}
|
||||||
|
if got := err.Error(); got != `gemini: llm: provider error: HTTP 503: {"error":"overloaded"}` {
|
||||||
|
t.Fatalf("Error() = %q", got)
|
||||||
|
}
|
||||||
|
if got := NewProviderError(500, nil, []string{"sk-test"}).Error(); got != "llm: provider error: HTTP 500" {
|
||||||
|
t.Fatalf("Error() without body = %q", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderErrorRedactsSecrets(t *testing.T) {
|
||||||
|
const key = "sk-live-0123456789abcdef"
|
||||||
|
body := []byte(`{"error":"invalid key sk-live-0123456789abcdef","echo":"Bearer sk-live-0123456789abcdef"}`)
|
||||||
|
pe := NewProviderError(400, body, []string{key, ""})
|
||||||
|
for _, s := range []string{pe.Error(), pe.Body()} {
|
||||||
|
if strings.Contains(s, key) {
|
||||||
|
t.Fatalf("secret leaked: %q", s)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
want := `{"error":"invalid key [REDACTED]","echo":"Bearer [REDACTED]"}`
|
||||||
|
if pe.Body() != want {
|
||||||
|
t.Fatalf("Body() = %q, want %q", pe.Body(), want)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderErrorRedactsBeforeTruncating(t *testing.T) {
|
||||||
|
const key = "sk-secret-straddling-the-cut"
|
||||||
|
// The key starts just before the excerpt limit, so truncating first
|
||||||
|
// would leave its prefix in the excerpt.
|
||||||
|
body := strings.Repeat("x", MaxBodyExcerpt-5) + key
|
||||||
|
pe := NewProviderError(500, []byte(body), []string{key})
|
||||||
|
if strings.Contains(pe.Body(), "sk-se") {
|
||||||
|
t.Fatalf("secret prefix leaked: %q", pe.Body())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderErrorTruncatesOnRuneBoundary(t *testing.T) {
|
||||||
|
body := strings.Repeat("a", MaxBodyExcerpt-1) + "é" + strings.Repeat("b", 100)
|
||||||
|
pe := NewProviderError(500, []byte(body), []string{"sk-test"})
|
||||||
|
want := strings.Repeat("a", MaxBodyExcerpt-1) + "…"
|
||||||
|
if pe.Body() != want {
|
||||||
|
t.Fatalf("Body() has %d bytes, want %d", len(pe.Body()), len(want))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestProviderErrorOmitsBodyWithoutSecret(t *testing.T) {
|
||||||
|
const key = "sk-live-0123456789abcdef"
|
||||||
|
body := []byte(`{"error":"invalid key sk-live-0123456789abcdef"}`)
|
||||||
|
for _, secrets := range [][]string{nil, {}, {""}} {
|
||||||
|
pe := NewProviderError(401, body, secrets)
|
||||||
|
if pe.Body() != "" {
|
||||||
|
t.Fatalf("secrets %q: Body() = %q, want empty", secrets, pe.Body())
|
||||||
|
}
|
||||||
|
if got := pe.Error(); got != "llm: provider error: HTTP 401" {
|
||||||
|
t.Fatalf("secrets %q: Error() = %q", secrets, got)
|
||||||
|
}
|
||||||
|
if strings.Contains(pe.Error(), key) {
|
||||||
|
t.Fatalf("secrets %q: secret leaked: %q", secrets, pe.Error())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,127 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
|
"slices"
|
||||||
|
"sync"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Fake is a scripted Provider for tests. Each method returns its queued
|
||||||
|
// results in order, one per call, and records every Request Generate
|
||||||
|
// receives. A call with nothing left in its queue fails, so an unexpected
|
||||||
|
// call is never silently answered. It is safe for concurrent use.
|
||||||
|
type Fake struct {
|
||||||
|
mu sync.Mutex
|
||||||
|
generate []GenerateResult
|
||||||
|
models []ListModelsResult
|
||||||
|
verify []error
|
||||||
|
requests []Request
|
||||||
|
}
|
||||||
|
|
||||||
|
var _ Provider = (*Fake)(nil)
|
||||||
|
|
||||||
|
// GenerateResult is one scripted Generate outcome.
|
||||||
|
type GenerateResult struct {
|
||||||
|
Response Response
|
||||||
|
Err error
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListModelsResult is one scripted ListModels outcome.
|
||||||
|
type ListModelsResult struct {
|
||||||
|
Models []Model
|
||||||
|
Err error
|
||||||
|
}
|
||||||
|
|
||||||
|
// QueueGenerate appends results to the Generate script.
|
||||||
|
func (f *Fake) QueueGenerate(results ...GenerateResult) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.generate = append(f.generate, results...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// QueueListModels appends results to the ListModels script.
|
||||||
|
func (f *Fake) QueueListModels(results ...ListModelsResult) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.models = append(f.models, results...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// QueueVerify appends results to the Verify script; a nil entry is success.
|
||||||
|
func (f *Fake) QueueVerify(results ...error) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.verify = append(f.verify, results...)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Generate records req and returns the next scripted result.
|
||||||
|
func (f *Fake) Generate(_ context.Context, req Request) (Response, error) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
f.requests = append(f.requests, cloneRequest(req))
|
||||||
|
if len(f.generate) == 0 {
|
||||||
|
return Response{}, fmt.Errorf("llm: fake: unscripted Generate call %d", len(f.requests))
|
||||||
|
}
|
||||||
|
r := f.generate[0]
|
||||||
|
f.generate = f.generate[1:]
|
||||||
|
return r.Response, r.Err
|
||||||
|
}
|
||||||
|
|
||||||
|
// ListModels returns the next scripted result.
|
||||||
|
func (f *Fake) ListModels(context.Context) ([]Model, error) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
if len(f.models) == 0 {
|
||||||
|
return nil, fmt.Errorf("llm: fake: unscripted ListModels call")
|
||||||
|
}
|
||||||
|
r := f.models[0]
|
||||||
|
f.models = f.models[1:]
|
||||||
|
return slices.Clone(r.Models), r.Err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Verify returns the next scripted result.
|
||||||
|
func (f *Fake) Verify(context.Context) error {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
if len(f.verify) == 0 {
|
||||||
|
return fmt.Errorf("llm: fake: unscripted Verify call")
|
||||||
|
}
|
||||||
|
err := f.verify[0]
|
||||||
|
f.verify = f.verify[1:]
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
// Requests returns a copy of every Request Generate received, in order.
|
||||||
|
func (f *Fake) Requests() []Request {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
out := make([]Request, len(f.requests))
|
||||||
|
for i, r := range f.requests {
|
||||||
|
out[i] = cloneRequest(r)
|
||||||
|
}
|
||||||
|
return out
|
||||||
|
}
|
||||||
|
|
||||||
|
// Remaining reports how many scripted results of each kind are unused, so a
|
||||||
|
// test can assert its script was consumed.
|
||||||
|
func (f *Fake) Remaining() (generate, listModels, verify int) {
|
||||||
|
f.mu.Lock()
|
||||||
|
defer f.mu.Unlock()
|
||||||
|
return len(f.generate), len(f.models), len(f.verify)
|
||||||
|
}
|
||||||
|
|
||||||
|
// cloneRequest deep-copies r so a recorded Request is unaffected by later
|
||||||
|
// changes the caller makes to its slices and pointers.
|
||||||
|
func cloneRequest(r Request) Request {
|
||||||
|
r.Messages = slices.Clone(r.Messages)
|
||||||
|
if r.Temperature != nil {
|
||||||
|
t := *r.Temperature
|
||||||
|
r.Temperature = &t
|
||||||
|
}
|
||||||
|
if r.Schema != nil {
|
||||||
|
s := JSONSchema{Name: r.Schema.Name, Schema: json.RawMessage(slices.Clone([]byte(r.Schema.Schema)))}
|
||||||
|
r.Schema = &s
|
||||||
|
}
|
||||||
|
return r
|
||||||
|
}
|
||||||
@@ -0,0 +1,138 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestFakeGenerateReturnsScriptInOrder(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
var f Fake
|
||||||
|
f.QueueGenerate(
|
||||||
|
GenerateResult{Response: Response{Text: "first", Model: "m1"}},
|
||||||
|
GenerateResult{Err: &RateLimitedError{RetryAfter: time.Second}},
|
||||||
|
)
|
||||||
|
f.QueueGenerate(GenerateResult{Response: Response{JSON: json.RawMessage(`{"ok":true}`), FinishReason: "stop"}})
|
||||||
|
|
||||||
|
req := Request{Messages: []Message{{Role: RoleUser, Content: "one"}}}
|
||||||
|
resp, err := f.Generate(ctx, req)
|
||||||
|
if err != nil || resp.Text != "first" || resp.Model != "m1" {
|
||||||
|
t.Fatalf("call 1 = %+v, %v", resp, err)
|
||||||
|
}
|
||||||
|
if _, err := f.Generate(ctx, req); !errors.Is(err, ErrRateLimited) {
|
||||||
|
t.Fatalf("call 2 err = %v, want ErrRateLimited", err)
|
||||||
|
}
|
||||||
|
resp, err = f.Generate(ctx, req)
|
||||||
|
if err != nil || string(resp.JSON) != `{"ok":true}` || resp.FinishReason != "stop" {
|
||||||
|
t.Fatalf("call 3 = %+v, %v", resp, err)
|
||||||
|
}
|
||||||
|
if _, err := f.Generate(ctx, req); err == nil {
|
||||||
|
t.Fatal("unscripted call 4 succeeded")
|
||||||
|
}
|
||||||
|
if g, l, v := f.Remaining(); g != 0 || l != 0 || v != 0 {
|
||||||
|
t.Fatalf("Remaining() = %d, %d, %d", g, l, v)
|
||||||
|
}
|
||||||
|
if n := len(f.Requests()); n != 4 {
|
||||||
|
t.Fatalf("recorded %d requests, want 4 (unscripted calls are recorded too)", n)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFakeRecordsRequests(t *testing.T) {
|
||||||
|
var f Fake
|
||||||
|
f.QueueGenerate(GenerateResult{}, GenerateResult{})
|
||||||
|
temp := 0.2
|
||||||
|
first := Request{
|
||||||
|
System: "be brief",
|
||||||
|
Messages: []Message{{Role: RoleUser, Content: "plan dinners"}},
|
||||||
|
Temperature: &temp,
|
||||||
|
MaxTokens: 100,
|
||||||
|
Schema: &JSONSchema{Name: "plan", Schema: json.RawMessage(`{"type":"object"}`)},
|
||||||
|
Timeout: 5 * time.Second,
|
||||||
|
}
|
||||||
|
if _, err := f.Generate(context.Background(), first); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if _, err := f.Generate(context.Background(), Request{System: "second"}); err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Later changes by the caller must not alter what was recorded.
|
||||||
|
first.Messages[0].Content = "changed"
|
||||||
|
temp = 1
|
||||||
|
first.Schema.Schema[2] = 'X'
|
||||||
|
|
||||||
|
got := f.Requests()
|
||||||
|
if len(got) != 2 || got[1].System != "second" {
|
||||||
|
t.Fatalf("Requests() = %+v", got)
|
||||||
|
}
|
||||||
|
r := got[0]
|
||||||
|
if r.System != "be brief" || r.Messages[0].Content != "plan dinners" || *r.Temperature != 0.2 ||
|
||||||
|
r.MaxTokens != 100 || r.Schema.Name != "plan" || string(r.Schema.Schema) != `{"type":"object"}` ||
|
||||||
|
r.Timeout != 5*time.Second {
|
||||||
|
t.Fatalf("recorded request = %+v", r)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Changing a returned copy must not alter the record either.
|
||||||
|
got[0].Messages[0].Content = "mutated"
|
||||||
|
if f.Requests()[0].Messages[0].Content != "plan dinners" {
|
||||||
|
t.Fatal("Requests() returned shared state")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFakeListModelsAndVerify(t *testing.T) {
|
||||||
|
ctx := context.Background()
|
||||||
|
var f Fake
|
||||||
|
f.QueueListModels(
|
||||||
|
ListModelsResult{Models: []Model{{ID: "a", DisplayName: "A"}, {ID: "b", DisplayName: "b"}}},
|
||||||
|
ListModelsResult{Err: ErrAuth},
|
||||||
|
)
|
||||||
|
f.QueueVerify(nil, fmt.Errorf("wrapped: %w", ErrAuth))
|
||||||
|
|
||||||
|
models, err := f.ListModels(ctx)
|
||||||
|
if err != nil || len(models) != 2 || models[0].ID != "a" {
|
||||||
|
t.Fatalf("ListModels 1 = %+v, %v", models, err)
|
||||||
|
}
|
||||||
|
if _, err := f.ListModels(ctx); !errors.Is(err, ErrAuth) {
|
||||||
|
t.Fatalf("ListModels 2 err = %v, want ErrAuth", err)
|
||||||
|
}
|
||||||
|
if _, err := f.ListModels(ctx); err == nil {
|
||||||
|
t.Fatal("unscripted ListModels succeeded")
|
||||||
|
}
|
||||||
|
if err := f.Verify(ctx); err != nil {
|
||||||
|
t.Fatalf("Verify 1 = %v", err)
|
||||||
|
}
|
||||||
|
if err := f.Verify(ctx); !errors.Is(err, ErrAuth) {
|
||||||
|
t.Fatalf("Verify 2 = %v, want ErrAuth", err)
|
||||||
|
}
|
||||||
|
if err := f.Verify(ctx); err == nil {
|
||||||
|
t.Fatal("unscripted Verify succeeded")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestFakeConcurrentUse(t *testing.T) {
|
||||||
|
const n = 50
|
||||||
|
var f Fake
|
||||||
|
for i := range n {
|
||||||
|
f.QueueGenerate(GenerateResult{Response: Response{Text: fmt.Sprint(i)}})
|
||||||
|
}
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
for range n {
|
||||||
|
wg.Go(func() {
|
||||||
|
if _, err := f.Generate(context.Background(), Request{}); err != nil {
|
||||||
|
t.Error(err)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
wg.Wait()
|
||||||
|
if len(f.Requests()) != n {
|
||||||
|
t.Fatalf("recorded %d requests, want %d", len(f.Requests()), n)
|
||||||
|
}
|
||||||
|
if g, _, _ := f.Remaining(); g != 0 {
|
||||||
|
t.Fatalf("%d Generate results unused", g)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,132 @@
|
|||||||
|
// Package llm is the single interface MealQuest uses to talk to a language
|
||||||
|
// model. Every vendor adapter implements Provider over plain net/http (no
|
||||||
|
// vendor SDKs) and reports failures with this package's typed errors, so
|
||||||
|
// callers decide what to do from errors.Is and errors.As alone.
|
||||||
|
//
|
||||||
|
// Errors from this package and its adapters never contain an API key.
|
||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Provider is a language-model vendor.
|
||||||
|
type Provider interface {
|
||||||
|
// Generate sends one request and returns the model's reply.
|
||||||
|
Generate(ctx context.Context, req Request) (Response, error)
|
||||||
|
// ListModels returns the models the configured credentials can use.
|
||||||
|
ListModels(ctx context.Context) ([]Model, error)
|
||||||
|
// Verify checks that the configured credentials are accepted.
|
||||||
|
Verify(ctx context.Context) error
|
||||||
|
}
|
||||||
|
|
||||||
|
// Role is the author of a Message.
|
||||||
|
type Role string
|
||||||
|
|
||||||
|
// The roles a Message may have. The system prompt is Request.System, not a
|
||||||
|
// message.
|
||||||
|
const (
|
||||||
|
RoleUser Role = "user"
|
||||||
|
RoleAssistant Role = "assistant"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Message is one turn of a conversation.
|
||||||
|
type Message struct {
|
||||||
|
Role Role
|
||||||
|
Content string
|
||||||
|
}
|
||||||
|
|
||||||
|
// JSONSchema is a named structured-output contract: when a Request carries
|
||||||
|
// one, the model must reply with JSON matching Schema.
|
||||||
|
type JSONSchema struct {
|
||||||
|
// Name identifies the contract to the vendor, e.g. "side_quests".
|
||||||
|
Name string
|
||||||
|
// Schema is a JSON Schema object.
|
||||||
|
Schema json.RawMessage
|
||||||
|
}
|
||||||
|
|
||||||
|
// Request is one generation request.
|
||||||
|
type Request struct {
|
||||||
|
// System is the system prompt; empty means none.
|
||||||
|
System string
|
||||||
|
// Messages is the conversation, oldest first. It must not be empty.
|
||||||
|
Messages []Message
|
||||||
|
// Temperature is the sampling temperature; nil means the vendor default.
|
||||||
|
Temperature *float64
|
||||||
|
// MaxTokens caps the reply length; 0 means the vendor default.
|
||||||
|
MaxTokens int
|
||||||
|
// Schema, when set, requests structured output; the reply is then in
|
||||||
|
// Response.JSON.
|
||||||
|
Schema *JSONSchema
|
||||||
|
// Timeout bounds the request; 0 means only ctx bounds it.
|
||||||
|
Timeout time.Duration
|
||||||
|
}
|
||||||
|
|
||||||
|
// Validate reports whether r is well formed. Adapters call it before
|
||||||
|
// contacting the vendor.
|
||||||
|
func (r Request) Validate() error {
|
||||||
|
if len(r.Messages) == 0 {
|
||||||
|
return errors.New("llm: request has no messages")
|
||||||
|
}
|
||||||
|
for i, m := range r.Messages {
|
||||||
|
if m.Role != RoleUser && m.Role != RoleAssistant {
|
||||||
|
return fmt.Errorf("llm: message %d has unknown role %q", i, m.Role)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if r.Temperature != nil && *r.Temperature < 0 {
|
||||||
|
return fmt.Errorf("llm: temperature %v is negative", *r.Temperature)
|
||||||
|
}
|
||||||
|
if r.MaxTokens < 0 {
|
||||||
|
return fmt.Errorf("llm: max tokens %d is negative", r.MaxTokens)
|
||||||
|
}
|
||||||
|
if r.Timeout < 0 {
|
||||||
|
return fmt.Errorf("llm: timeout %v is negative", r.Timeout)
|
||||||
|
}
|
||||||
|
if r.Schema != nil {
|
||||||
|
if r.Schema.Name == "" {
|
||||||
|
return errors.New("llm: schema has no name")
|
||||||
|
}
|
||||||
|
var obj map[string]json.RawMessage
|
||||||
|
if err := json.Unmarshal(r.Schema.Schema, &obj); err != nil || obj == nil {
|
||||||
|
return fmt.Errorf("llm: schema %q is not a JSON object", r.Schema.Name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// Usage is the token count of one Generate call.
|
||||||
|
type Usage struct {
|
||||||
|
InputTokens int
|
||||||
|
OutputTokens int
|
||||||
|
}
|
||||||
|
|
||||||
|
// Total returns InputTokens + OutputTokens.
|
||||||
|
func (u Usage) Total() int { return u.InputTokens + u.OutputTokens }
|
||||||
|
|
||||||
|
// Response is the model's reply to a Request.
|
||||||
|
type Response struct {
|
||||||
|
// Text is the reply text.
|
||||||
|
Text string
|
||||||
|
// JSON is the structured reply when the Request carried a Schema, and
|
||||||
|
// nil otherwise.
|
||||||
|
JSON json.RawMessage
|
||||||
|
// Usage is the tokens the call consumed.
|
||||||
|
Usage Usage
|
||||||
|
// Model is the model that answered, as the vendor reports it.
|
||||||
|
Model string
|
||||||
|
// FinishReason is why generation stopped, as the vendor reports it.
|
||||||
|
FinishReason string
|
||||||
|
}
|
||||||
|
|
||||||
|
// Model is a model a Provider offers.
|
||||||
|
type Model struct {
|
||||||
|
// ID is the vendor's model identifier, used in requests.
|
||||||
|
ID string
|
||||||
|
// DisplayName is a human-readable name; it equals ID when the vendor
|
||||||
|
// gives none.
|
||||||
|
DisplayName string
|
||||||
|
}
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
package llm
|
||||||
|
|
||||||
|
import (
|
||||||
|
"encoding/json"
|
||||||
|
"testing"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRequestValidate(t *testing.T) {
|
||||||
|
neg := -0.1
|
||||||
|
valid := func() Request {
|
||||||
|
return Request{
|
||||||
|
Messages: []Message{{Role: RoleUser, Content: "hi"}, {Role: RoleAssistant, Content: "hello"}},
|
||||||
|
Schema: &JSONSchema{Name: "plan", Schema: json.RawMessage(`{"type":"object"}`)},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if err := valid().Validate(); err != nil {
|
||||||
|
t.Fatalf("valid request: %v", err)
|
||||||
|
}
|
||||||
|
cases := map[string]func(*Request){
|
||||||
|
"no messages": func(r *Request) { r.Messages = nil },
|
||||||
|
"unknown role": func(r *Request) { r.Messages[0].Role = "system" },
|
||||||
|
"negative temp": func(r *Request) { r.Temperature = &neg },
|
||||||
|
"negative max": func(r *Request) { r.MaxTokens = -1 },
|
||||||
|
"negative time": func(r *Request) { r.Timeout = -1 },
|
||||||
|
"unnamed schema": func(r *Request) { r.Schema.Name = "" },
|
||||||
|
"schema array": func(r *Request) { r.Schema.Schema = json.RawMessage(`[]`) },
|
||||||
|
"schema null": func(r *Request) { r.Schema.Schema = json.RawMessage(`null`) },
|
||||||
|
"schema invalid": func(r *Request) { r.Schema.Schema = json.RawMessage(`{`) },
|
||||||
|
"schema empty": func(r *Request) { r.Schema.Schema = nil },
|
||||||
|
}
|
||||||
|
for name, mutate := range cases {
|
||||||
|
r := valid()
|
||||||
|
mutate(&r)
|
||||||
|
if err := r.Validate(); err == nil {
|
||||||
|
t.Errorf("%s: Validate() = nil", name)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestUsageTotal(t *testing.T) {
|
||||||
|
if got := (Usage{InputTokens: 12, OutputTokens: 30}).Total(); got != 42 {
|
||||||
|
t.Fatalf("Total() = %d, want 42", got)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user