Add go-llm v1 chat Client with OpenAI-compatible and Ollama adapters.
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
+25
@@ -0,0 +1,25 @@
|
||||
# go-llm
|
||||
|
||||
Shared Go Library for calling a chat LLM: one Client interface with an Ollama-native Adapter and an OpenAI-compatible Adapter. ExApps and Services import `gitea.neitzel.de/konrad/go-llm`. It has no Nextcloud types and no story types.
|
||||
|
||||
## Language
|
||||
|
||||
**Client**:
|
||||
A configured caller bound to one Adapter, base URL, optional API key, and HTTP behaviour. Chat, Stream, and List Models go through that Client.
|
||||
_Avoid_: a chat-history store, Nextcloud AppAPI credentials
|
||||
|
||||
**Adapter**:
|
||||
Which native protocol the Client speaks: Ollama (`/api/chat`) or OpenAI-compatible (`/v1/chat/completions`). vLLM and similar servers use the OpenAI-compatible Adapter.
|
||||
_Avoid_: a third vLLM-native protocol, one Adapter type per model name
|
||||
|
||||
**Message**:
|
||||
One chat turn with a role (system, user, or assistant) and text content. The Library sends and streams Messages; it does not build product prompts.
|
||||
_Avoid_: Story, Chapter, prompt template
|
||||
|
||||
**Stream**:
|
||||
The same Chat request delivered as incremental text chunks instead of one complete reply.
|
||||
_Avoid_: writing chunks to files
|
||||
|
||||
**Model**:
|
||||
A server-side model name as listed by the Adapter. The Library lists names and accepts the name the caller passes; it does not download or install models.
|
||||
_Avoid_: GGUF files, Ollama Modelfile management
|
||||
@@ -0,0 +1,10 @@
|
||||
# go-llm
|
||||
|
||||
Go Library for chat LLM calls (`gitea.neitzel.de/konrad/go-llm`). Domain language is in `CONTEXT.md`. Decisions live in the Workspace at `docs/adr/go-llm/`.
|
||||
|
||||
v1: Chat, Stream, List Models, temperature, and max tokens. Adapters: Ollama native (`/api/chat`, `/api/tags`) and OpenAI-compatible (`/v1/chat/completions`, `/v1/models`). No embeddings, tools, raw generate, Nextcloud types, or story types.
|
||||
|
||||
```go
|
||||
client, err := gollm.NewClient(gollm.AdapterOpenAI, "http://localhost:8000", "")
|
||||
text, err := client.Chat(ctx, []gollm.Message{{Role: gollm.RoleUser, Content: "Hello"}}, gollm.Options{Model: "llama"})
|
||||
```
|
||||
@@ -0,0 +1,279 @@
|
||||
// Package gollm is a chat LLM wire client: Chat, Stream, and List Models
|
||||
// through an Ollama-native or OpenAI-compatible Adapter.
|
||||
package gollm
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Adapter is which native protocol the Client speaks.
|
||||
type Adapter string
|
||||
|
||||
const (
|
||||
AdapterOpenAI Adapter = "openai"
|
||||
AdapterOllama Adapter = "ollama"
|
||||
)
|
||||
|
||||
// Role is a Message role.
|
||||
type Role string
|
||||
|
||||
const (
|
||||
RoleSystem Role = "system"
|
||||
RoleUser Role = "user"
|
||||
RoleAssistant Role = "assistant"
|
||||
)
|
||||
|
||||
// Message is one chat turn.
|
||||
type Message struct {
|
||||
Role Role
|
||||
Content string
|
||||
}
|
||||
|
||||
// Options are Chat/Stream parameters. Temperature and MaxTokens are omitted when nil.
|
||||
type Options struct {
|
||||
Model string
|
||||
Temperature *float64
|
||||
MaxTokens *int
|
||||
}
|
||||
|
||||
// Client is a configured caller bound to one Adapter, base URL, and optional API key.
|
||||
type Client struct {
|
||||
adapter Adapter
|
||||
baseURL string
|
||||
apiKey string
|
||||
http *http.Client
|
||||
}
|
||||
|
||||
// NewClient validates Adapter and base URL. apiKey may be empty.
|
||||
func NewClient(adapter Adapter, baseURL, apiKey string) (*Client, error) {
|
||||
switch adapter {
|
||||
case AdapterOpenAI, AdapterOllama:
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown Adapter %q", adapter)
|
||||
}
|
||||
base := strings.TrimRight(strings.TrimSpace(baseURL), "/")
|
||||
if base == "" {
|
||||
return nil, fmt.Errorf("base URL is required")
|
||||
}
|
||||
return &Client{adapter: adapter, baseURL: base, apiKey: apiKey, http: http.DefaultClient}, nil
|
||||
}
|
||||
|
||||
// Chat returns the assistant text for messages.
|
||||
func (c *Client) Chat(ctx context.Context, messages []Message, opt Options) (string, error) {
|
||||
switch c.adapter {
|
||||
case AdapterOpenAI:
|
||||
return c.openaiChat(ctx, messages, opt)
|
||||
case AdapterOllama:
|
||||
return c.ollamaChat(ctx, messages, opt)
|
||||
default:
|
||||
return "", fmt.Errorf("unknown Adapter %q", c.adapter)
|
||||
}
|
||||
}
|
||||
|
||||
// Stream yields assistant text chunks for the same Chat request.
|
||||
func (c *Client) Stream(ctx context.Context, messages []Message, opt Options, yield func(chunk string) error) error {
|
||||
switch c.adapter {
|
||||
case AdapterOpenAI:
|
||||
return c.openaiStream(ctx, messages, opt, yield)
|
||||
case AdapterOllama:
|
||||
return c.ollamaStream(ctx, messages, opt, yield)
|
||||
default:
|
||||
return fmt.Errorf("unknown Adapter %q", c.adapter)
|
||||
}
|
||||
}
|
||||
|
||||
// ListModels returns server-side Model names.
|
||||
func (c *Client) ListModels(ctx context.Context) ([]string, error) {
|
||||
switch c.adapter {
|
||||
case AdapterOpenAI:
|
||||
return c.openaiListModels(ctx)
|
||||
case AdapterOllama:
|
||||
return c.ollamaListModels(ctx)
|
||||
default:
|
||||
return nil, fmt.Errorf("unknown Adapter %q", c.adapter)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Client) auth(req *http.Request) {
|
||||
if c.apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+c.apiKey)
|
||||
}
|
||||
}
|
||||
|
||||
type openaiChatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []openaiMessage `json:"messages"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
MaxTokens *int `json:"max_tokens,omitempty"`
|
||||
}
|
||||
|
||||
type openaiMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
func openaiMessages(messages []Message) []openaiMessage {
|
||||
out := make([]openaiMessage, len(messages))
|
||||
for i, m := range messages {
|
||||
out[i] = openaiMessage{Role: string(m.Role), Content: m.Content}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (c *Client) openaiChat(ctx context.Context, messages []Message, opt Options) (string, error) {
|
||||
body, err := json.Marshal(openaiChatRequest{
|
||||
Model: opt.Model,
|
||||
Messages: openaiMessages(messages),
|
||||
Temperature: opt.Temperature,
|
||||
MaxTokens: opt.MaxTokens,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.auth(req)
|
||||
res, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
raw, err := io.ReadAll(res.Body)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||||
return "", fmt.Errorf("openai chat: %s: %s", res.Status, truncate(raw))
|
||||
}
|
||||
var parsed struct {
|
||||
Choices []struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &parsed); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if len(parsed.Choices) == 0 {
|
||||
return "", fmt.Errorf("openai chat: no choices")
|
||||
}
|
||||
return parsed.Choices[0].Message.Content, nil
|
||||
}
|
||||
|
||||
func (c *Client) openaiStream(ctx context.Context, messages []Message, opt Options, yield func(string) error) error {
|
||||
body, err := json.Marshal(openaiChatRequest{
|
||||
Model: opt.Model,
|
||||
Messages: openaiMessages(messages),
|
||||
Stream: true,
|
||||
Temperature: opt.Temperature,
|
||||
MaxTokens: opt.MaxTokens,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Accept", "text/event-stream")
|
||||
c.auth(req)
|
||||
res, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||||
raw, _ := io.ReadAll(res.Body)
|
||||
return fmt.Errorf("openai stream: %s: %s", res.Status, truncate(raw))
|
||||
}
|
||||
sc := bufio.NewScanner(res.Body)
|
||||
sc.Buffer(nil, 1024*1024)
|
||||
for sc.Scan() {
|
||||
line := strings.TrimSpace(sc.Text())
|
||||
if line == "" || !strings.HasPrefix(line, "data:") {
|
||||
continue
|
||||
}
|
||||
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
|
||||
if payload == "[DONE]" {
|
||||
return nil
|
||||
}
|
||||
var parsed struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"delta"`
|
||||
} `json:"choices"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(payload), &parsed); err != nil {
|
||||
return err
|
||||
}
|
||||
if len(parsed.Choices) == 0 {
|
||||
continue
|
||||
}
|
||||
chunk := parsed.Choices[0].Delta.Content
|
||||
if chunk == "" {
|
||||
continue
|
||||
}
|
||||
if err := yield(chunk); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return sc.Err()
|
||||
}
|
||||
|
||||
func (c *Client) openaiListModels(ctx context.Context) ([]string, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+"/v1/models", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.auth(req)
|
||||
res, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
raw, err := io.ReadAll(res.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("openai models: %s: %s", res.Status, truncate(raw))
|
||||
}
|
||||
var parsed struct {
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &parsed); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
names := make([]string, 0, len(parsed.Data))
|
||||
for _, d := range parsed.Data {
|
||||
if d.ID != "" {
|
||||
names = append(names, d.ID)
|
||||
}
|
||||
}
|
||||
return names, nil
|
||||
}
|
||||
|
||||
func truncate(b []byte) string {
|
||||
s := strings.TrimSpace(string(b))
|
||||
if len(s) > 200 {
|
||||
return s[:200]
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
package gollm
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type ollamaChatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []openaiMessage `json:"messages"`
|
||||
Stream bool `json:"stream"`
|
||||
Options map[string]any `json:"options,omitempty"`
|
||||
}
|
||||
|
||||
func ollamaOptions(opt Options) map[string]any {
|
||||
if opt.Temperature == nil && opt.MaxTokens == nil {
|
||||
return nil
|
||||
}
|
||||
m := map[string]any{}
|
||||
if opt.Temperature != nil {
|
||||
m["temperature"] = *opt.Temperature
|
||||
}
|
||||
if opt.MaxTokens != nil {
|
||||
m["num_predict"] = *opt.MaxTokens
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func (c *Client) ollamaChat(ctx context.Context, messages []Message, opt Options) (string, error) {
|
||||
body, err := json.Marshal(ollamaChatRequest{
|
||||
Model: opt.Model,
|
||||
Messages: openaiMessages(messages),
|
||||
Stream: false,
|
||||
Options: ollamaOptions(opt),
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/chat", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.auth(req)
|
||||
res, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
raw, err := io.ReadAll(res.Body)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||||
return "", fmt.Errorf("ollama chat: %s: %s", res.Status, truncate(raw))
|
||||
}
|
||||
var parsed struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &parsed); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return parsed.Message.Content, nil
|
||||
}
|
||||
|
||||
func (c *Client) ollamaStream(ctx context.Context, messages []Message, opt Options, yield func(string) error) error {
|
||||
body, err := json.Marshal(ollamaChatRequest{
|
||||
Model: opt.Model,
|
||||
Messages: openaiMessages(messages),
|
||||
Stream: true,
|
||||
Options: ollamaOptions(opt),
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/chat", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
c.auth(req)
|
||||
res, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||||
raw, _ := io.ReadAll(res.Body)
|
||||
return fmt.Errorf("ollama stream: %s: %s", res.Status, truncate(raw))
|
||||
}
|
||||
sc := bufio.NewScanner(res.Body)
|
||||
sc.Buffer(nil, 1024*1024)
|
||||
for sc.Scan() {
|
||||
line := strings.TrimSpace(sc.Text())
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
var parsed struct {
|
||||
Message struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"message"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(line), &parsed); err != nil {
|
||||
return err
|
||||
}
|
||||
if parsed.Message.Content == "" {
|
||||
continue
|
||||
}
|
||||
if err := yield(parsed.Message.Content); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return sc.Err()
|
||||
}
|
||||
|
||||
func (c *Client) ollamaListModels(ctx context.Context) ([]string, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+"/api/tags", nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
c.auth(req)
|
||||
res, err := c.http.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer res.Body.Close()
|
||||
raw, err := io.ReadAll(res.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if res.StatusCode < 200 || res.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("ollama tags: %s: %s", res.Status, truncate(raw))
|
||||
}
|
||||
var parsed struct {
|
||||
Models []struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"models"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &parsed); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
names := make([]string, 0, len(parsed.Models))
|
||||
for _, m := range parsed.Models {
|
||||
if m.Name != "" {
|
||||
names = append(names, m.Name)
|
||||
}
|
||||
}
|
||||
return names, nil
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -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