package llm import ( "context" "net/http" "net/http/httptest" "strings" "testing" "github.com/stretchr/testify/require" ) const openaiStreamFixture = `data: {"id":"c1","object":"chat.completion.chunk","model":"gpt-4o","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]} data: {"id":"c1","object":"chat.completion.chunk","model":"gpt-4o","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":null}]} data: {"id":"c1","object":"chat.completion.chunk","model":"gpt-4o","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}} data: [DONE] ` const openaiReasoningFixture = `data: {"id":"c1","object":"chat.completion.chunk","model":"deepseek-reasoner","choices":[{"index":0,"delta":{"reasoning_content":"Let me think"},"finish_reason":null}]} data: {"id":"c1","object":"chat.completion.chunk","model":"deepseek-reasoner","choices":[{"index":0,"delta":{"content":"Done"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":1}} data: [DONE] ` func TestOpenAI_Stream_TextDelta(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte(openaiStreamFixture)) })) defer srv.Close() c, _ := NewOpenAI("k", srv.URL) ch, err := c.Stream(context.Background(), "gpt-4o", Request{ Messages: []Message{{Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hi"}}}}, MaxOutputTokens: 100, }) require.NoError(t, err) var deltas []string var sawDone, sawUsage bool for ev := range ch { switch ev.Type { case EventTextDelta: deltas = append(deltas, ev.Text) case EventUsage: sawUsage = true require.Equal(t, 10, ev.Usage.PromptTokens) require.Equal(t, 2, ev.Usage.CompletionTokens) case EventDone: sawDone = true } } require.Equal(t, "Hello world", strings.Join(deltas, "")) require.True(t, sawUsage) require.True(t, sawDone) } func TestOpenAI_Stream_ReasoningContent_Compatible(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") _, _ = w.Write([]byte(openaiReasoningFixture)) })) defer srv.Close() c, _ := NewOpenAI("k", srv.URL) ch, _ := c.Stream(context.Background(), "deepseek-reasoner", Request{ Messages: []Message{{Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hi"}}}}, }) var thinking, text string for ev := range ch { switch ev.Type { case EventThinkingDelta: thinking += ev.Text case EventTextDelta: text += ev.Text } } require.Equal(t, "Let me think", thinking) require.Equal(t, "Done", text) } func TestOpenAI_CountTokens_LocalTiktoken(t *testing.T) { c, _ := NewOpenAI("k", "") got, err := c.CountTokens(context.Background(), "gpt-4o", Request{ SystemPrompt: "You are helpful.", Messages: []Message{{Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hello world"}}}}, }) require.NoError(t, err) require.Greater(t, got, 0) }