You've already forked agentic-coding-workflow
81 lines
2.3 KiB
Go
81 lines
2.3 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
const anthropicStreamFixture = `event: message_start
|
|
data: {"type":"message_start","message":{"id":"m1","content":[],"model":"claude-opus","role":"assistant","stop_reason":null,"stop_sequence":null,"type":"message","usage":{"input_tokens":10,"output_tokens":0}}}
|
|
|
|
event: content_block_start
|
|
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
|
|
|
|
event: content_block_delta
|
|
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
|
|
|
|
event: content_block_delta
|
|
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" world"}}
|
|
|
|
event: content_block_stop
|
|
data: {"type":"content_block_stop","index":0}
|
|
|
|
event: message_delta
|
|
data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":2}}
|
|
|
|
event: message_stop
|
|
data: {"type":"message_stop"}
|
|
|
|
`
|
|
|
|
func TestAnthropic_Stream_TextDeltaAndDone(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
require.Equal(t, "/v1/messages", r.URL.Path)
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.WriteHeader(200)
|
|
_, _ = w.Write([]byte(anthropicStreamFixture))
|
|
}))
|
|
defer srv.Close()
|
|
|
|
c, err := NewAnthropic("test-key", srv.URL)
|
|
require.NoError(t, err)
|
|
|
|
req := Request{
|
|
Messages: []Message{{Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hi"}}}},
|
|
MaxOutputTokens: 100,
|
|
}
|
|
ch, err := c.Stream(context.Background(), "claude-opus-4-7", req)
|
|
require.NoError(t, err)
|
|
|
|
var deltas []string
|
|
var sawUsage, sawDone bool
|
|
var promptTokens, completionTokens int
|
|
for ev := range ch {
|
|
switch ev.Type {
|
|
case EventTextDelta:
|
|
deltas = append(deltas, ev.Text)
|
|
case EventUsage:
|
|
sawUsage = true
|
|
promptTokens = ev.Usage.PromptTokens
|
|
completionTokens = ev.Usage.CompletionTokens
|
|
case EventDone:
|
|
sawDone = true
|
|
}
|
|
}
|
|
require.Equal(t, "Hello world", strings.Join(deltas, ""))
|
|
require.True(t, sawUsage)
|
|
require.Equal(t, 10, promptTokens)
|
|
require.Equal(t, 2, completionTokens)
|
|
require.True(t, sawDone)
|
|
}
|
|
|
|
func TestAnthropic_Provider(t *testing.T) {
|
|
c, _ := NewAnthropic("k", "http://localhost")
|
|
require.Equal(t, "anthropic", c.Provider())
|
|
}
|