diff --git a/internal/entity/models/n1n.go b/internal/entity/models/n1n.go index a5d58d526f..2c0d4078a4 100644 --- a/internal/entity/models/n1n.go +++ b/internal/entity/models/n1n.go @@ -54,7 +54,15 @@ type N1NModel struct { // long-lived SSE streams in ChatStreamlyWithSender. Non-streaming // callers wrap each request with context.WithTimeout instead. func NewN1NModel(baseURL map[string]string, urlSuffix URLSuffix) *N1NModel { - transport := http.DefaultTransport.(*http.Transport).Clone() + defaultTransport, ok := http.DefaultTransport.(*http.Transport) + var transport *http.Transport + if ok { + transport = defaultTransport.Clone() + } else { + transport = &http.Transport{ + Proxy: http.ProxyFromEnvironment, + } + } transport.MaxIdleConns = 100 transport.MaxIdleConnsPerHost = 10 transport.IdleConnTimeout = 90 * time.Second diff --git a/internal/entity/models/n1n_test.go b/internal/entity/models/n1n_test.go new file mode 100644 index 0000000000..6fdb90ddaf --- /dev/null +++ b/internal/entity/models/n1n_test.go @@ -0,0 +1,49 @@ +package models + +import ( + "net/http" + "testing" +) + +type roundTripperFunc func(*http.Request) (*http.Response, error) + +func (f roundTripperFunc) RoundTrip(r *http.Request) (*http.Response, error) { + return f(r) +} + +func newN1NForTest(baseURL string) *N1NModel { + return NewN1NModel( + map[string]string{"default": baseURL}, + URLSuffix{Chat: "chat/completions", Models: "models"}, + ) +} + +func TestN1NName(t *testing.T) { + if got := newN1NForTest("http://unused").Name(); got != "n1n" { + t.Errorf("Name()=%q, want %q", got, "n1n") + } +} + +func TestN1NFactory(t *testing.T) { + driver, err := NewModelFactory().CreateModelDriver("n1n", map[string]string{"default": "http://unused"}, URLSuffix{}) + if err != nil { + t.Fatalf("CreateModelDriver: %v", err) + } + if _, ok := driver.(*N1NModel); !ok { + t.Fatalf("driver type=%T, want *N1NModel", driver) + } +} + +func TestN1NNewModelWithCustomDefaultTransport(t *testing.T) { + original := http.DefaultTransport + http.DefaultTransport = roundTripperFunc(func(*http.Request) (*http.Response, error) { + return nil, nil + }) + t.Cleanup(func() { + http.DefaultTransport = original + }) + + if model := NewN1NModel(map[string]string{"default": "http://unused"}, URLSuffix{}); model == nil { + t.Fatal("NewN1NModel returned nil") + } +}