llmgw/internal/router/messages_test.go

package router

import (
	"bytes"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/config"
)

func TestAMessagesRequestBecomesAChatCompletion(t *testing.T) {
	var in msgRequest
	_ = json.Unmarshal([]byte(`{"model":"m","max_tokens":100,"system":[{"type":"text","text":"be brief"}],
		"stop_sequences":["END"],"tool_choice":{"type":"any"},
		"tools":[{"name":"get_weather","description":"w","input_schema":{"type":"object"}}],
		"messages":[
		  {"role":"user","content":[{"type":"text","text":"look"},{"type":"image","source":{"type":"base64","media_type":"image/png","data":"AAAA"}}]},
		  {"role":"assistant","content":[{"type":"text","text":"checking"},{"type":"tool_use","id":"toolu_1","name":"get_weather","input":{"city":"Ankara"}}]},
		  {"role":"user","content":[{"type":"tool_result","tool_use_id":"toolu_1","content":"sunny"},{"type":"text","text":"and tomorrow?"}]}]}`), &in)
	chat, err := messagesToChat(in)
	if err != nil {
		t.Fatal(err)
	}
	b, _ := json.Marshal(chat)
	for _, want := range []string{
		`{"content":"be brief","role":"system"}`,
		`"image_url":{"url":"data:image/png;base64,AAAA"}`,
		`"tool_calls":[{"function":{"arguments":"{\"city\":\"Ankara\"}","name":"get_weather"},"id":"toolu_1","type":"function"}]`,
		`{"content":"sunny","role":"tool","tool_call_id":"toolu_1"}`,
		`{"content":"and tomorrow?","role":"user"}`,
		`"max_tokens":100`, `"stop":["END"]`, `"tool_choice":"required"`,
	} {
		if !strings.Contains(string(b), want) {
			t.Errorf("lacks %s:\n%s", want, b)
		}
	}
}

func messagesRouter(t *testing.T, answer string) (*Router, http.Handler) {
	t.Helper()
	srv, _, _ := engineFor(t, "m", map[string]string{"/v1/chat/completions": answer})
	cfg := baseConfig(config.Endpoint{ID: "a", URL: srv.URL, Models: []string{"m"}})
	rt := newTestRouter(t, cfg)
	return rt, rt.Handler()
}

func sendMessages(t *testing.T, h http.Handler, body map[string]any) *httptest.ResponseRecorder {
	t.Helper()
	b, _ := json.Marshal(body)
	req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(b))
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	return rec
}

func TestAMessageComesBackInAnthropicsShape(t *testing.T) {
	_, h := messagesRouter(t, `{"choices":[{"message":{"content":"Hello","tool_calls":[{"id":"call_9","function":{"name":"f","arguments":"{\"a\":1}"}}]},"finish_reason":"tool_calls"}],"usage":{"prompt_tokens":5,"completion_tokens":3}}`)
	rec := sendMessages(t, h, map[string]any{"model": "m", "max_tokens": 10,
		"messages": []any{map[string]any{"role": "user", "content": "hi"}}})
	var out struct {
		Type       string           `json:"type"`
		Role       string           `json:"role"`
		Content    []map[string]any `json:"content"`
		StopReason string           `json:"stop_reason"`
		Usage      map[string]int   `json:"usage"`
	}
	_ = json.Unmarshal(rec.Body.Bytes(), &out)
	if rec.Code != 200 || out.Type != "message" || out.Role != "assistant" || len(out.Content) != 2 ||
		out.Content[0]["text"] != "Hello" || out.Content[1]["type"] != "tool_use" || out.StopReason != "tool_use" ||
		out.Usage["input_tokens"] != 5 {
		t.Errorf("%d %s", rec.Code, rec.Body.String())
	}

	rec = sendMessages(t, h, map[string]any{"model": "nope", "max_tokens": 10,
		"messages": []any{map[string]any{"role": "user", "content": "hi"}}})
	var e struct {
		Type  string `json:"type"`
		Error struct {
			Type, Message string
		} `json:"error"`
	}
	_ = json.Unmarshal(rec.Body.Bytes(), &e)
	if rec.Code < 400 || e.Type != "error" || e.Error.Type == "" || e.Error.Message == "" {
		t.Errorf("error shape: %d %s", rec.Code, rec.Body.String())
	}
}

func TestAStreamedMessageIsAnthropicsEvents(t *testing.T) {
	stream := "data: {\"choices\":[{\"delta\":{\"content\":\"Hel\"}}]}\n\n" +
		"data: {\"choices\":[{\"delta\":{\"content\":\"lo\"}}]}\n\n" +
		"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"id\":\"call_1\",\"function\":{\"name\":\"f\",\"arguments\":\"{\\\"a\\\"\"}}]}}]}\n\n" +
		"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\":1}\"}}]},\"finish_reason\":\"tool_calls\"}]}\n\n" +
		"data: {\"choices\":[],\"usage\":{\"prompt_tokens\":4,\"completion_tokens\":6}}\n\n" +
		"data: [DONE]\n\n"
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/chat/completions" {
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = w.Write([]byte(stream))
	}))
	t.Cleanup(srv.Close)
	rt := newTestRouter(t, baseConfig(config.Endpoint{ID: "a", URL: srv.URL, Models: []string{"m"}}))
	rec := sendMessages(t, rt.Handler(), map[string]any{"model": "m", "max_tokens": 10, "stream": true,
		"messages": []any{map[string]any{"role": "user", "content": "hi"}}})
	got := rec.Body.String()
	var events []string
	for _, l := range strings.Split(got, "\n") {
		if e, ok := strings.CutPrefix(l, "event: "); ok {
			events = append(events, e)
		}
	}
	want := "message_start content_block_start content_block_delta content_block_delta content_block_stop " +
		"content_block_start content_block_delta content_block_delta content_block_stop message_delta message_stop"
	if strings.Join(events, " ") != want {
		t.Errorf("events:\n%s\nwant:\n%s\n%s", strings.Join(events, " "), want, got)
	}
	for _, s := range []string{`"text":"Hel","type":"text_delta"`, `"partial_json":":1}"`, `"stop_reason":"tool_use"`, `"output_tokens":6`} {
		if !strings.Contains(got, s) {
			t.Errorf("lacks %s", s)
		}
	}
}

func TestAnAnthropicClientsKeyIsAccepted(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	b, _ := json.Marshal(map[string]any{"model": "q14", "max_tokens": 5, "messages": []any{map[string]any{"role": "user", "content": "hi"}}})
	req := httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(b))
	req.Header.Set("x-api-key", secret)
	req.Header.Set("anthropic-version", "2023-06-01")
	rec := httptest.NewRecorder()
	rt.Handler().ServeHTTP(rec, req)
	if rec.Code != 200 || !strings.Contains(rec.Body.String(), `"type":"message"`) {
		t.Errorf("%d %s", rec.Code, rec.Body.String())
	}
}