Add go-llm v1 chat Client with OpenAI-compatible and Ollama adapters.
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -0,0 +1,122 @@
|
||||
package gollm_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
gollm "gitea.neitzel.de/konrad/go-llm"
|
||||
)
|
||||
|
||||
func TestChatReturnsAssistantTextFromOpenAICompatibleServer(t *testing.T) {
|
||||
var gotBody map[string]any
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/chat/completions" || r.Method != http.MethodPost {
|
||||
t.Errorf("got %s %s", r.Method, r.URL.Path)
|
||||
}
|
||||
if r.Header.Get("Authorization") != "Bearer secret-key" {
|
||||
t.Errorf("Authorization %q", r.Header.Get("Authorization"))
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
io.WriteString(w, `{"choices":[{"message":{"role":"assistant","content":"Once upon a time."}}]}`)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
client, err := gollm.NewClient(gollm.AdapterOpenAI, srv.URL, "secret-key")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
temp := 0.2
|
||||
maxTok := 64
|
||||
text, err := client.Chat(context.Background(), []gollm.Message{
|
||||
{Role: gollm.RoleSystem, Content: "You write prose."},
|
||||
{Role: gollm.RoleUser, Content: "Begin."},
|
||||
}, gollm.Options{Model: "llama", Temperature: &temp, MaxTokens: &maxTok})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if text != "Once upon a time." {
|
||||
t.Fatalf("got %q", text)
|
||||
}
|
||||
if gotBody["model"] != "llama" {
|
||||
t.Fatalf("model %#v", gotBody["model"])
|
||||
}
|
||||
if gotBody["temperature"] != 0.2 {
|
||||
t.Fatalf("temperature %#v", gotBody["temperature"])
|
||||
}
|
||||
if gotBody["max_tokens"] != float64(64) {
|
||||
t.Fatalf("max_tokens %#v", gotBody["max_tokens"])
|
||||
}
|
||||
msgs, _ := gotBody["messages"].([]any)
|
||||
if len(msgs) != 2 {
|
||||
t.Fatalf("messages %#v", gotBody["messages"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamYieldsOpenAICompatibleChunks(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/chat/completions" {
|
||||
t.Errorf("path %s", r.URL.Path)
|
||||
}
|
||||
var body map[string]any
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if body["stream"] != true {
|
||||
t.Fatalf("stream %#v", body["stream"])
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"Once\"}}]}\n\n")
|
||||
io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\" upon\"}}]}\n\n")
|
||||
io.WriteString(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
client, err := gollm.NewClient(gollm.AdapterOpenAI, srv.URL, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var chunks []string
|
||||
err = client.Stream(context.Background(), []gollm.Message{
|
||||
{Role: gollm.RoleUser, Content: "Begin."},
|
||||
}, gollm.Options{Model: "llama"}, func(chunk string) error {
|
||||
chunks = append(chunks, chunk)
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Join(chunks, "") != "Once upon" {
|
||||
t.Fatalf("chunks %#v", chunks)
|
||||
}
|
||||
}
|
||||
|
||||
func TestListModelsReturnsOpenAICompatibleNames(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/models" || r.Method != http.MethodGet {
|
||||
t.Errorf("got %s %s", r.Method, r.URL.Path)
|
||||
}
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
io.WriteString(w, `{"data":[{"id":"llama"},{"id":"mistral"}]}`)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
|
||||
client, err := gollm.NewClient(gollm.AdapterOpenAI, srv.URL, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
names, err := client.ListModels(context.Background())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(names) != 2 || names[0] != "llama" || names[1] != "mistral" {
|
||||
t.Fatalf("got %#v", names)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user