Changeset 0.14.0 (#67)

server/knowledge/embedder.go

@@ -0,0 +1,184 @@
+package knowledge
+
+import (
+	"context"
+	"fmt"
+	"log"
+
+	"git.gobha.me/xcaliber/chat-switchboard/models"
+	"git.gobha.me/xcaliber/chat-switchboard/roles"
+	"git.gobha.me/xcaliber/chat-switchboard/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
+}