agentu/pkg/llm/openai_test.go

package llm

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestOpenAICompatibleClientStreamsContentAndToolCalls(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/v1/chat/completions" {
			t.Fatalf("path = %s", r.URL.Path)
		}
		if got := r.Header.Get("Authorization"); got != "Bearer sk-test" {
			t.Fatalf("authorization = %q", got)
		}
		body := mustReadBody(t, r)
		var raw map[string]any
		if err := json.NewDecoder(strings.NewReader(body)).Decode(&raw); err != nil {
			t.Fatal(err)
		}
		var req ChatRequest
		if err := json.NewDecoder(strings.NewReader(body)).Decode(&req); err != nil {
			t.Fatal(err)
		}
		if raw["thinking"] != "high" {
			t.Fatalf("thinking = %#v", raw["thinking"])
		}
		streamOptions, ok := raw["stream_options"].(map[string]any)
		if !ok || streamOptions["include_usage"] != true {
			t.Fatalf("stream_options = %#v", raw["stream_options"])
		}
		if !req.Stream {
			t.Fatal("request did not enable stream")
		}
		if req.ToolChoice != "auto" && len(req.Tools) > 0 {
			t.Fatalf("tool_choice = %q", req.ToolChoice)
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"hi "},"finish_reason":""}]}` + "\n\n"))
		_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"file_read","arguments":"{\"path\""}}]},"finish_reason":""}]}` + "\n\n"))
		_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":":\"README.md\"}"}}]},"finish_reason":"tool_calls"}]}` + "\n\n"))
		_, _ = w.Write([]byte("data: [DONE]\n\n"))
	}))
	defer server.Close()

	client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
	var content strings.Builder
	var calls []ToolCallDelta
	err := client.ChatStream(context.Background(), ChatRequest{Model: "test", Extra: map[string]any{"thinking": "high"}}, func(event StreamEvent) error {
		content.WriteString(event.Content)
		calls = append(calls, event.ToolCalls...)
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	if content.String() != "hi " {
		t.Fatalf("content = %q", content.String())
	}
	if len(calls) != 2 {
		t.Fatalf("calls = %#v", calls)
	}
	if calls[0].Name != "file_read" || calls[1].Arguments != ":\"README.md\"}" {
		t.Fatalf("unexpected calls: %#v", calls)
	}
}

func TestChatStreamCapturesUsage(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		body := mustReadBody(t, r)
		var raw map[string]any
		if err := json.NewDecoder(strings.NewReader(body)).Decode(&raw); err != nil {
			t.Fatal(err)
		}
		streamOptions, ok := raw["stream_options"].(map[string]any)
		if !ok || streamOptions["include_usage"] != true {
			t.Fatalf("stream_options = %#v", raw["stream_options"])
		}

		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"ok"},"finish_reason":""}]}` + "\n\n"))
		_, _ = w.Write([]byte(`data: {"choices":[],"usage":{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15}}` + "\n\n"))
		_, _ = w.Write([]byte("data: [DONE]\n\n"))
	}))
	defer server.Close()

	client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
	var usage *Usage
	err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
		if event.Usage != nil {
			usage = event.Usage
		}
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	if usage == nil {
		t.Fatal("usage was not emitted")
	}
	if usage.PromptTokens != 12 || usage.CompletionTokens != 3 || usage.TotalTokens != 15 {
		t.Fatalf("usage = %#v", usage)
	}
}

func TestOpenAICompatibleClientRetriesTransientHTTPError(t *testing.T) {
	attempts := 0
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		attempts++
		if attempts == 1 {
			http.Error(w, "temporary", http.StatusBadGateway)
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		_, _ = w.Write([]byte(`data: {"choices":[{"delta":{"content":"ok"},"finish_reason":""}]}` + "\n\n"))
		_, _ = w.Write([]byte("data: [DONE]\n\n"))
	}))
	defer server.Close()

	client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
	var content strings.Builder
	err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(event StreamEvent) error {
		content.WriteString(event.Content)
		return nil
	})
	if err != nil {
		t.Fatal(err)
	}
	if attempts != 2 {
		t.Fatalf("attempts = %d, want 2", attempts)
	}
	if content.String() != "ok" {
		t.Fatalf("content = %q", content.String())
	}
}

func mustReadBody(t *testing.T, r *http.Request) string {
	t.Helper()
	data, err := io.ReadAll(r.Body)
	if err != nil {
		t.Fatal(err)
	}
	return string(data)
}

func TestOpenAICompatibleClientHTTPError(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		http.Error(w, "bad key", http.StatusUnauthorized)
	}))
	defer server.Close()

	client := NewOpenAICompatibleClient(server.URL, "sk-test", server.Client())
	err := client.ChatStream(context.Background(), ChatRequest{Model: "test"}, func(StreamEvent) error {
		return nil
	})
	if err == nil || !strings.Contains(err.Error(), "status=401") {
		t.Fatalf("err = %v", err)
	}
}