You've already forked agentic-coding-workflow
97 lines
3.1 KiB
Go
97 lines
3.1 KiB
Go
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)
|
|
}
|