114 lines
3.6 KiB
Go
114 lines
3.6 KiB
Go
package model
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
|
)
|
|
|
|
func TestChain_Chat_Fallback(t *testing.T) {
|
|
t.Parallel()
|
|
var primaryCalls, backupCalls atomic.Int64
|
|
primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
primaryCalls.Add(1)
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
_, _ = io.WriteString(w, `{"error":"down"}`)
|
|
}))
|
|
t.Cleanup(primary.Close)
|
|
backup := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
backupCalls.Add(1)
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
_, _ = io.WriteString(w, okCompletion)
|
|
}))
|
|
t.Cleanup(backup.Close)
|
|
|
|
records := make([]CallRecord, 0, 1)
|
|
chain := NewChain(
|
|
Step{Model: "main", Adapter: NewOpenAIChat(
|
|
config.LLMConfig{BaseURL: primary.URL, Model: "main"},
|
|
WithRecorder(func(cr CallRecord) { records = append(records, cr) }),
|
|
)},
|
|
Step{Model: "backup", Adapter: NewOpenAIChat(
|
|
config.LLMConfig{BaseURL: backup.URL, Model: "backup"},
|
|
WithRecorder(func(cr CallRecord) { records = append(records, cr) }),
|
|
)},
|
|
)
|
|
|
|
resp, err := chain.Chat(context.Background(),
|
|
Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Purpose: "chat"}, Options{})
|
|
if err != nil {
|
|
t.Fatalf("Chain.Chat error: %v", err)
|
|
}
|
|
if resp.Content != "Hello!" {
|
|
t.Errorf("resp.Content = %q, want Hello!", resp.Content)
|
|
}
|
|
if primaryCalls.Load() != 1 {
|
|
t.Errorf("primary calls = %d, want 1", primaryCalls.Load())
|
|
}
|
|
if backupCalls.Load() != 1 {
|
|
t.Errorf("backup calls = %d, want 1", backupCalls.Load())
|
|
}
|
|
if len(records) != 1 {
|
|
t.Fatalf("records = %d, want just the successful backup record", len(records))
|
|
}
|
|
r := records[0]
|
|
if r.Model != "backup" || r.Tokens != 7 || r.Purpose != "chat" || r.TookMs < 0 {
|
|
t.Errorf("record = %+v, want model=backup tokens=7 purpose=chat took_ms>=0", r)
|
|
}
|
|
}
|
|
|
|
func TestChain_Chat_AllStepsFail(t *testing.T) {
|
|
t.Parallel()
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
_, _ = io.WriteString(w, `{"error":"down"}`)
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
chain := NewChain(Step{Model: "m", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "m"})})
|
|
|
|
if _, err := chain.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrUpstream) {
|
|
t.Fatalf("err = %v, want ErrUpstream", err)
|
|
}
|
|
}
|
|
|
|
func TestChain_Chat_EmptyChain(t *testing.T) {
|
|
t.Parallel()
|
|
chain := NewChain()
|
|
if _, err := chain.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrNoSteps) {
|
|
t.Fatalf("err = %v, want ErrNoSteps", err)
|
|
}
|
|
}
|
|
|
|
func TestNewOpenAIChatChain_Steps(t *testing.T) {
|
|
t.Parallel()
|
|
chain, err := NewOpenAIChatChain(config.LLMConfig{BaseURL: "http://x", Model: "m", Fallback: "fb"})
|
|
if err != nil {
|
|
t.Fatalf("NewOpenAIChatChain error: %v", err)
|
|
}
|
|
if len(chain.steps) != 2 {
|
|
t.Fatalf("steps = %d, want 2 (primary+fallback)", len(chain.steps))
|
|
}
|
|
if chain.steps[0].Model != "m" || chain.steps[1].Model != "fb" {
|
|
t.Errorf("steps = %+v, want models m then fb", chain.steps)
|
|
}
|
|
|
|
noFB, err := NewOpenAIChatChain(config.LLMConfig{BaseURL: "http://x", Model: "m"})
|
|
if err != nil {
|
|
t.Fatalf("NewOpenAIChatChain error: %v", err)
|
|
}
|
|
if len(noFB.steps) != 1 {
|
|
t.Fatalf("steps = %d, want 1 (no fallback)", len(noFB.steps))
|
|
}
|
|
|
|
if _, err := NewOpenAIChatChain(config.LLMConfig{BaseURL: "http://x"}); !errors.Is(err, ErrNoModel) {
|
|
t.Fatalf("err = %v, want ErrNoModel", err)
|
|
}
|
|
}
|