66 lines
2.5 KiB
Go
66 lines
2.5 KiB
Go
package proxy
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"net/url"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/example/ollama-fair-gateway/internal/cost"
|
|
)
|
|
|
|
func TestNativeMeter(t *testing.T) {
|
|
m := newMeter("ollama", 0)
|
|
m.Feed([]byte("{\"response\":\"x\",\"done\":false}\n{\"done\":true,\"prompt_eval_count\":12,\"eval_count\":4,\"eval_duration\":100}\n"))
|
|
u := m.Finish(100)
|
|
if u.PromptTokens != 12 || u.CompletionTokens != 4 || u.EvalNS != 100 || u.Approximate {
|
|
t.Fatalf("bad usage %#v", u)
|
|
}
|
|
}
|
|
func TestOpenAIMeter(t *testing.T) {
|
|
m := newMeter("openai", 9)
|
|
m.Feed([]byte("data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\ndata: {\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":2}}\n"))
|
|
u := m.Finish(100)
|
|
if u.PromptTokens != 10 || u.CompletionTokens != 2 || u.Approximate {
|
|
t.Fatalf("bad usage %#v", u)
|
|
}
|
|
}
|
|
|
|
func TestAnthropicMeter(t *testing.T) {
|
|
m := newMeter("anthropic", 0)
|
|
m.Feed([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":11,\"output_tokens\":0}}}\n\nevent: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\nevent: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":5}}\n\n"))
|
|
u := m.Finish(100)
|
|
if u.PromptTokens != 11 || u.CompletionTokens != 5 || u.Approximate {
|
|
t.Fatalf("bad anthropic usage %#v", u)
|
|
}
|
|
}
|
|
|
|
func TestForwardProgressObserver(t *testing.T) {
|
|
backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/x-ndjson")
|
|
_, _ = io.WriteString(w, "{\"message\":{\"content\":\"hi\"},\"done\":false}\n")
|
|
if f, ok := w.(http.Flusher); ok {
|
|
f.Flush()
|
|
}
|
|
_, _ = io.WriteString(w, "{\"done\":true,\"prompt_eval_count\":10,\"eval_count\":2}\n")
|
|
}))
|
|
defer backend.Close()
|
|
target, _ := url.Parse(backend.URL)
|
|
in := httptest.NewRequest(http.MethodPost, "http://gateway/api/chat", strings.NewReader(`{"model":"x"}`))
|
|
out := httptest.NewRecorder()
|
|
var calls int
|
|
var lastBytes int64
|
|
var last cost.Usage
|
|
res := New().Forward(context.Background(), out, in, target, strings.NewReader(`{"model":"x"}`), "ollama", 0, func(n int64, u cost.Usage) {
|
|
calls++
|
|
lastBytes = n
|
|
last = u
|
|
})
|
|
if res.Status != http.StatusOK || calls == 0 || lastBytes == 0 || last.PromptTokens != 10 || last.CompletionTokens != 2 {
|
|
t.Fatalf("status=%d calls=%d bytes=%d usage=%#v", res.Status, calls, lastBytes, last)
|
|
}
|
|
}
|