95 lines
2.7 KiB
Go
95 lines
2.7 KiB
Go
package ollama
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type Client struct {
|
|
BaseURL, ChatModel, EmbeddingModel string
|
|
HTTP *http.Client
|
|
}
|
|
|
|
func New(base, chat, embed string) *Client {
|
|
return &Client{BaseURL: strings.TrimRight(base, "/"), ChatModel: chat, EmbeddingModel: embed, HTTP: &http.Client{Timeout: 8 * time.Minute}}
|
|
}
|
|
func (c *Client) Embed(ctx context.Context, texts []string) ([][]float64, error) {
|
|
body := map[string]any{"model": c.EmbeddingModel, "input": texts, "truncate": true}
|
|
var out struct {
|
|
Embeddings [][]float64 `json:"embeddings"`
|
|
}
|
|
if err := c.post(ctx, "/api/embed", body, &out); err != nil {
|
|
return nil, err
|
|
}
|
|
if len(out.Embeddings) != len(texts) {
|
|
return nil, fmt.Errorf("ollama returned %d embeddings for %d inputs", len(out.Embeddings), len(texts))
|
|
}
|
|
return out.Embeddings, nil
|
|
}
|
|
func (c *Client) ChatJSON(ctx context.Context, system, user string, schema any, target any) error {
|
|
body := map[string]any{"model": c.ChatModel, "messages": []map[string]string{{"role": "system", "content": system}, {"role": "user", "content": user}}, "stream": false, "think": false, "format": schema, "options": map[string]any{"temperature": 0}}
|
|
var env struct {
|
|
Message struct {
|
|
Content string `json:"content"`
|
|
} `json:"message"`
|
|
}
|
|
if err := c.post(ctx, "/api/chat", body, &env); err != nil {
|
|
return err
|
|
}
|
|
if strings.TrimSpace(env.Message.Content) == "" {
|
|
return fmt.Errorf("empty Ollama response")
|
|
}
|
|
if err := json.Unmarshal([]byte(env.Message.Content), target); err != nil {
|
|
return fmt.Errorf("decode structured response: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
func (c *Client) Ping(ctx context.Context) error {
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.BaseURL+"/api/tags", nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
resp, err := c.HTTP.Do(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode/100 != 2 {
|
|
return fmt.Errorf("ollama HTTP %d", resp.StatusCode)
|
|
}
|
|
return nil
|
|
}
|
|
func (c *Client) post(ctx context.Context, path string, in, out any) error {
|
|
b, err := json.Marshal(in)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+path, bytes.NewReader(b))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
resp, err := c.HTTP.Do(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer resp.Body.Close()
|
|
data, err := io.ReadAll(io.LimitReader(resp.Body, 16<<20))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if resp.StatusCode/100 != 2 {
|
|
return fmt.Errorf("ollama %s HTTP %d: %s", path, resp.StatusCode, strings.TrimSpace(string(data)))
|
|
}
|
|
if err := json.Unmarshal(data, out); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|