67 lines
2.2 KiB
Go
67 lines
2.2 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestOpenAIStreamTextAndToolCalls(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.Header.Get("Authorization") != "Bearer k" || r.URL.Path != "/v1/chat/completions" {
|
|
http.Error(w, "bad", 400)
|
|
return
|
|
}
|
|
var body map[string]any
|
|
json.NewDecoder(r.Body).Decode(&body)
|
|
if body["stream"] != true || len(body["tools"].([]any)) != 1 {
|
|
http.Error(w, "bad body", 400)
|
|
return
|
|
}
|
|
chunks := []string{
|
|
`{"choices":[{"delta":{"content":"Hel"}}]}`,
|
|
`{"choices":[{"delta":{"content":"lo"}}]}`,
|
|
`{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"c1","function":{"name":"time_now","arguments":"{\"tz\":"}}]}}]}`,
|
|
`{"choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":"\"UTC\"}"}}]},"finish_reason":"tool_calls"}]}`,
|
|
}
|
|
for _, c := range chunks {
|
|
fmt.Fprintf(w, "data: %s\n\n", c)
|
|
}
|
|
fmt.Fprint(w, "data: [DONE]\n\n")
|
|
}))
|
|
defer srv.Close()
|
|
p := NewOpenAI(srv.URL+"/v1/", "k", "m", 0, 5*time.Second)
|
|
var streamed strings.Builder
|
|
resp, err := p.Chat(context.Background(), Request{
|
|
Messages: []Message{{Role: RoleUser, Content: "hi"}},
|
|
Tools: []ToolSpec{{Name: "time_now", Parameters: map[string]any{"type": "object"}}},
|
|
}, func(s string) { streamed.WriteString(s) })
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp.Content != "Hello" || streamed.String() != "Hello" {
|
|
t.Fatalf("content %q streamed %q", resp.Content, streamed.String())
|
|
}
|
|
if len(resp.ToolCalls) != 1 || resp.ToolCalls[0].Arguments != `{"tz":"UTC"}` || resp.ToolCalls[0].ID != "c1" {
|
|
t.Fatalf("tool calls %+v", resp.ToolCalls)
|
|
}
|
|
if resp.FinishReason != "tool_calls" {
|
|
t.Fatal(resp.FinishReason)
|
|
}
|
|
}
|
|
|
|
func TestOpenAIHTTPError(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Error(w, "nope", 401)
|
|
}))
|
|
defer srv.Close()
|
|
_, err := NewOpenAI(srv.URL, "", "m", 0, time.Second).Chat(context.Background(), Request{}, nil)
|
|
if err == nil || !strings.Contains(err.Error(), "401") {
|
|
t.Fatalf("want 401 err, got %v", err)
|
|
}
|
|
}
|