core/server/knowledge/embedder.go

package knowledge

import (
	"context"
	"fmt"
	"log"

	"switchboard-core/models"
	"switchboard-core/roles"
	"switchboard-core/store"
)

// ── Configuration ────────────────────────────

const (
	// DefaultBatchSize is the max chunks per embedding API call.
	// Most providers handle 100 inputs per request comfortably.
	DefaultBatchSize = 100

	// MaxVectorDim is the column width in pgvector (vector(3072)).
	// Vectors shorter than this are zero-padded on storage.
	MaxVectorDim = 3072
)

// ── Embedder ─────────────────────────────────

// Embedder generates vector embeddings for text chunks using the
// embedding role resolver. It handles batching and dimension
// normalization (zero-padding to MaxVectorDim).
type Embedder struct {
	resolver  *roles.Resolver
	stores    store.Stores // optional — for usage tracking
	batchSize int
}

// NewEmbedder creates an embedder that dispatches through the role resolver.
func NewEmbedder(resolver *roles.Resolver) *Embedder {
	return &Embedder{
		resolver:  resolver,
		batchSize: DefaultBatchSize,
	}
}

// WithStores attaches stores for usage tracking. Returns self for chaining.
func (e *Embedder) WithStores(s store.Stores) *Embedder {
	e.stores = s
	return e
}

// EmbedResult holds the output of a batch embedding operation.
type EmbedResult struct {
	Vectors       [][]float64 // one per input chunk, zero-padded to MaxVectorDim
	Model         string      // model that produced the embeddings
	Dimensions    int         // native dimension before padding
	InputTokens   int         // total tokens consumed across all batches
	ConfigID      string      // provider config that was used
	ProviderScope string      // "personal", "team", or "global"
}

// EmbedChunks generates embeddings for a slice of text chunks.
// Chunks are batched in groups of batchSize to avoid provider limits.
// Uses the embedding role resolution chain: personal → team → global.
//
// Returns vectors in the same order as input chunks.
func (e *Embedder) EmbedChunks(ctx context.Context, userID string, teamID *string, texts []string) (*EmbedResult, error) {
	if len(texts) == 0 {
		return &EmbedResult{Vectors: [][]float64{}}, nil
	}

	// Verify embedding role is configured before starting
	cfg, err := e.resolver.GetConfig(ctx, roles.RoleEmbedding, userID, teamID)
	if err != nil {
		return nil, fmt.Errorf("embedding role not configured: %w", err)
	}
	if cfg.Primary == nil {
		return nil, fmt.Errorf("embedding role has no primary binding")
	}

	allVectors := make([][]float64, 0, len(texts))
	var model string
	var nativeDim int
	var configID, provScope string
	totalTokens := 0

	for i := 0; i < len(texts); i += e.batchSize {
		end := i + e.batchSize
		if end > len(texts) {
			end = len(texts)
		}
		batch := texts[i:end]

		log.Printf("  embed batch %d/%d (%d chunks)",
			(i/e.batchSize)+1, (len(texts)+e.batchSize-1)/e.batchSize, len(batch))

		result, err := e.resolver.Embed(ctx, roles.RoleEmbedding, userID, teamID, batch)
		if err != nil {
			return nil, fmt.Errorf("embed batch starting at chunk %d: %w", i, err)
		}

		if len(result.Embeddings) != len(batch) {
			return nil, fmt.Errorf("embed batch %d: expected %d vectors, got %d",
				i/e.batchSize, len(batch), len(result.Embeddings))
		}

		// Capture model info from first batch
		if model == "" {
			model = result.Model
			configID = result.ConfigID
			provScope = result.ProviderScope
			if len(result.Embeddings) > 0 {
				nativeDim = len(result.Embeddings[0])
			}
		}

		totalTokens += result.InputTokens
		allVectors = append(allVectors, result.Embeddings...)
	}

	// Zero-pad to MaxVectorDim for uniform storage
	for i, vec := range allVectors {
		allVectors[i] = padVector(vec, MaxVectorDim)
	}

	return &EmbedResult{
		Vectors:       allVectors,
		Model:         model,
		Dimensions:    nativeDim,
		InputTokens:   totalTokens,
		ConfigID:      configID,
		ProviderScope: provScope,
	}, nil
}

// LogUsage records an embedding usage entry. Silently no-ops if stores
// were not attached via WithStores or if there are no tokens to log.
func (e *Embedder) LogUsage(ctx context.Context, userID string, channelID *string, result *EmbedResult) {
	if e.stores.Usage == nil || result == nil || result.InputTokens == 0 {
		return
	}

	role := "embedding"
	entry := &models.UsageEntry{
		ChannelID:        channelID,
		UserID:           userID,
		ProviderConfigID: strPtr(result.ConfigID),
		ProviderScope:    result.ProviderScope,
		ModelID:          result.Model,
		Role:             &role,
		PromptTokens:     result.InputTokens,
		CompletionTokens: 0,
	}

	if err := e.stores.Usage.Log(ctx, entry); err != nil {
		log.Printf("⚠ Failed to log embedding usage: %v", err)
	}
}

// IsConfigured returns true if the embedding role has at least a primary binding.
func (e *Embedder) IsConfigured(ctx context.Context) bool {
	return e.resolver.IsConfigured(ctx, roles.RoleEmbedding)
}

// ── Helpers ──────────────────────────────────

// padVector zero-pads a vector to the target dimension.
// If the vector is already the target size or larger, it's truncated.
func padVector(vec []float64, targetDim int) []float64 {
	if len(vec) == targetDim {
		return vec
	}
	if len(vec) > targetDim {
		return vec[:targetDim]
	}
	padded := make([]float64, targetDim)
	copy(padded, vec)
	return padded
}

func strPtr(s string) *string {
	if s == "" {
		return nil
	}
	return &s
}