102 lines
2.6 KiB
Go
102 lines
2.6 KiB
Go
package gollama
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/url"
|
|
"os"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
)
|
|
|
|
// Client wraps the official Ollama API client with ergonomic defaults.
|
|
type Client struct {
|
|
api *api.Client
|
|
host *url.URL
|
|
defaultModel string
|
|
defaultOpts OptionSet
|
|
}
|
|
|
|
// ClientOption configures a Client.
|
|
type ClientOption func(*Client)
|
|
|
|
// WithDefaultModel sets the model used when calls do not specify one.
|
|
func WithDefaultModel(name string) ClientOption {
|
|
return func(c *Client) { c.defaultModel = name }
|
|
}
|
|
|
|
// WithDefaultOptions sets client-wide generation options.
|
|
func WithDefaultOptions(opts OptionSet) ClientOption {
|
|
return func(c *Client) { c.defaultOpts = opts }
|
|
}
|
|
|
|
// WithHTTPClient sets the HTTP client used for API requests.
|
|
func WithHTTPClient(httpClient *http.Client) ClientOption {
|
|
return func(c *Client) {
|
|
if httpClient != nil && c.host != nil {
|
|
c.api = api.NewClient(c.host, httpClient)
|
|
}
|
|
}
|
|
}
|
|
|
|
// FromEnv creates a Client using OLLAMA_HOST (default http://127.0.0.1:11434).
|
|
func FromEnv(opts ...ClientOption) (*Client, error) {
|
|
apiClient, err := api.ClientFromEnvironment()
|
|
if err != nil {
|
|
return nil, wrapError("connect", "", err)
|
|
}
|
|
c := newClient(apiClient, nil, opts...)
|
|
// Re-read host from env for WithHTTPClient; api.ClientFromEnvironment uses envconfig.Host().
|
|
if u, err := url.Parse(hostFromEnv()); err == nil {
|
|
c.host = u
|
|
}
|
|
return c, nil
|
|
}
|
|
|
|
// New creates a Client for the given Ollama host URL.
|
|
func New(host string, opts ...ClientOption) (*Client, error) {
|
|
u, err := url.Parse(host)
|
|
if err != nil {
|
|
return nil, wrapError("connect", "", fmt.Errorf("invalid host %q: %w", host, err))
|
|
}
|
|
apiClient := api.NewClient(u, http.DefaultClient)
|
|
return newClient(apiClient, u, opts...), nil
|
|
}
|
|
|
|
func newClient(apiClient *api.Client, host *url.URL, opts ...ClientOption) *Client {
|
|
c := &Client{api: apiClient, host: host}
|
|
for _, opt := range opts {
|
|
opt(c)
|
|
}
|
|
return c
|
|
}
|
|
|
|
func hostFromEnv() string {
|
|
if v := os.Getenv("OLLAMA_HOST"); v != "" {
|
|
return v
|
|
}
|
|
return "http://127.0.0.1:11434"
|
|
}
|
|
|
|
// Ping checks whether the Ollama server is reachable.
|
|
func (c *Client) Ping(ctx context.Context) error {
|
|
return wrapError("ping", "", c.api.Heartbeat(ctx))
|
|
}
|
|
|
|
// DefaultModel returns the client's default model name.
|
|
func (c *Client) DefaultModel() string { return c.defaultModel }
|
|
|
|
// API returns the underlying official api.Client for advanced use.
|
|
func (c *Client) API() *api.Client { return c.api }
|
|
|
|
func (c *Client) resolveModel(model string) (string, error) {
|
|
if model != "" {
|
|
return model, nil
|
|
}
|
|
if c.defaultModel != "" {
|
|
return c.defaultModel, nil
|
|
}
|
|
return "", errNoModel()
|
|
}
|