ollie/skills/skills.go

// Package skills provides skill discovery and matching.
// Skills are markdown knowledge modules with YAML frontmatter.
package skills

import (
	"bufio"
	"bytes"
	"fmt"
	"os"
	"path/filepath"
	"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
	index  *embedding.Index
	mu     sync.RWMutex
}

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

// 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
}

// 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 idx.index == nil {
		return nil, nil
	}

	results, err := idx.index.Match(query, threshold, limit)
	if err != nil {
		return nil, err
	}

	// Convert embedding.MatchResult to MatchResult with full Skill
	out := make([]MatchResult, len(results))
	for i, r := range results {
		// Find the skill by name
		for _, skill := range idx.skills {
			if skill.Name == r.Item.Name {
				out[i] = MatchResult{Skill: skill, Score: r.Score}
				break
			}
		}
	}
	return out, nil
}

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)
		}
	}

	// Convert to embedding.Item for the generic index
	items := make([]embedding.Item, len(skills))
	for i, skill := range skills {
		items[i] = embedding.Item{
			Name:        skill.Name,
			Description: skill.Description,
		}
	}

	// Build embedding index
	embIdx, err := embedding.NewIndex(idx.model, items)
	if err != nil {
		return err
	}

	idx.skills = skills
	idx.index = embIdx
	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")
}