package model import ( "context" "encoding/json" "errors" "io" "net/http" "net/http/httptest" "reflect" "strings" "testing" "git.yeij.top/AskaEth/Cyrene_v2/internal/config" ) const okCompletion = `{ "id": "cmpl-1", "model": "gpt-test", "choices": [ {"index": 0, "message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop"} ], "usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7} }` // capturedRequest records what an OpenAIChatAdapter sent to a stub server. type capturedRequest struct { method string path string contentType string body map[string]any } func newCompletionServer(t *testing.T, status int, contentType, body string, ch chan<- capturedRequest) *httptest.Server { t.Helper() srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var reqBody map[string]any if ch != nil { _ = json.NewDecoder(r.Body).Decode(&reqBody) } if contentType != "" { w.Header().Set("Content-Type", contentType) } w.WriteHeader(status) if body != "" { _, _ = io.WriteString(w, body) } if ch != nil { ch <- capturedRequest{ method: r.Method, path: r.URL.Path, contentType: r.Header.Get("Content-Type"), body: reqBody, } } })) t.Cleanup(srv.Close) return srv } func TestOpenAIChat_Chat_NonStreaming(t *testing.T) { t.Parallel() ch := make(chan capturedRequest, 1) srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch) recs := make([]CallRecord, 0, 1) temp := 1.0 ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, APIKey: "sk-test", Model: "gpt-test"}, WithRecorder(func(cr CallRecord) { recs = append(recs, cr) })) ctx := context.Background() resp, err := ad.Chat(ctx, Request{ Messages: []Message{TextMessage(RoleUser, "hi")}, Model: "gpt-test", Temperature: &temp, MaxTokens: 64, Purpose: "chat", }, Options{}) if err != nil { t.Fatalf("Chat error: %v", err) } req := <-ch if req.method != http.MethodPost { t.Errorf("method = %s, want POST", req.method) } if req.path != "/chat/completions" { t.Errorf("path = %s, want /chat/completions", req.path) } if req.contentType != "application/json" { t.Errorf("content-type = %s, want application/json", req.contentType) } if req.body["model"] != "gpt-test" { t.Errorf("body.model = %v, want gpt-test", req.body["model"]) } if req.body["stream"] != nil && req.body["stream"] != false { t.Errorf("body.stream = %v, want omitted(false)", req.body["stream"]) } if req.body["temperature"] != 1.0 { t.Errorf("body.temperature = %v, want 1.0", req.body["temperature"]) } if req.body["max_tokens"] != float64(64) { t.Errorf("body.max_tokens = %v, want 64", req.body["max_tokens"]) } msgs, _ := req.body["messages"].([]any) if len(msgs) != 1 { t.Fatalf("body.messages = %v, want 1 message", req.body["messages"]) } if resp.Model != "gpt-test" { t.Errorf("resp.Model = %q, want gpt-test", resp.Model) } if resp.Content != "Hello!" { t.Errorf("resp.Content = %q, want Hello!", resp.Content) } if resp.TotalTokens != 7 || resp.PromptTokens != 5 || resp.CompletionTokens != 2 { t.Errorf("resp usage = (%d,%d,%d), want (5,2,7)", resp.PromptTokens, resp.CompletionTokens, resp.TotalTokens) } if resp.Message.Role != RoleAssistant || len(resp.Message.Content) != 1 || resp.Message.Content[0].Text != "Hello!" { t.Errorf("resp.Message = %+v", resp.Message) } if len(recs) != 1 { t.Fatalf("records = %d, want 1", len(recs)) } if recs[0].Model != "gpt-test" || recs[0].Purpose != "chat" || recs[0].Tokens != 7 || recs[0].TookMs < 0 { t.Errorf("record = %+v, want model=gpt-test purpose=chat tokens=7 took_ms>=0", recs[0]) } } func TestOpenAIChat_Chat_ImageContent(t *testing.T) { t.Parallel() ch := make(chan capturedRequest, 1) srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch) ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"}) req := Request{Messages: []Message{ImageMessage(RoleUser, "what is this", "data:image/png;base64,abc")}} if _, err := ad.Chat(context.Background(), req, Options{}); err != nil { t.Fatalf("Chat error: %v", err) } got := <-ch msgs, _ := got.body["messages"].([]any) part0, _ := msgs[0].(map[string]any) content, _ := part0["content"].([]any) if len(content) != 2 { t.Fatalf("content parts = %d, want 2 (text+image)", len(content)) } found := false for _, c := range content { cp, _ := c.(map[string]any) if cp["type"] != "image_url" { continue } found = true iu, _ := cp["image_url"].(map[string]any) if iu["url"] != "data:image/png;base64,abc" { t.Fatalf("image url = %v, want data:image/png;base64,abc", iu["url"]) } } if !found { t.Fatalf("no image_url part found in %v", content) } } func TestOpenAIChat_Chat_ToolsPlaceholder(t *testing.T) { t.Parallel() ch := make(chan capturedRequest, 1) srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch) ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"}) req := Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Tools: []ToolSpec{{Name: "lookup", Description: "look it up", Parameters: map[string]any{"type": "object"}}}} if _, err := ad.Chat(context.Background(), req, Options{}); err != nil { t.Fatalf("Chat error: %v", err) } got := <-ch tools, ok := got.body["tools"].([]any) if !ok || len(tools) != 1 { t.Fatalf("body.tools = %v, want 1 entry", got.body["tools"]) } } func TestOpenAIChat_ChatStream_SSE(t *testing.T) { t.Parallel() sse := strings.Join([]string{ `data: {"model":"gpt-test","choices":[{"index":0,"delta":{"role":"assistant","content":"Hel"}}]}`, ``, `data: {"choices":[{"delta":{"content":"lo"}}]}`, ``, `data: {"model":"gpt-test","choices":[],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`, ``, `data: [DONE]`, ``, }, "\n") srv := newCompletionServer(t, http.StatusOK, "text/event-stream", sse, nil) recs := make([]CallRecord, 0, 1) ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"}, WithRecorder(func(cr CallRecord) { recs = append(recs, cr) })) var deltas []string resp, err := ad.ChatStream(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Model: "gpt-test", Purpose: "chat"}, func(d string) { deltas = append(deltas, d) }) if err != nil { t.Fatalf("ChatStream error: %v", err) } if resp.Content != "Hello" { t.Errorf("resp.Content = %q, want Hello", resp.Content) } if resp.Model != "gpt-test" { t.Errorf("resp.Model = %q, want gpt-test", resp.Model) } if resp.TotalTokens != 7 { t.Errorf("resp.TotalTokens = %d, want 7", resp.TotalTokens) } if !reflect.DeepEqual(deltas, []string{"Hel", "lo"}) { t.Errorf("deltas = %v, want [Hel lo]", deltas) } if len(recs) != 1 || recs[0].Tokens != 7 || recs[0].Purpose != "chat" { t.Errorf("records = %+v, want 1 record tokens=7 purpose=chat", recs) } } func TestOpenAIChat_Chat_UpstreamError(t *testing.T) { t.Parallel() srv := newCompletionServer(t, http.StatusInternalServerError, "application/json", `{"error":"boom"}`, nil) ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "m"}) _, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}) if err == nil { t.Fatal("want error, got nil") } if !errors.Is(err, ErrUpstream) { t.Fatalf("err = %v, want ErrUpstream", err) } } func TestOpenAIChat_Chat_NoBaseURL(t *testing.T) { t.Parallel() ad := NewOpenAIChat(config.LLMConfig{}) if _, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrUpstream) { t.Fatalf("err = %v, want ErrUpstream", err) } } func TestOpenAIChat_Chat_NoModel(t *testing.T) { t.Parallel() srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, nil) ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL}) if _, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrNoModel) { t.Fatalf("err = %v, want ErrNoModel", err) } }