archivdms/internal/llm/ollama.go

// Package llm is a minimal HTTP client for an EXTERNAL, already-running Ollama
// server (never installed on the archivdms host — the base URL is provided per
// tenant, see internal/storage/ollama_config.go). It intentionally does no
// retrying, no connection pooling and no streaming: a single direct call to
// Ollama's /api/generate endpoint, CGO-free, net/http only.
package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

// generateRequest is the JSON body POSTed to <base_url>/api/generate. stream is
// always false (we want the whole answer at once) and format is "json" so the
// model is nudged to emit valid JSON in the response field.
type generateRequest struct {
	Model  string `json:"model"`
	Prompt string `json:"prompt"`
	Stream bool   `json:"stream"`
	Format string `json:"format"`
}

// generateResponse is the (non-streaming) envelope Ollama returns; the actual
// model output is the Response string, which — because we requested
// format=json — is itself a JSON document the caller parses structurally.
type generateResponse struct {
	Response string `json:"response"`
	Done     bool   `json:"done"`
	Error    string `json:"error"`
}

// tagsResponse is the JSON envelope Ollama returns from GET /api/tags: a list
// of the models installed on that server. Only the name is consumed here.
type tagsResponse struct {
	Models []struct {
		Name string `json:"name"`
	} `json:"models"`
}

// ListModels performs a single blocking GET /api/tags call against the given
// Ollama base URL and returns the names of the models installed on that server,
// so the frontend can offer a picklist instead of a free-text model field. Any
// network error, timeout or non-200 status yields a clear error — there is NO
// silent fallback. The returned slice is always non-nil (make, never nil) so it
// JSON-encodes as [] rather than null.
func ListModels(ctx context.Context, baseURL string, timeout time.Duration) ([]string, error) {
	baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
	if baseURL == "" {
		return nil, fmt.Errorf("llm: ollama base_url is empty")
	}
	if timeout <= 0 {
		timeout = 10 * time.Second
	}

	reqCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, baseURL+"/api/tags", nil)
	if err != nil {
		return nil, fmt.Errorf("llm: build tags request: %w", err)
	}

	client := &http.Client{Timeout: timeout}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("llm: ollama tags request failed: %w", err)
	}
	defer resp.Body.Close()

	raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if err != nil {
		return nil, fmt.Errorf("llm: read ollama tags response: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("llm: ollama returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
	}

	var env tagsResponse
	if err := json.Unmarshal(raw, &env); err != nil {
		return nil, fmt.Errorf("llm: parse ollama tags envelope: %w", err)
	}

	out := make([]string, 0, len(env.Models))
	for _, m := range env.Models {
		if name := strings.TrimSpace(m.Name); name != "" {
			out = append(out, name)
		}
	}
	return out, nil
}

// GenerateJSON performs a single blocking /api/generate call against the given
// Ollama base URL and returns the model's inner `response` field as a
// json.RawMessage (the caller unmarshals it into its own schema). Any network
// error, timeout, non-200 status, Ollama-reported error or empty/invalid outer
// response yields a clear error — there is NO silent fallback, so the caller
// can report to the frontend exactly that Ollama did not answer.
func GenerateJSON(ctx context.Context, baseURL, model string, timeout time.Duration, prompt string) (json.RawMessage, error) {
	baseURL = strings.TrimRight(strings.TrimSpace(baseURL), "/")
	if baseURL == "" {
		return nil, fmt.Errorf("llm: ollama base_url is empty")
	}
	if model == "" {
		return nil, fmt.Errorf("llm: ollama model is empty")
	}
	if timeout <= 0 {
		timeout = 30 * time.Second
	}

	body, err := json.Marshal(generateRequest{
		Model:  model,
		Prompt: prompt,
		Stream: false,
		Format: "json",
	})
	if err != nil {
		return nil, fmt.Errorf("llm: marshal request: %w", err)
	}

	reqCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, baseURL+"/api/generate", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("llm: build request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	client := &http.Client{Timeout: timeout}
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("llm: ollama request failed: %w", err)
	}
	defer resp.Body.Close()

	// Cap the read so a misbehaving/unexpected endpoint cannot exhaust memory.
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if err != nil {
		return nil, fmt.Errorf("llm: read ollama response: %w", err)
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("llm: ollama returned status %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
	}

	var env generateResponse
	if err := json.Unmarshal(raw, &env); err != nil {
		return nil, fmt.Errorf("llm: parse ollama envelope: %w", err)
	}
	if env.Error != "" {
		return nil, fmt.Errorf("llm: ollama error: %s", env.Error)
	}
	inner := strings.TrimSpace(env.Response)
	if inner == "" {
		return nil, fmt.Errorf("llm: ollama returned an empty response")
	}
	if !json.Valid([]byte(inner)) {
		return nil, fmt.Errorf("llm: ollama response field is not valid JSON")
	}
	return json.RawMessage(inner), nil
}