@pi-in-go/pigpen-ollama-native 0.1.0

package/extensions/ollama-native/extension.go

@@ -0,0 +1,170 @@
+// Package ollamanative adds Ollama as a model provider for PiG, through Ollama's
+// native /api/chat endpoint.
+//
+// It uses the native API instead of the OpenAI-compatible /v1 one so that tool
+// calls, thinking and images go through Ollama's own message format. The model
+// list is never written down here: it comes from Ollama's /api/tags (with
+// /api/show for each model's context length and capabilities) whenever PiG
+// refreshes provider catalogs. Nothing touches the network at startup, so pig
+// starts the same whether or not Ollama is running; the extension stays silent
+// until PiG asks it for models or for a response.
+//
+// The server address is OLLAMA_HOST, the variable Ollama's own tools read, and
+// defaults to http://127.0.0.1:11434.
+//
+// The first version of this extension was contributed by Peder Munksgaard
+// (pigpen pull request 3); see CREDITS.md.
+package ollamanative
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"os"
+	"strings"
+	"sync"
+	"time"
+
+	sdk "github.com/MichaelKinsy/PiG/extensions/sdk"
+)
+
+// ProviderID is the provider name: models are selected as ollama-native/<model>.
+const ProviderID = "ollama-native"
+
+// Extension returns the extension that registers the ollama-native provider.
+// It does no I/O.
+func Extension() *sdk.Extension {
+	ext := sdk.New(ProviderID)
+	c := &client{
+		host:  func() string { return os.Getenv("OLLAMA_HOST") },
+		http:  &http.Client{Transport: transport()},
+		newID: randomID,
+	}
+	if err := ext.RegisterNativeProvider(newProvider(c)); err != nil {
+		panic(err) // a static declaration: this is a programming error, caught by the tests
+	}
+	ext.Command("ollama", "Ollama: /ollama refresh reads the installed models from the server", refreshCommand)
+	return ext
+}
+
+// refreshCommand is /ollama [refresh]. PiG refreshes provider catalogs when an
+// interactive session starts; this does it on request, which is also how a
+// one-shot `pig -p` run gets its first catalog.
+func refreshCommand(ctx sdk.Context, args string) error {
+	if a := strings.TrimSpace(args); a != "" && a != "refresh" {
+		ctx.Notify("Usage: /ollama refresh", "warning")
+		return nil
+	}
+	allow := true
+	registry := ctx.ModelRegistry()
+	result, err := registry.Refresh(sdk.ModelsRefreshOptions{AllowNetwork: &allow, Providers: []string{ProviderID}})
+	if err != nil {
+		return err
+	}
+	if msg := result.Errors[ProviderID]; msg != "" {
+		ctx.Notify(msg, "error")
+		return nil
+	}
+	all, err := registry.GetAll()
+	if err != nil {
+		return err
+	}
+	n := 0
+	for _, m := range all {
+		if m["provider"] == ProviderID {
+			n++
+		}
+	}
+	ctx.Notify(fmt.Sprintf("Ollama: %d model(s) available as %s/<model>", n, ProviderID), "info")
+	return nil
+}
+
+// newProvider describes the provider. Its catalog starts empty and only changes
+// when the host calls RefreshModels with the network allowed.
+func newProvider(c *client) *sdk.Provider {
+	var mu sync.Mutex
+	models := []map[string]any{}
+	base := c.baseURL()
+	p := &sdk.Provider{
+		ID:      ProviderID,
+		Name:    "Ollama (native API)",
+		BaseURL: &base,
+		Auth: sdk.ProviderAuth{APIKey: &sdk.APIKeyAuth{
+			Name: "Ollama (no key needed; OLLAMA_API_KEY for a protected server)",
+			Resolve: func(sdk.APIKeyAuthInput) (*sdk.AuthResult, error) {
+				key := os.Getenv("OLLAMA_API_KEY")
+				if key == "" {
+					key = localKey
+				}
+				return &sdk.AuthResult{Auth: map[string]any{"apiKey": key}}, nil
+			},
+		}},
+		GetModels: func() ([]map[string]any, error) {
+			mu.Lock()
+			defer mu.Unlock()
+			return append([]map[string]any{}, models...), nil
+		},
+		RefreshModels: func(in sdk.RefreshModelsContext) error {
+			apply := func(next []map[string]any) func() error {
+				return func() error {
+					mu.Lock()
+					models = next
+					mu.Unlock()
+					return nil
+				}
+			}
+			if !in.AllowNetwork {
+				// PiG's start-up pass: no network, and the catalog stored by the last
+				// refresh comes back in Stored. Restoring it is how a one-shot
+				// `pig -p --model ollama-native/<model>` finds its model.
+				restored := restore(in.Stored)
+				if len(restored) == 0 {
+					return nil
+				}
+				_, err := in.Publish(sdk.ModelsPublication{Update: apply(restored)})
+				return err
+			}
+			found, err := c.discover(in.Signal)
+			if err != nil {
+				return err
+			}
+			// The host stores Persist for the next start and applies the new catalog
+			// by calling Update, then reads GetModels.
+			stored, err := json.Marshal(map[string]any{"models": found, "checkedAt": time.Now().Unix()})
+			if err != nil {
+				return err
+			}
+			_, err = in.Publish(sdk.ModelsPublication{Persist: stored, Update: apply(found)})
+			return err
+		},
+		Stream:       c.stream,
+		StreamSimple: c.stream,
+	}
+	return p
+}
+
+// localKey is what Ollama gets as an API key when none is configured: PiG wants
+// a key for every provider, a local Ollama ignores it, and it is never sent.
+const localKey = "ollama"
+
+func randomID() string {
+	b := make([]byte, 6)
+	_, _ = rand.Read(b)
+	return "call_" + hex.EncodeToString(b)
+}
+
+// restore reads back the catalog a refresh stored. Entries without an id are
+// dropped, so a damaged store cannot put a nameless model in the list.
+func restore(stored map[string]any) []map[string]any {
+	items, _ := stored["models"].([]any)
+	var out []map[string]any
+	for _, item := range items {
+		m, _ := item.(map[string]any)
+		if id, _ := m["id"].(string); id != "" {
+			out = append(out, m)
+		}
+	}
+	return out
+}

package/extensions/ollama-native/go.mod

@@ -0,0 +1,6 @@
+module github.com/MichaelKinsy/pigpen/components/ollama-native/extensions/ollama-native
+
+go 1.26
+
+// Pig resolves this requirement to the version-matched staged SDK at build time.
+require github.com/MichaelKinsy/PiG/extensions/sdk v0.4.1

package/extensions/ollama-native/go.sum

@@ -0,0 +1,2 @@
+github.com/MichaelKinsy/PiG/extensions/sdk v0.4.1/go.mod h1:4j0XhCW0EPOBWvqYE7XKjTvIAahI7K1VyyW3KpjTKj8=
+github.com/MichaelKinsy/PiG/extensions/sdk v0.4.1 h1:Ql9sTkc6dvAXaYQLq5nxqmSZgqu8S0L53ZTs8MkOuxw=

package/extensions/ollama-native/helpers_test.go

@@ -0,0 +1,144 @@
+package ollamanative
+
+import (
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	sdk "github.com/MichaelKinsy/PiG/extensions/sdk"
+)
+
+// fixture reads a recorded Ollama response from testdata. The files follow the
+// wire format in Ollama's API documentation (docs/api.md): /api/tags and
+// /api/show answer one JSON document, /api/chat streams newline-delimited JSON.
+func fixture(t *testing.T, name string) []byte {
+	t.Helper()
+	b, err := os.ReadFile(filepath.Join("testdata", name))
+	if err != nil {
+		t.Fatal(err)
+	}
+	return b
+}
+
+// fakeOllama serves the recorded responses and counts every request it gets.
+type fakeOllama struct {
+	*httptest.Server
+	mu       sync.Mutex
+	requests []string
+	chats    []map[string]any
+	// chat is the fixture /api/chat streams, chatStatus its HTTP status.
+	chat       string
+	chatStatus int
+	tags       string
+	hold       chan struct{} // when set, /api/chat sends the first line then waits for it
+}
+
+func newFake(t *testing.T) *fakeOllama {
+	t.Helper()
+	f := &fakeOllama{chat: "chat_text.ndjson", tags: "tags.json", chatStatus: 200}
+	f.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		f.mu.Lock()
+		f.requests = append(f.requests, r.Method+" "+r.URL.Path)
+		f.mu.Unlock()
+		switch r.URL.Path {
+		case "/api/tags":
+			w.Write(fixture(t, f.tags))
+		case "/api/show":
+			var in struct {
+				Model string `json:"model"`
+			}
+			json.NewDecoder(r.Body).Decode(&in)
+			name := strings.NewReplacer(":", "_", "/", "_").Replace(in.Model)
+			b, err := os.ReadFile(filepath.Join("testdata", "show_"+name+".json"))
+			if err != nil {
+				http.Error(w, `{"error":"model not found"}`, 404)
+				return
+			}
+			w.Write(b)
+		case "/api/chat":
+			var in map[string]any
+			json.NewDecoder(r.Body).Decode(&in)
+			f.mu.Lock()
+			f.chats = append(f.chats, in)
+			f.mu.Unlock()
+			w.WriteHeader(f.chatStatus)
+			body := fixture(t, f.chat)
+			if f.hold != nil {
+				line, rest, _ := strings.Cut(string(body), "\n")
+				io.WriteString(w, line+"\n")
+				w.(http.Flusher).Flush()
+				select {
+				case <-f.hold:
+				case <-r.Context().Done():
+					return
+				}
+				body = []byte(rest)
+			}
+			w.Write(body)
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	t.Cleanup(f.Close)
+	return f
+}
+
+func (f *fakeOllama) count() int { f.mu.Lock(); defer f.mu.Unlock(); return len(f.requests) }
+
+// newClient builds a client aimed at base with deterministic tool-call ids.
+func newClient(base string) *client {
+	n := 0
+	return &client{host: func() string { return base }, http: &http.Client{}, newID: func() string { n++; return "call_t" + string(rune('0'+n)) }}
+}
+
+type streamed struct {
+	events []map[string]any
+	result map[string]any
+}
+
+func (s streamed) types() []string {
+	var out []string
+	for _, e := range s.events {
+		out = append(out, e["type"].(string))
+	}
+	return out
+}
+
+func (s streamed) terminal() map[string]any { return s.events[len(s.events)-1] }
+
+// collect drains a provider stream.
+func collect(t *testing.T, stream *sdk.ModelEventStream) streamed {
+	t.Helper()
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	var s streamed
+	for e := range stream.Events(ctx) {
+		s.events = append(s.events, e)
+	}
+	if ctx.Err() != nil {
+		t.Fatalf("stream did not finish: %v (events %v)", ctx.Err(), s.types())
+	}
+	s.result = stream.Result()
+	return s
+}
+
+func userTranscript(text string) map[string]any {
+	return map[string]any{"messages": []any{map[string]any{"role": "user", "content": text, "timestamp": 1}}}
+}
+
+var graniteModel = map[string]any{"id": "granite4.1:3b", "name": "granite4.1:3b", "api": "ollama-native", "provider": "ollama-native", "reasoning": false}
+
+func httpHandler(see func(map[string][]string), next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		see(r.Header)
+		next.ServeHTTP(w, r)
+	})
+}

package/extensions/ollama-native/models.go

@@ -0,0 +1,235 @@
+package ollamanative
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"net/url"
+	"slices"
+	"strings"
+	"time"
+)
+
+const (
+	defaultHost          = "http://127.0.0.1:11434"
+	defaultPort          = "11434"
+	defaultContextWindow = 32768
+	defaultMaxTokens     = 8192
+)
+
+// client talks to one Ollama server. host is read on every use, so a changed
+// OLLAMA_HOST applies without restarting pig.
+type client struct {
+	host  func() string
+	http  *http.Client
+	newID func() string
+}
+
+// transport bounds connecting, not the response: a large model can take a long
+// time to load before the first token, and a long answer is a long stream.
+func transport() *http.Transport {
+	t := http.DefaultTransport.(*http.Transport).Clone()
+	t.DialContext = (&net.Dialer{Timeout: 5 * time.Second}).DialContext
+	return t
+}
+
+func (c *client) baseURL() string { return normalizeHost(c.host()) }
+
+// normalizeHost turns an OLLAMA_HOST value into a base URL the way Ollama's own
+// client reads it: a bare host or host:port means http on port 11434, ":port"
+// means this machine, and an explicit scheme without a port means that scheme's
+// default port.
+func normalizeHost(value string) string {
+	value = strings.TrimSpace(value)
+	if value == "" {
+		return defaultHost
+	}
+	scheme, rest := "http", value
+	explicit := false
+	if s, r, ok := strings.Cut(value, "://"); ok {
+		scheme, rest, explicit = strings.ToLower(s), r, true
+	}
+	u, err := url.Parse(scheme + "://" + rest)
+	if err != nil || (u.Scheme != "http" && u.Scheme != "https") {
+		return defaultHost
+	}
+	host, port := u.Hostname(), u.Port()
+	if host == "" {
+		host = "127.0.0.1"
+	}
+	if port == "" {
+		switch {
+		case !explicit:
+			port = defaultPort
+		case scheme == "https":
+			port = "443"
+		default:
+			port = "80"
+		}
+	}
+	return scheme + "://" + net.JoinHostPort(host, port) + strings.TrimRight(u.Path, "/")
+}
+
+// unreachable is the error for an Ollama that is not running (or not there).
+func unreachable(base string, err error) error {
+	return fmt.Errorf("Ollama is not reachable at %s (%v). Start it with \"ollama serve\", or set OLLAMA_HOST to where it listens.", base, rootCause(err))
+}
+
+func rootCause(err error) error {
+	var op *net.OpError
+	if errors.As(err, &op) && op.Err != nil {
+		return fmt.Errorf("%s: %w", op.Op, op.Err)
+	}
+	var ue *url.Error
+	if errors.As(err, &ue) {
+		return ue.Err
+	}
+	return err
+}
+
+// isDialFailure reports a request that never reached a server.
+func isDialFailure(err error) bool {
+	var op *net.OpError
+	return errors.As(err, &op) && op.Op == "dial"
+}
+
+func (c *client) do(ctx context.Context, method, path string, body any, timeout time.Duration) (*http.Response, context.CancelFunc, error) {
+	if ctx == nil {
+		ctx = context.Background()
+	}
+	cancel := func() {}
+	if timeout > 0 {
+		ctx, cancel = context.WithTimeout(ctx, timeout)
+	}
+	var reader io.Reader
+	if body != nil {
+		raw, err := json.Marshal(body)
+		if err != nil {
+			cancel()
+			return nil, nil, err
+		}
+		reader = bytes.NewReader(raw)
+	}
+	req, err := http.NewRequestWithContext(ctx, method, c.baseURL()+path, reader)
+	if err != nil {
+		cancel()
+		return nil, nil, err
+	}
+	if body != nil {
+		req.Header.Set("Content-Type", "application/json")
+	}
+	resp, err := c.http.Do(req)
+	if err != nil {
+		cancel()
+		return nil, nil, err
+	}
+	return resp, cancel, nil
+}
+
+// apiError reads Ollama's {"error": "..."} body, or falls back to the raw text.
+func apiError(resp *http.Response) string {
+	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 64<<10))
+	var e struct {
+		Error string `json:"error"`
+	}
+	if json.Unmarshal(raw, &e) == nil && e.Error != "" {
+		return e.Error
+	}
+	if text := strings.TrimSpace(string(raw)); text != "" {
+		return text
+	}
+	return http.StatusText(resp.StatusCode)
+}
+
+// discover lists the models Ollama has installed. It reads /api/tags, then asks
+// /api/show about each model for its context length and capabilities. A model
+// /api/show cannot describe is still listed, with defaults.
+func (c *client) discover(ctx context.Context) ([]map[string]any, error) {
+	resp, cancel, err := c.do(ctx, http.MethodGet, "/api/tags", nil, 10*time.Second)
+	if err != nil {
+		if ctx != nil && ctx.Err() != nil {
+			return nil, ctx.Err()
+		}
+		return nil, unreachable(c.baseURL(), err)
+	}
+	defer cancel()
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("Ollama at %s answered %d to /api/tags: %s", c.baseURL(), resp.StatusCode, apiError(resp))
+	}
+	var tags struct {
+		Models []struct {
+			Name string `json:"name"`
+		} `json:"models"`
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&tags); err != nil {
+		return nil, fmt.Errorf("Ollama at %s sent a model list that is not JSON: %w", c.baseURL(), err)
+	}
+	models := make([]map[string]any, 0, len(tags.Models))
+	for _, m := range tags.Models {
+		if m.Name == "" {
+			continue
+		}
+		info := c.show(ctx, m.Name)
+		if info.capabilities != nil && !slices.Contains(info.capabilities, "completion") {
+			continue // embedding-only models cannot chat
+		}
+		input := []string{"text"}
+		if slices.Contains(info.capabilities, "vision") {
+			input = append(input, "image")
+		}
+		window := info.contextLength
+		if window <= 0 {
+			window = defaultContextWindow
+		}
+		models = append(models, map[string]any{
+			"id":            m.Name,
+			"name":          m.Name,
+			"api":           ProviderID,
+			"reasoning":     slices.Contains(info.capabilities, "thinking"),
+			"input":         input,
+			"contextWindow": window,
+			"maxTokens":     min(window, defaultMaxTokens),
+			"cost":          map[string]any{"input": 0, "output": 0, "cacheRead": 0, "cacheWrite": 0},
+		})
+	}
+	return models, nil
+}
+
+type modelInfo struct {
+	contextLength int
+	capabilities  []string
+}
+
+// show reads a model's details. Failure is not an error: the zero value means
+// "unknown", and discover then uses defaults.
+func (c *client) show(ctx context.Context, name string) modelInfo {
+	resp, cancel, err := c.do(ctx, http.MethodPost, "/api/show", map[string]any{"model": name}, 5*time.Second)
+	if err != nil {
+		return modelInfo{}
+	}
+	defer cancel()
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return modelInfo{}
+	}
+	var out struct {
+		ModelInfo    map[string]any `json:"model_info"`
+		Capabilities []string       `json:"capabilities"`
+	}
+	if json.NewDecoder(resp.Body).Decode(&out) != nil {
+		return modelInfo{}
+	}
+	info := modelInfo{capabilities: out.Capabilities}
+	if arch, _ := out.ModelInfo["general.architecture"].(string); arch != "" {
+		if n, ok := out.ModelInfo[arch+".context_length"].(float64); ok {
+			info.contextLength = int(n)
+		}
+	}
+	return info
+}