118 lines
3.3 KiB
Go
118 lines
3.3 KiB
Go
package gollm_test
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
gollm "gitea.neitzel.de/konrad/go-llm"
|
|
)
|
|
|
|
func TestChatReturnsAssistantTextFromOllama(t *testing.T) {
|
|
var gotBody map[string]any
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/api/chat" || r.Method != http.MethodPost {
|
|
t.Errorf("got %s %s", r.Method, r.URL.Path)
|
|
}
|
|
if err := json.NewDecoder(r.Body).Decode(&gotBody); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
io.WriteString(w, `{"message":{"role":"assistant","content":"Once upon a time."},"done":true}`)
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
|
|
client, err := gollm.NewClient(gollm.AdapterOllama, srv.URL, "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
temp := 0.2
|
|
maxTok := 64
|
|
text, err := client.Chat(context.Background(), []gollm.Message{
|
|
{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["stream"] != false {
|
|
t.Fatalf("stream %#v", gotBody["stream"])
|
|
}
|
|
opts, _ := gotBody["options"].(map[string]any)
|
|
if opts["temperature"] != 0.2 {
|
|
t.Fatalf("temperature %#v", opts["temperature"])
|
|
}
|
|
if opts["num_predict"] != float64(64) {
|
|
t.Fatalf("num_predict %#v", opts["num_predict"])
|
|
}
|
|
}
|
|
|
|
func TestStreamYieldsOllamaChunks(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/api/chat" {
|
|
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", "application/x-ndjson")
|
|
io.WriteString(w, `{"message":{"content":"Once"},"done":false}`+"\n")
|
|
io.WriteString(w, `{"message":{"content":" upon"},"done":true}`+"\n")
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
|
|
client, err := gollm.NewClient(gollm.AdapterOllama, 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 TestListModelsReturnsOllamaNames(t *testing.T) {
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if r.URL.Path != "/api/tags" || r.Method != http.MethodGet {
|
|
t.Errorf("got %s %s", r.Method, r.URL.Path)
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
io.WriteString(w, `{"models":[{"name":"llama3:latest"},{"name":"mistral"}]}`)
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
|
|
client, err := gollm.NewClient(gollm.AdapterOllama, 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] != "llama3:latest" || names[1] != "mistral" {
|
|
t.Fatalf("got %#v", names)
|
|
}
|
|
}
|