ollie/skills/skills.go

// Package skills provides skill discovery, embedding, and matching.
package skills

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"sync"

	"ollie/embedding"
	"ollie/util"
)

// Skill represents a discovered skill with its metadata and content.
type Skill struct {
	Name        string
	Description string
	Path        string // path to SKILL.md
	Content     string // full content (including frontmatter)
}

// Index holds precomputed skill embeddings for fast matching.
type Index struct {
	model  *embedding.Model
	skills []Skill
	vecs   []embedding.Vector
	mu     sync.RWMutex
}

// NewIndex creates a new skill index with precomputed embeddings.
// modelDir is the path to the directory containing the ONNX model.
// skillDirs is a list of directories to scan for skills.
func NewIndex(modelDir string, skillDirs []string) (*Index, error) {
	model, err := embedding.LoadModel(modelDir)
	if err != nil {
		return nil, fmt.Errorf("load embedding model: %w", err)
	}

	idx := &Index{model: model}
	if err := idx.loadSkills(skillDirs); err != nil {
		model.Close()
		return nil, err
	}
	return idx, nil
}

// NewIndexFromSkills creates an index from pre-loaded skills.
// Takes ownership of the model.
func NewIndexFromSkills(model *embedding.Model, skills []Skill) (*Index, error) {
	idx := &Index{model: model, skills: skills}

	// Compute embeddings
	vecs := make([]embedding.Vector, len(skills))
	for i, skill := range skills {
		vec, err := model.Embed(skill.Description)
		if err != nil {
			return nil, fmt.Errorf("embed skill %q: %w", skill.Name, err)
		}
		vecs[i] = vec
	}
	idx.vecs = vecs

	return idx, nil
}

// Close releases resources.
func (idx *Index) Close() error {
	if idx.model != nil {
		return idx.model.Close()
	}
	return nil
}

// Match returns skills matching the query, sorted by relevance.
// threshold is the minimum cosine similarity (0-1) to include.
// limit is the maximum number of skills to return (0 for no limit).
func (idx *Index) Match(query string, threshold float32, limit int) ([]MatchResult, error) {
	idx.mu.RLock()
	defer idx.mu.RUnlock()

	if len(idx.skills) == 0 {
		return nil, nil
	}

	qvec, err := idx.model.Embed(query)
	if err != nil {
		return nil, fmt.Errorf("embed query: %w", err)
	}

	var results []MatchResult
	for i, skill := range idx.skills {
		sim := embedding.CosineSimilarity(qvec, idx.vecs[i])
		if sim >= threshold {
			results = append(results, MatchResult{
				Skill: skill,
				Score: sim,
			})
		}
	}

	// Sort by score descending
	sort.Slice(results, func(i, j int) bool {
		return results[i].Score > results[j].Score
	})

	if limit > 0 && len(results) > limit {
		results = results[:limit]
	}
	return results, nil
}

// MatchResult holds a matched skill and its similarity score.
type MatchResult struct {
	Skill Skill
	Score float32
}

// All returns all indexed skills.
func (idx *Index) All() []Skill {
	idx.mu.RLock()
	defer idx.mu.RUnlock()
	return append([]Skill(nil), idx.skills...)
}

// Reload rescans skill directories and updates embeddings.
func (idx *Index) Reload(skillDirs []string) error {
	idx.mu.Lock()
	defer idx.mu.Unlock()
	return idx.loadSkillsLocked(skillDirs)
}

func (idx *Index) loadSkills(dirs []string) error {
	idx.mu.Lock()
	defer idx.mu.Unlock()
	return idx.loadSkillsLocked(dirs)
}

func (idx *Index) loadSkillsLocked(dirs []string) error {
	var skills []Skill
	seen := make(map[string]bool)

	for _, dir := range dirs {
		dir = util.ExpandHome(dir)
		entries, err := os.ReadDir(dir)
		if err != nil {
			continue // skip missing directories
		}
		for _, entry := range entries {
			if !entry.IsDir() {
				continue
			}
			name := entry.Name()
			if seen[name] {
				continue // first dir wins
			}
			skillPath := filepath.Join(dir, name, "SKILL.md")
			skill, err := loadSkill(skillPath)
			if err != nil {
				continue // skip invalid skills
			}
			seen[name] = true
			skills = append(skills, skill)
		}
	}

	// Compute embeddings
	vecs := make([]embedding.Vector, len(skills))
	for i, skill := range skills {
		vec, err := idx.model.Embed(skill.Description)
		if err != nil {
			return fmt.Errorf("embed skill %q: %w", skill.Name, err)
		}
		vecs[i] = vec
	}

	idx.skills = skills
	idx.vecs = vecs
	return nil
}

func loadSkill(path string) (Skill, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return Skill{}, err
	}

	name, desc, err := parseFrontmatter(data)
	if err != nil {
		return Skill{}, err
	}

	return Skill{
		Name:        name,
		Description: desc,
		Path:        path,
		Content:     string(data),
	}, nil
}

// parseFrontmatter extracts name and description from YAML frontmatter.
func parseFrontmatter(data []byte) (name, description string, err error) {
	scanner := bufio.NewScanner(bytes.NewReader(data))

	// First line must be ---
	if !scanner.Scan() || strings.TrimSpace(scanner.Text()) != "---" {
		return "", "", fmt.Errorf("missing frontmatter")
	}

	// Read until closing ---
	for scanner.Scan() {
		line := scanner.Text()
		if strings.TrimSpace(line) == "---" {
			break
		}
		if key, val, ok := strings.Cut(line, ":"); ok {
			key = strings.TrimSpace(key)
			val = strings.TrimSpace(val)
			switch key {
			case "name":
				name = val
			case "description":
				description = val
			}
		}
	}

	if name == "" {
		return "", "", fmt.Errorf("missing name in frontmatter")
	}
	if description == "" {
		return "", "", fmt.Errorf("missing description in frontmatter")
	}
	return name, description, nil
}

// DefaultSkillDirs returns the skill directories to search.
// Reads from ~/.config/ollie/embedding.conf if it exists (one path per line).
// Supports ~ and $VAR expansion. Otherwise, defaults to ~/.config/ollie/skills.
func DefaultSkillDirs() []string {
	confPath := filepath.Join(util.CfgDir(), "embedding.conf")
	data, err := os.ReadFile(confPath)
	if err == nil {
		var dirs []string
		for _, line := range strings.Split(string(data), "\n") {
			line = strings.TrimSpace(line)
			if line != "" && !strings.HasPrefix(line, "#") {
				line = os.ExpandEnv(line)
				dirs = append(dirs, util.ExpandHome(line))
			}
		}
		if len(dirs) > 0 {
			return dirs
		}
	}
	return []string{filepath.Join(util.CfgDir(), "skills")}
}

// DefaultModelDir returns the default model directory.
func DefaultModelDir() string {
	return filepath.Join(util.DataDir(), "models")
}