NextExpense/internal/ai/ollama.go

package ai

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"image"
	"image/jpeg"
	"io"
	"log"
	"net/http"
	"os"
	"strings"
	"time"
)

const (
	ollamaTimeout = 300 * time.Second
	ollamaMaxSize = 800
	ollamaQuality = 60
)

type ollamaRequest struct {
	Model    string          `json:"model"`
	Messages []ollamaMessage `json:"messages"`
	Stream   bool            `json:"stream"`
	Options  map[string]any  `json:"options,omitempty"`
}

type ollamaMessage struct {
	Role    string   `json:"role"`
	Content string   `json:"content"`
	Images  []string `json:"images,omitempty"`
}

type ollamaResponse struct {
	Message struct {
		Content  string `json:"content"`
		Thinking string `json:"thinking"`
	} `json:"message"`
	Error string `json:"error,omitempty"`
}

type ollamaProvider struct{}

func (p *ollamaProvider) ExtractReceipt(imagePath string) (*ReceiptData, error) {
	imageData, err := readFile(imagePath)
	if err != nil {
		return &ReceiptData{}, fmt.Errorf("read file: %w", err)
	}

	// Compress the image before base64 encoding: resize to max 800px
	// and JPEG-compress at quality 60 to keep the payload manageable
	// for a local 2.3B model.
	compressed, err := compressImageForOllama(imageData)
	if err == nil && len(compressed) < len(imageData) {
		imageData = compressed
		log.Printf("ExtractReceipt [ollama]: compressed %d -> %d bytes", len(imageData), len(compressed))
	}

	baseURL := os.Getenv("AI_BASE_URL")
	if baseURL == "" {
		baseURL = "http://localhost:11434"
	}
	baseURL = strings.TrimRight(baseURL, "/")

	model := os.Getenv("AI_MODEL")
	if model == "" {
		model = "glm-ocr"
	}

	b64Data := base64.StdEncoding.EncodeToString(imageData)

	payload := ollamaRequest{
		Model:  model,
		Stream: false,
		Messages: []ollamaMessage{{
			Role:    "user",
			Content: "Analyze this receipt image. Extract as strict JSON with keys: \"merchant\" (string), \"amount\" (number), \"currency\" (3-letter code), \"category\" (Food/Travel/Lodging/Software/Other), \"date\" (YYYY-MM-DD). Return ONLY valid JSON. No markdown.",
			Images:  []string{b64Data},
		}},
		Options: map[string]any{"num_predict": 256},
	}

	body, _ := json.Marshal(payload)
	apiURL := baseURL + "/api/chat"
	req, _ := http.NewRequest(http.MethodPost, apiURL, bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")

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

	respBody, _ := io.ReadAll(resp.Body)
	if resp.StatusCode != http.StatusOK {
		return &ReceiptData{}, fmt.Errorf("Ollama status %d: %s", resp.StatusCode, string(respBody))
	}

	var apiResp ollamaResponse
	json.Unmarshal(respBody, &apiResp)

	if apiResp.Error != "" {
		return &ReceiptData{}, fmt.Errorf("Ollama error: %s", apiResp.Error)
	}

	contentStr := apiResp.Message.Content
	if contentStr == "" {
		contentStr = apiResp.Message.Thinking
	}
	contentStr = stripMarkdownFences(contentStr)
	var receipt ReceiptData
	if err := json.Unmarshal([]byte(contentStr), &receipt); err != nil {
		return &ReceiptData{}, fmt.Errorf("parse JSON: %w (content: %s)", err, contentStr)
	}

	log.Printf("ExtractReceipt [ollama-%s]: merchant=%q amount=%.2f %s",
		model, receipt.Merchant, receipt.Amount, receipt.Currency)
	return &receipt, nil
}

// compressImageForOllama resizes the image to fit within maxSize and
// re-encodes as JPEG at the given quality to reduce the base64 payload.
func compressImageForOllama(data []byte) ([]byte, error) {
	img, _, err := image.Decode(bytes.NewReader(data))
	if err != nil {
		return nil, err
	}

	bounds := img.Bounds()
	w := bounds.Dx()
	h := bounds.Dy()

	// If already small enough, re-encode at lower quality
	if w <= ollamaMaxSize && h <= ollamaMaxSize {
		var buf bytes.Buffer
		if err := jpeg.Encode(&buf, img, &jpeg.Options{Quality: ollamaQuality}); err != nil {
			return nil, err
		}
		return buf.Bytes(), nil
	}

	// Scale down
	ratio := float64(ollamaMaxSize) / float64(max(w, h))
	nw := int(float64(w) * ratio)
	nh := int(float64(h) * ratio)
	if nw < 1 {
		nw = 1
	}
	if nh < 1 {
		nh = 1
	}

	dst := image.NewRGBA(image.Rect(0, 0, nw, nh))
	for y := 0; y < nh; y++ {
		for x := 0; x < nw; x++ {
			sx := x * w / nw
			sy := y * h / nh
			dst.Set(x, y, img.At(sx, sy))
		}
	}

	var buf bytes.Buffer
	if err := jpeg.Encode(&buf, dst, &jpeg.Options{Quality: ollamaQuality}); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}