Cyrene_v2/internal/runtime/model/openai_test.go

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)
	}
}